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

package rpc

import (
	"bufio"
	"encoding/gob"
	"errors"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
)

// ServerError represents an error that has been returned from
// the remote side of the RPC connection.
type ServerError string

func (e ServerError) Error() string {
	return string(e)
}

var ErrShutdown = errors.New("connection is shut down")

// Call represents an active RPC.
type Call struct {
	ServiceMethod string      // The name of the service and method to call.
	Args          interface{} // The argument to the function (*struct).
	Reply         interface{} // The reply from the function (*struct).
	Error         error       // After completion, the error status.
	Done          chan *Call  // Strobes when call is complete; value is the error status.
	seq           uint64
}

// Client represents an RPC Client.
// There may be multiple outstanding Calls associated
// with a single Client.
type Client struct {
	mutex    sync.Mutex // protects pending, seq, request
	sending  sync.Mutex
	request  Request
	seq      uint64
	codec    ClientCodec
	pending  map[uint64]*Call
	closing  bool
	shutdown bool
}

// A ClientCodec implements writing of RPC requests and
// reading of RPC responses for the client side of an RPC session.
// The client calls WriteRequest to write a request to the connection
// and calls ReadResponseHeader and ReadResponseBody in pairs
// to read responses.  The client calls Close when finished with the
// connection. ReadResponseBody may be called with a nil
// argument to force the body of the response to be read and then
// discarded.
type ClientCodec interface {
	WriteRequest(*Request, interface{}) error
	ReadResponseHeader(*Response) error
	ReadResponseBody(interface{}) error

	Close() error
}

func (client *Client) send(c *Call) {
	// Register this call.
	client.mutex.Lock()
	if client.shutdown {
		c.Error = ErrShutdown
		client.mutex.Unlock()
		c.done()
		return
	}
	c.seq = client.seq
	client.seq++
	client.pending[c.seq] = c
	client.mutex.Unlock()

	// Encode and send the request.
	client.sending.Lock()
	defer client.sending.Unlock()
	client.request.Seq = c.seq
	client.request.ServiceMethod = c.ServiceMethod
	if err := client.codec.WriteRequest(&client.request, c.Args); err != nil {
		c.Error = err
		c.done()
	}
}

func (client *Client) input() {
	var err error
	var response Response
	for err == nil {
		response = Response{}
		err = client.codec.ReadResponseHeader(&response)
		if err != nil {
			if err == io.EOF && !client.closing {
				err = io.ErrUnexpectedEOF
			}
			break
		}
		seq := response.Seq
		client.mutex.Lock()
		c := client.pending[seq]
		delete(client.pending, seq)
		client.mutex.Unlock()

		if response.Error == "" {
			err = client.codec.ReadResponseBody(c.Reply)
			if err != nil {
				c.Error = errors.New("reading body " + err.Error())
			}
		} else {
			// We've got an error response. Give this to the request;
			// any subsequent requests will get the ReadResponseBody
			// error if there is one.
			c.Error = ServerError(response.Error)
			err = client.codec.ReadResponseBody(nil)
			if err != nil {
				err = errors.New("reading error body: " + err.Error())
			}
		}
		c.done()
	}
	// Terminate pending calls.
	client.mutex.Lock()
	client.shutdown = true
	for _, call := range client.pending {
		call.Error = err
		call.done()
	}
	client.mutex.Unlock()
	if err != io.EOF || !client.closing {
		log.Println("rpc: client protocol error:", err)
	}
}

func (call *Call) done() {
	select {
	case call.Done <- call:
		// ok
	default:
		// We don't want to block here.  It is the caller's responsibility to make
		// sure the channel has enough buffer space. See comment in Go().
		log.Println("rpc: discarding Call reply due to insufficient Done chan capacity")
	}
}

// NewClient returns a new Client to handle requests to the
// set of services at the other end of the connection.
// It adds a buffer to the write side of the connection so
// the header and payload are sent as a unit.
func NewClient(conn io.ReadWriteCloser) *Client {
	encBuf := bufio.NewWriter(conn)
	client := &gobClientCodec{conn, gob.NewDecoder(conn), gob.NewEncoder(encBuf), encBuf}
	return NewClientWithCodec(client)
}

// NewClientWithCodec is like NewClient but uses the specified
// codec to encode requests and decode responses.
func NewClientWithCodec(codec ClientCodec) *Client {
	client := &Client{
		codec:   codec,
		pending: make(map[uint64]*Call),
	}
	go client.input()
	return client
}

type gobClientCodec struct {
	rwc    io.ReadWriteCloser
	dec    *gob.Decoder
	enc    *gob.Encoder
	encBuf *bufio.Writer
}

func (c *gobClientCodec) WriteRequest(r *Request, body interface{}) (err error) {
	if err = c.enc.Encode(r); err != nil {
		return
	}
	if err = c.enc.Encode(body); err != nil {
		return
	}
	return c.encBuf.Flush()
}

func (c *gobClientCodec) ReadResponseHeader(r *Response) error {
	return c.dec.Decode(r)
}

func (c *gobClientCodec) ReadResponseBody(body interface{}) error {
	return c.dec.Decode(body)
}

func (c *gobClientCodec) Close() error {
	return c.rwc.Close()
}

// DialHTTP connects to an HTTP RPC server at the specified network address
// listening on the default HTTP RPC path.
func DialHTTP(network, address string) (*Client, error) {
	return DialHTTPPath(network, address, DefaultRPCPath)
}

// DialHTTPPath connects to an HTTP RPC server
// at the specified network address and path.
func DialHTTPPath(network, address, path string) (*Client, error) {
	var err error
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	io.WriteString(conn, "CONNECT "+path+" HTTP/1.0\n\n")

	// Require successful HTTP response
	// before switching to RPC protocol.
	resp, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: "CONNECT"})
	if err == nil && resp.Status == connected {
		return NewClient(conn), nil
	}
	if err == nil {
		err = errors.New("unexpected HTTP response: " + resp.Status)
	}
	conn.Close()
	return nil, &net.OpError{"dial-http", network + " " + address, nil, err}
}

// Dial connects to an RPC server at the specified network address.
func Dial(network, address string) (*Client, error) {
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return NewClient(conn), nil
}

func (client *Client) Close() error {
	client.mutex.Lock()
	if client.shutdown || client.closing {
		client.mutex.Unlock()
		return ErrShutdown
	}
	client.closing = true
	client.mutex.Unlock()
	return client.codec.Close()
}

// Go invokes the function asynchronously.  It returns the Call structure representing
// the invocation.  The done channel will signal when the call is complete by returning
// the same Call object.  If done is nil, Go will allocate a new channel.
// If non-nil, done must be buffered or Go will deliberately crash.
func (client *Client) Go(serviceMethod string, args interface{}, reply interface{}, done chan *Call) *Call {
	call := new(Call)
	call.ServiceMethod = serviceMethod
	call.Args = args
	call.Reply = reply
	if done == nil {
		done = make(chan *Call, 10) // buffered.
	} else {
		// If caller passes done != nil, it must arrange that
		// done has enough buffer for the number of simultaneous
		// RPCs that will be using that channel.  If the channel
		// is totally unbuffered, it's best not to run at all.
		if cap(done) == 0 {
			log.Panic("rpc: done channel is unbuffered")
		}
	}
	call.Done = done
	if client.shutdown {
		call.Error = ErrShutdown
		call.done()
		return call
	}
	client.send(call)
	return call
}

// Call invokes the named function, waits for it to complete, and returns its error status.
func (client *Client) Call(serviceMethod string, args interface{}, reply interface{}) error {
	if client.shutdown {
		return ErrShutdown
	}
	call := <-client.Go(serviceMethod, args, reply, make(chan *Call, 1)).Done
	return call.Error
}
<commit_msg>net/rpc: fix data race on Call.Error +eliminates a possibility of sending a call to Done several times. +fixes memory leak in case of temporal Write errors. +fixes data race on Client.shutdown. +fixes data race on Client.closing. +fixes comments. Fixes issue 2780.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rpc

import (
	"bufio"
	"encoding/gob"
	"errors"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
)

// ServerError represents an error that has been returned from
// the remote side of the RPC connection.
type ServerError string

func (e ServerError) Error() string {
	return string(e)
}

var ErrShutdown = errors.New("connection is shut down")

// Call represents an active RPC.
type Call struct {
	ServiceMethod string      // The name of the service and method to call.
	Args          interface{} // The argument to the function (*struct).
	Reply         interface{} // The reply from the function (*struct).
	Error         error       // After completion, the error status.
	Done          chan *Call  // Strobes when call is complete.
}

// Client represents an RPC Client.
// There may be multiple outstanding Calls associated
// with a single Client.
type Client struct {
	mutex    sync.Mutex // protects pending, seq, request
	sending  sync.Mutex
	request  Request
	seq      uint64
	codec    ClientCodec
	pending  map[uint64]*Call
	closing  bool
	shutdown bool
}

// A ClientCodec implements writing of RPC requests and
// reading of RPC responses for the client side of an RPC session.
// The client calls WriteRequest to write a request to the connection
// and calls ReadResponseHeader and ReadResponseBody in pairs
// to read responses.  The client calls Close when finished with the
// connection. ReadResponseBody may be called with a nil
// argument to force the body of the response to be read and then
// discarded.
type ClientCodec interface {
	WriteRequest(*Request, interface{}) error
	ReadResponseHeader(*Response) error
	ReadResponseBody(interface{}) error

	Close() error
}

func (client *Client) send(call *Call) {
	client.sending.Lock()
	defer client.sending.Unlock()

	// Register this call.
	client.mutex.Lock()
	if client.shutdown {
		call.Error = ErrShutdown
		client.mutex.Unlock()
		call.done()
		return
	}
	seq := client.seq
	client.seq++
	client.pending[seq] = call
	client.mutex.Unlock()

	// Encode and send the request.
	client.request.Seq = seq
	client.request.ServiceMethod = call.ServiceMethod
	err := client.codec.WriteRequest(&client.request, call.Args)
	if err != nil {
		client.mutex.Lock()
		delete(client.pending, seq)
		client.mutex.Unlock()
		call.Error = err
		call.done()
	}
}

func (client *Client) input() {
	var err error
	var response Response
	for err == nil {
		response = Response{}
		err = client.codec.ReadResponseHeader(&response)
		if err != nil {
			if err == io.EOF && !client.closing {
				err = io.ErrUnexpectedEOF
			}
			break
		}
		seq := response.Seq
		client.mutex.Lock()
		call := client.pending[seq]
		delete(client.pending, seq)
		client.mutex.Unlock()

		if response.Error == "" {
			err = client.codec.ReadResponseBody(call.Reply)
			if err != nil {
				call.Error = errors.New("reading body " + err.Error())
			}
		} else {
			// We've got an error response. Give this to the request;
			// any subsequent requests will get the ReadResponseBody
			// error if there is one.
			call.Error = ServerError(response.Error)
			err = client.codec.ReadResponseBody(nil)
			if err != nil {
				err = errors.New("reading error body: " + err.Error())
			}
		}
		call.done()
	}
	// Terminate pending calls.
	client.sending.Lock()
	client.mutex.Lock()
	client.shutdown = true
	closing := client.closing
	for _, call := range client.pending {
		call.Error = err
		call.done()
	}
	client.mutex.Unlock()
	client.sending.Unlock()
	if err != io.EOF || !closing {
		log.Println("rpc: client protocol error:", err)
	}
}

func (call *Call) done() {
	select {
	case call.Done <- call:
		// ok
	default:
		// We don't want to block here.  It is the caller's responsibility to make
		// sure the channel has enough buffer space. See comment in Go().
		log.Println("rpc: discarding Call reply due to insufficient Done chan capacity")
	}
}

// NewClient returns a new Client to handle requests to the
// set of services at the other end of the connection.
// It adds a buffer to the write side of the connection so
// the header and payload are sent as a unit.
func NewClient(conn io.ReadWriteCloser) *Client {
	encBuf := bufio.NewWriter(conn)
	client := &gobClientCodec{conn, gob.NewDecoder(conn), gob.NewEncoder(encBuf), encBuf}
	return NewClientWithCodec(client)
}

// NewClientWithCodec is like NewClient but uses the specified
// codec to encode requests and decode responses.
func NewClientWithCodec(codec ClientCodec) *Client {
	client := &Client{
		codec:   codec,
		pending: make(map[uint64]*Call),
	}
	go client.input()
	return client
}

type gobClientCodec struct {
	rwc    io.ReadWriteCloser
	dec    *gob.Decoder
	enc    *gob.Encoder
	encBuf *bufio.Writer
}

func (c *gobClientCodec) WriteRequest(r *Request, body interface{}) (err error) {
	if err = c.enc.Encode(r); err != nil {
		return
	}
	if err = c.enc.Encode(body); err != nil {
		return
	}
	return c.encBuf.Flush()
}

func (c *gobClientCodec) ReadResponseHeader(r *Response) error {
	return c.dec.Decode(r)
}

func (c *gobClientCodec) ReadResponseBody(body interface{}) error {
	return c.dec.Decode(body)
}

func (c *gobClientCodec) Close() error {
	return c.rwc.Close()
}

// DialHTTP connects to an HTTP RPC server at the specified network address
// listening on the default HTTP RPC path.
func DialHTTP(network, address string) (*Client, error) {
	return DialHTTPPath(network, address, DefaultRPCPath)
}

// DialHTTPPath connects to an HTTP RPC server
// at the specified network address and path.
func DialHTTPPath(network, address, path string) (*Client, error) {
	var err error
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	io.WriteString(conn, "CONNECT "+path+" HTTP/1.0\n\n")

	// Require successful HTTP response
	// before switching to RPC protocol.
	resp, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: "CONNECT"})
	if err == nil && resp.Status == connected {
		return NewClient(conn), nil
	}
	if err == nil {
		err = errors.New("unexpected HTTP response: " + resp.Status)
	}
	conn.Close()
	return nil, &net.OpError{"dial-http", network + " " + address, nil, err}
}

// Dial connects to an RPC server at the specified network address.
func Dial(network, address string) (*Client, error) {
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return NewClient(conn), nil
}

func (client *Client) Close() error {
	client.mutex.Lock()
	if client.shutdown || client.closing {
		client.mutex.Unlock()
		return ErrShutdown
	}
	client.closing = true
	client.mutex.Unlock()
	return client.codec.Close()
}

// Go invokes the function asynchronously.  It returns the Call structure representing
// the invocation.  The done channel will signal when the call is complete by returning
// the same Call object.  If done is nil, Go will allocate a new channel.
// If non-nil, done must be buffered or Go will deliberately crash.
func (client *Client) Go(serviceMethod string, args interface{}, reply interface{}, done chan *Call) *Call {
	call := new(Call)
	call.ServiceMethod = serviceMethod
	call.Args = args
	call.Reply = reply
	if done == nil {
		done = make(chan *Call, 10) // buffered.
	} else {
		// If caller passes done != nil, it must arrange that
		// done has enough buffer for the number of simultaneous
		// RPCs that will be using that channel.  If the channel
		// is totally unbuffered, it's best not to run at all.
		if cap(done) == 0 {
			log.Panic("rpc: done channel is unbuffered")
		}
	}
	call.Done = done
	client.send(call)
	return call
}

// Call invokes the named function, waits for it to complete, and returns its error status.
func (client *Client) Call(serviceMethod string, args interface{}, reply interface{}) error {
	call := <-client.Go(serviceMethod, args, reply, make(chan *Call, 1)).Done
	return call.Error
}
<|endoftext|>


<commit_before>// Copyright 2014 Chris Monson <shiblon@gmail.com>
//
// 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 keyheap implements a library for a simple heap that allows peeking and
popping from the middle based on a Key() in the stored interface.
*/
package keyheap // import "entrogo.com/taskstore/keyheap"

import (
	"container/heap"
	"fmt"
	"math"
	"math/rand"
	"strings"
)

type Item interface {
	// Priority is used to determine which element (lowest priority) is at the
	// top of the heap.
	Priority() int64

	// Key is used to look up individual items inside of the heap.
	Key() int64
}

type KeyHeap struct {
	itemHeap heapImpl
	itemMap  heapMap
}

// New creates a new empty KeyHeap.
func New() *KeyHeap {
	return NewFromItems(nil)
}

// NewFromItems creates a KeyHeap from a slice of Item.
func NewFromItems(items []Item) *KeyHeap {
	q := &KeyHeap{
		itemHeap: make(heapImpl, len(items)),
		itemMap:  make(heapMap),
	}
	for i, item := range items {
		ti := &indexedItem{index: i, item: item}
		q.itemHeap[i] = ti
		q.itemMap[item.Key()] = ti
	}

	if q.Len() > 1 {
		heap.Init(&q.itemHeap)
	}
	return q
}

// String formats this heap into a string.
func (q *KeyHeap) String() string {
	hpieces := []string{"["}
	for _, v := range q.itemHeap {
		hpieces = append(hpieces, fmt.Sprintf("   %s", v))
	}
	if len(hpieces) == 1 {
		hpieces[0] += "]"
	} else {
		hpieces = append(hpieces, "]")
	}
	return fmt.Sprintf("KeyHeap(%v)", strings.Join(hpieces, "\n     "))
}

// Push adds an Item to the heap.
func (q *KeyHeap) Push(item Item) {
	ti := &indexedItem{item: item, index: -1}
	heap.Push(&q.itemHeap, ti)
	q.itemMap[item.Key()] = ti
}

// Pop removes the Item with the lowest Priority() from the KeyHeap.
func (q *KeyHeap) Pop() Item {
	ti := heap.Pop(&q.itemHeap).(*indexedItem)
	delete(q.itemMap, ti.item.Key())
	return ti.item
}

// PopAt removes an element from the specified index in the heap in O(log(n)) time.
func (q *KeyHeap) PopAt(idx int) Item {
	ti := heap.Remove(&q.itemHeap, idx).(*indexedItem)
	if ti == nil {
		return nil
	}
	delete(q.itemMap, ti.item.Key())
	return ti.item
}

// Len returns the size of the heap.
func (q *KeyHeap) Len() int {
	return len(q.itemHeap)
}

// Peek returns the top element in the heap (with the smallest Priority()), or nil if the heap is empty.
func (q *KeyHeap) Peek() Item {
	return q.PeekAt(0)
}

// PeekAt finds the item at index idx in the heap and returns it. Returns nil if idx is out of bounds.
func (q *KeyHeap) PeekAt(idx int) Item {
	if idx >= q.Len() {
		return nil
	}
	return q.itemHeap[idx].item
}

// PeekByKey finds the item with the given Key() and returns it, or nil if not found.
func (q *KeyHeap) PeekByKey(key int64) Item {
	ti := q.itemMap[key]
	if ti == nil {
		return nil
	}
	return ti.item
}

// PopByKey finds the item with the given Key() and returns it, removing it
// from the data structure.
func (q *KeyHeap) PopByKey(key int64) Item {
	ti := q.itemMap[key]
	if ti == nil {
		return nil
	}
	return q.PopAt(ti.index)
}

// PopRandomConstrained walks the heap randomly choosing a child until it either
// picks one or runs out (and picks the last one before the maxPriority). If
// maxPriority <= 0, then there is no constraint.
// Note that this greatly favors items near the top, because the probability of
// traversing the tree very far quickly gets vanishingly small. There are
// undoubtedly other interesting approaches to doing this.
func (q *KeyHeap) PopRandomConstrained(maxPriority int64) Item {
	// Start at the leftmost location (the lowest value), and randomly jump to
	// children so long as they are earlier than the maxPriority.
	idx := 0
	chosen := -1
	for idx < q.Len() && q.PeekAt(idx).Priority() <= maxPriority {
		left := idx*2 + 1
		right := left + 1
		choices := make([]int, 1, 3)
		choices[0] = idx
		if left < q.Len() && q.PeekAt(left).Priority() <= maxPriority {
			choices = append(choices, left)
		}
		if right < q.Len() && q.PeekAt(right).Priority() <= maxPriority {
			choices = append(choices, right)
		}
		if len(choices) == 0 {
			break
		}
		choiceIndex := int(math.Floor(rand.Float64() * float64(len(choices))))
		if choiceIndex == 0 {
			chosen = choices[choiceIndex]
			break
		}
		// If we didn't choose the current node, redo the random draw with one of
		// the children as the new heap root.
		idx = choices[choiceIndex]
	}
	if chosen < 0 {
		return nil
	}
	return q.PopAt(chosen)
}

type indexedItem struct {
	index int
	item  Item
}

func (ti *indexedItem) String() string {
	return fmt.Sprintf("{%d:%v}", ti.index, ti.item)
}

type heapImpl []*indexedItem
type heapMap map[int64]*indexedItem

func (h heapImpl) Len() int {
	return len(h)
}

func (h heapImpl) Less(i, j int) bool {
	if h[i].item == nil {
		return true
	} else if h[j].item == nil {
		return false
	}
	return h[i].item.Priority() < h[j].item.Priority()
}

func (h heapImpl) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
	h[i].index = i
	h[j].index = j
}

func (h *heapImpl) Push(x interface{}) {
	item := x.(*indexedItem)
	item.index = len(*h)
	*h = append(*h, item)
}

func (h *heapImpl) Pop() interface{} {
	n := len(*h)
	item := (*h)[n-1]
	item.index = -1
	*h = (*h)[:n-1]
	return item
}
<commit_msg>Simplified Less, no longer needs to know about nil due to heap.Remove change.<commit_after>// Copyright 2014 Chris Monson <shiblon@gmail.com>
//
// 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 keyheap implements a library for a simple heap that allows peeking and
popping from the middle based on a Key() in the stored interface.
*/
package keyheap // import "entrogo.com/taskstore/keyheap"

import (
	"container/heap"
	"fmt"
	"math"
	"math/rand"
	"strings"
)

type Item interface {
	// Priority is used to determine which element (lowest priority) is at the
	// top of the heap.
	Priority() int64

	// Key is used to look up individual items inside of the heap.
	Key() int64
}

type KeyHeap struct {
	itemHeap heapImpl
	itemMap  heapMap
}

// New creates a new empty KeyHeap.
func New() *KeyHeap {
	return NewFromItems(nil)
}

// NewFromItems creates a KeyHeap from a slice of Item.
func NewFromItems(items []Item) *KeyHeap {
	q := &KeyHeap{
		itemHeap: make(heapImpl, len(items)),
		itemMap:  make(heapMap),
	}
	for i, item := range items {
		ti := &indexedItem{index: i, item: item}
		q.itemHeap[i] = ti
		q.itemMap[item.Key()] = ti
	}

	if q.Len() > 1 {
		heap.Init(&q.itemHeap)
	}
	return q
}

// String formats this heap into a string.
func (q *KeyHeap) String() string {
	hpieces := []string{"["}
	for _, v := range q.itemHeap {
		hpieces = append(hpieces, fmt.Sprintf("   %s", v))
	}
	if len(hpieces) == 1 {
		hpieces[0] += "]"
	} else {
		hpieces = append(hpieces, "]")
	}
	return fmt.Sprintf("KeyHeap(%v)", strings.Join(hpieces, "\n     "))
}

// Push adds an Item to the heap.
func (q *KeyHeap) Push(item Item) {
	ti := &indexedItem{item: item, index: -1}
	heap.Push(&q.itemHeap, ti)
	q.itemMap[item.Key()] = ti
}

// Pop removes the Item with the lowest Priority() from the KeyHeap.
func (q *KeyHeap) Pop() Item {
	ti := heap.Pop(&q.itemHeap).(*indexedItem)
	delete(q.itemMap, ti.item.Key())
	return ti.item
}

// PopAt removes an element from the specified index in the heap in O(log(n)) time.
func (q *KeyHeap) PopAt(idx int) Item {
	ti := heap.Remove(&q.itemHeap, idx).(*indexedItem)
	if ti == nil {
		return nil
	}
	delete(q.itemMap, ti.item.Key())
	return ti.item
}

// Len returns the size of the heap.
func (q *KeyHeap) Len() int {
	return len(q.itemHeap)
}

// Peek returns the top element in the heap (with the smallest Priority()), or nil if the heap is empty.
func (q *KeyHeap) Peek() Item {
	return q.PeekAt(0)
}

// PeekAt finds the item at index idx in the heap and returns it. Returns nil if idx is out of bounds.
func (q *KeyHeap) PeekAt(idx int) Item {
	if idx >= q.Len() {
		return nil
	}
	return q.itemHeap[idx].item
}

// PeekByKey finds the item with the given Key() and returns it, or nil if not found.
func (q *KeyHeap) PeekByKey(key int64) Item {
	ti := q.itemMap[key]
	if ti == nil {
		return nil
	}
	return ti.item
}

// PopByKey finds the item with the given Key() and returns it, removing it
// from the data structure.
func (q *KeyHeap) PopByKey(key int64) Item {
	ti := q.itemMap[key]
	if ti == nil {
		return nil
	}
	return q.PopAt(ti.index)
}

// PopRandomConstrained walks the heap randomly choosing a child until it either
// picks one or runs out (and picks the last one before the maxPriority). If
// maxPriority <= 0, then there is no constraint.
// Note that this greatly favors items near the top, because the probability of
// traversing the tree very far quickly gets vanishingly small. There are
// undoubtedly other interesting approaches to doing this.
func (q *KeyHeap) PopRandomConstrained(maxPriority int64) Item {
	// Start at the leftmost location (the lowest value), and randomly jump to
	// children so long as they are earlier than the maxPriority.
	idx := 0
	chosen := -1
	for idx < q.Len() && q.PeekAt(idx).Priority() <= maxPriority {
		left := idx*2 + 1
		right := left + 1
		choices := make([]int, 1, 3)
		choices[0] = idx
		if left < q.Len() && q.PeekAt(left).Priority() <= maxPriority {
			choices = append(choices, left)
		}
		if right < q.Len() && q.PeekAt(right).Priority() <= maxPriority {
			choices = append(choices, right)
		}
		if len(choices) == 0 {
			break
		}
		choiceIndex := int(math.Floor(rand.Float64() * float64(len(choices))))
		if choiceIndex == 0 {
			chosen = choices[choiceIndex]
			break
		}
		// If we didn't choose the current node, redo the random draw with one of
		// the children as the new heap root.
		idx = choices[choiceIndex]
	}
	if chosen < 0 {
		return nil
	}
	return q.PopAt(chosen)
}

type indexedItem struct {
	index int
	item  Item
}

func (ti *indexedItem) String() string {
	return fmt.Sprintf("{%d:%v}", ti.index, ti.item)
}

type heapImpl []*indexedItem
type heapMap map[int64]*indexedItem

func (h heapImpl) Len() int {
	return len(h)
}

func (h heapImpl) Less(i, j int) bool {
	return h[i].item.Priority() < h[j].item.Priority()
}

func (h heapImpl) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
	h[i].index = i
	h[j].index = j
}

func (h *heapImpl) Push(x interface{}) {
	item := x.(*indexedItem)
	item.index = len(*h)
	*h = append(*h, item)
}

func (h *heapImpl) Pop() interface{} {
	n := len(*h)
	item := (*h)[n-1]
	item.index = -1
	*h = (*h)[:n-1]
	return item
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package keymutex

import (
	"hash/fnv"
	"runtime"
	"sync"

	"github.com/golang/glog"
)

// NewHashed returns a new instance of KeyMutex which hashes arbitrary keys to
// a fixed set of locks. `n` specifies number of locks, if n <= 0, we use
// number of cpus.
// Note that because it uses fixed set of locks, different keys may share same
// lock, so it's possible to wait on same lock.
func NewHashed(n int) KeyMutex {
	if n <= 0 {
		n = runtime.NumCPU()
	}
	return &hashedKeyMutex{
		mutexes: make([]sync.Mutex, n),
	}
}

type hashedKeyMutex struct {
	mutexes []sync.Mutex
}

// Acquires a lock associated with the specified ID.
func (km *hashedKeyMutex) LockKey(id string) {
	glog.V(5).Infof("hashedKeyMutex.LockKey(...) called for id %q\r\n", id)
	km.mutexes[km.hash(id)%len(km.mutexes)].Lock()
	glog.V(5).Infof("hashedKeyMutex.LockKey(...) for id %q completed.\r\n", id)
}

// Releases the lock associated with the specified ID.
func (km *hashedKeyMutex) UnlockKey(id string) error {
	glog.V(5).Infof("hashedKeyMutex.UnlockKey(...) called for id %q\r\n", id)
	km.mutexes[km.hash(id)%len(km.mutexes)].Unlock()
	glog.V(5).Infof("hashedKeyMutex.UnlockKey(...) for id %q completed.\r\n", id)
	return nil
}

func (km *hashedKeyMutex) hash(id string) int {
	h := fnv.New32a()
	h.Write([]byte(id))
	return int(h.Sum32())
}
<commit_msg>keymutex: remove glog usage<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package keymutex

import (
	"hash/fnv"
	"runtime"
	"sync"
)

// NewHashed returns a new instance of KeyMutex which hashes arbitrary keys to
// a fixed set of locks. `n` specifies number of locks, if n <= 0, we use
// number of cpus.
// Note that because it uses fixed set of locks, different keys may share same
// lock, so it's possible to wait on same lock.
func NewHashed(n int) KeyMutex {
	if n <= 0 {
		n = runtime.NumCPU()
	}
	return &hashedKeyMutex{
		mutexes: make([]sync.Mutex, n),
	}
}

type hashedKeyMutex struct {
	mutexes []sync.Mutex
}

// Acquires a lock associated with the specified ID.
func (km *hashedKeyMutex) LockKey(id string) {
	km.mutexes[km.hash(id)%len(km.mutexes)].Lock()
}

// Releases the lock associated with the specified ID.
func (km *hashedKeyMutex) UnlockKey(id string) error {
	km.mutexes[km.hash(id)%len(km.mutexes)].Unlock()
	return nil
}

func (km *hashedKeyMutex) hash(id string) int {
	h := fnv.New32a()
	h.Write([]byte(id))
	return int(h.Sum32())
}
<|endoftext|>


<commit_before>package imageserver

import (
	"fmt"
	"sync"
	"testing"
)

func TestNewCacheMissError(t *testing.T) {
	key := "foobar"
	cache := newCacheMap()
	previousErr := fmt.Errorf("not found")

	err := NewCacheMissError(key, cache, previousErr)
	err.Error()
}

type cacheMap struct {
	mutex sync.RWMutex
	data  map[string]*Image
}

func newCacheMap() *cacheMap {
	return &cacheMap{
		data: make(map[string]*Image),
	}
}

func (cache *cacheMap) Get(key string, parameters Parameters) (*Image, error) {
	cache.mutex.RLock()
	defer cache.mutex.RUnlock()

	image, ok := cache.data[key]
	if !ok {
		return nil, NewCacheMissError(key, cache, nil)
	}

	return image, nil
}

func (cache *cacheMap) Set(key string, image *Image, parameters Parameters) error {
	cache.mutex.Lock()
	defer cache.mutex.Unlock()

	cache.data[key] = image

	return nil
}

type cacheFunc struct {
	GetFunc func(key string, parameters Parameters) (*Image, error)
	SetFunc func(key string, image *Image, parameters Parameters) error
}

func (cache *cacheFunc) Get(key string, parameters Parameters) (*Image, error) {
	return cache.GetFunc(key, parameters)
}

func (cache *cacheFunc) Set(key string, image *Image, parameters Parameters) error {
	return cache.SetFunc(key, image, parameters)
}
<commit_msg>Add test on NewParametersHashCacheKeyFunc<commit_after>package imageserver

import (
	"crypto/sha256"
	"fmt"
	"sync"
	"testing"
)

func TestNewCacheMissError(t *testing.T) {
	key := "foobar"
	cache := newCacheMap()
	previousErr := fmt.Errorf("not found")

	err := NewCacheMissError(key, cache, previousErr)
	err.Error()
}

func TestNewParametersHashCacheKeyFunc(t *testing.T) {
	f := NewParametersHashCacheKeyFunc(sha256.New)
	parameters := Parameters{
		"foo": "bar",
	}
	f(parameters)
}

type cacheMap struct {
	mutex sync.RWMutex
	data  map[string]*Image
}

func newCacheMap() *cacheMap {
	return &cacheMap{
		data: make(map[string]*Image),
	}
}

func (cache *cacheMap) Get(key string, parameters Parameters) (*Image, error) {
	cache.mutex.RLock()
	defer cache.mutex.RUnlock()

	image, ok := cache.data[key]
	if !ok {
		return nil, NewCacheMissError(key, cache, nil)
	}

	return image, nil
}

func (cache *cacheMap) Set(key string, image *Image, parameters Parameters) error {
	cache.mutex.Lock()
	defer cache.mutex.Unlock()

	cache.data[key] = image

	return nil
}

type cacheFunc struct {
	GetFunc func(key string, parameters Parameters) (*Image, error)
	SetFunc func(key string, image *Image, parameters Parameters) error
}

func (cache *cacheFunc) Get(key string, parameters Parameters) (*Image, error) {
	return cache.GetFunc(key, parameters)
}

func (cache *cacheFunc) Set(key string, image *Image, parameters Parameters) error {
	return cache.SetFunc(key, image, parameters)
}
<|endoftext|>


<commit_before>package sshvault

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"
)

func TestCacheIsFile(t *testing.T) {
	cache := &cache{}
	if cache.IsFile("/") {
		t.Errorf("Expecting false")
	}
	if !cache.IsFile("cache_test.go") {
		t.Errorf("Expecting true")
	}
}

type mockSchlosser struct{}

func (m mockSchlosser) GetKey(u string) ([]string, error) {
	switch u {
	case "alice":
		return []string{"ssh-rsa ABC"}, nil
	case "bob":
		return nil, nil
	case "matilde":
		return []string{"ssh-rsa ABC", "ssh-rsa ABC", "ssh-rsa ABC"}, nil
	default:
		return nil, fmt.Errorf("Error")
	}
}

func TestCacheGet(t *testing.T) {
	dir, err := ioutil.TempDir("", "cache")
	if err != nil {
		t.Error(err)
	}
	defer os.RemoveAll(dir) // clean up
	var testTable = []struct {
		user string
		key  int
		out  string
		err  bool
	}{
		{"alice", 0, "alice.key-1", false},
		{"alice", 1, "alice.key-1", false},
		{"alice", 2, "", true},
		{"bob", 1, "", true},
		{"matilde", 3, "matilde.key-3", false},
		{"matilde", 2, "matilde.key-2", false},
		{"matilde", 0, "matilde.key-1", false},
		{"matilde", 4, "", true},
	}
	cache := &cache{dir}
	gk := mockSchlosser{}
	for _, tt := range testTable {
		out, err := cache.Get(gk, tt.user, "", tt.key)
		if tt.err {
			if err == nil {
				t.Error("Expecting error")
			}
		} else if strings.HasPrefix(out, tt.out) {
			t.Errorf("%q != %q", tt.out, out)
		}
	}
}
<commit_msg>test cacheFingerprint<commit_after>package sshvault

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"
)

func TestCacheIsFile(t *testing.T) {
	cache := &cache{}
	if cache.IsFile("/") {
		t.Errorf("Expecting false")
	}
	if !cache.IsFile("cache_test.go") {
		t.Errorf("Expecting true")
	}
}

type mockSchlosser struct{}

func (m mockSchlosser) GetKey(u string) ([]string, error) {
	switch u {
	case "alice":
		return []string{"ssh-rsa ABC"}, nil
	case "bob":
		return nil, nil
	case "matilde":
		return []string{"ssh-rsa ABC", "ssh-rsa ABC", "ssh-rsa ABC"}, nil
	case "pedro":
		return []string{"ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCrrjZ4Hw9wj/RXaNmwAS0eAxub9LYYCv4bsfxE4UmXcLQSj4YIM8+GfsPkykKZNl5+iatzeKrolYCHLIjC1xwsC199o5lpEBskV1g0uFhRiuguUJxM2r66bbxOfuSZcY9tHD/NkgLg0rTqDzGXtkWbBbjtam9N0H4dbCfgVpGVI8feZqFR5uiukG2eDJKn+0S4UTwZgO7TvSxpMl31xqlPy9EsgEhb+19YYuvSQOXWBX6yuKr1gjY7g3/wmtXRdrZbTjZmIeACITNWgWM7TFEqYf88bHHAMz1pSj5V8Uu0k/yEd2RRIHoMc1fMq5ygMEU6mcEf3C8zy6w5r3rRms2n", "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDGOhBrPToSBJCblZoK44w3/ub3K6Vx39ilHB/2sJIDqLZTx8I1U2l2RD3WhwKXdqqpH6RZh0piGlWuGV/E7xOseH9qEOKZMgscdvNO9nzD8jkSlShhZQUmhWOqLPcVUDlgIubxrFRVODcFxqgJwjm+qR2X2GaHJottrn5jFhNBEYcjdnuDKXZQ7Cr+K2bOcD+pvhMI7/qtR7jKa7Q5BoRxQEsNQEZvvgJpen2CqAsnjpJXjAXttnXJnAXcyYyOe8ZOCY/tkmXWvn9Fkd1EYmK14rB8WNEe+vraNCS9tSi1PyLMJWr3XNeluLr2/y7gHSyO6xzQNoXiTDDBFW2y3VK5", "ssh-rsa AAA", "ssh-rsa BBB"}, nil
	default:
		return nil, fmt.Errorf("Error")
	}
}

func TestCacheGet(t *testing.T) {
	dir, err := ioutil.TempDir("", "cache")
	if err != nil {
		t.Error(err)
	}
	defer os.RemoveAll(dir) // clean up
	var testTable = []struct {
		user string
		key  int
		out  string
		err  bool
	}{
		{"alice", 0, "alice.key-1", false},
		{"alice", 1, "alice.key-1", false},
		{"alice", 2, "", true},
		{"bob", 1, "", true},
		{"matilde", 3, "matilde.key-3", false},
		{"matilde", 2, "matilde.key-2", false},
		{"matilde", 0, "matilde.key-1", false},
		{"matilde", 4, "", true},
	}
	cache := &cache{dir}
	gk := mockSchlosser{}
	for _, tt := range testTable {
		out, err := cache.Get(gk, tt.user, "", tt.key)
		if tt.err {
			if err == nil {
				t.Error("Expecting error")
			}
		} else if strings.HasPrefix(out, tt.out) {
			t.Errorf("%q != %q", tt.out, out)
		}
	}
}

func TestCacheGetFingerprint(t *testing.T) {
	dir, err := ioutil.TempDir("", "cacheFingerprint")
	if err != nil {
		t.Error(err)
	}
	defer os.RemoveAll(dir) // clean up
	cache := &cache{dir}
	gk := mockSchlosser{}
	out, err := cache.Get(gk, "pedro", "4a:5e:4b:4d:81:2c:de:db:d5:1d:c3:f9:6e:85:d6:ad", 0)
	if err != nil {
		t.Error(err)
	}
	pubKey, err := ioutil.ReadFile(out)
	expectedKey := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDGOhBrPToSBJCblZoK44w3/ub3K6Vx39ilHB/2sJIDqLZTx8I1U2l2RD3WhwKXdqqpH6RZh0piGlWuGV/E7xOseH9qEOKZMgscdvNO9nzD8jkSlShhZQUmhWOqLPcVUDlgIubxrFRVODcFxqgJwjm+qR2X2GaHJottrn5jFhNBEYcjdnuDKXZQ7Cr+K2bOcD+pvhMI7/qtR7jKa7Q5BoRxQEsNQEZvvgJpen2CqAsnjpJXjAXttnXJnAXcyYyOe8ZOCY/tkmXWvn9Fkd1EYmK14rB8WNEe+vraNCS9tSi1PyLMJWr3XNeluLr2/y7gHSyO6xzQNoXiTDDBFW2y3VK5"
	if string(pubKey) != expectedKey {
		t.Error("Expecting %q got %q", expectedKey, pubKey)
	}
}
<|endoftext|>


<commit_before>package auth

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"
	libkb "github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol"
	context "golang.org/x/net/context"
	"sync"
	"testing"
	"time"
)

type testUser struct {
	uid      keybase1.UID
	username libkb.NormalizedUsername
	sibkeys  []keybase1.KID
	subkeys  []keybase1.KID
}

type testState struct {
	sync.Mutex

	users     map[keybase1.UID](*testUser)
	changes   []keybase1.UID
	now       time.Time
	evictCh   chan keybase1.UID
	pokeCh    chan struct{}
	startOnce sync.Once
	numGets   int
}

var seq uint32

func genKID() keybase1.KID {
	var kid [35]byte
	kid[0] = 0x1
	kid[1] = 0x20
	binary.BigEndian.PutUint32(kid[30:34], seq)
	seq++
	kid[34] = 0xa0
	return keybase1.KIDFromSlice(kid[:])
}

func genUsername() string {
	w, _ := libkb.SecWordList(1)
	var buf [4]byte
	rand.Read(buf[:])
	return fmt.Sprintf("%s%x", w[0], buf)
}

func newTestUser(nKeys int) *testUser {
	un := genUsername()
	ret := testUser{
		username: libkb.NewNormalizedUsername(un),
		uid:      libkb.UsernameToUID(un),
		sibkeys:  make([]keybase1.KID, nKeys),
		subkeys:  make([]keybase1.KID, nKeys),
	}
	for i := 0; i < nKeys; i++ {
		ret.sibkeys[i] = genKID()
		ret.subkeys[i] = genKID()
	}
	return &ret
}

func (ts *testState) newTestUser(nKeys int) *testUser {
	ts.Lock()
	defer ts.Unlock()
	ret := newTestUser(nKeys)
	ts.users[ret.uid] = ret
	return ret
}

func (ts *testState) mutateUser(uid keybase1.UID, mutator func(u *testUser)) bool {
	ts.Lock()
	defer ts.Unlock()
	u := ts.users[uid]
	if u == nil {
		return false
	}
	mutator(u)
	ts.changes = append(ts.changes, uid)
	return true
}

func newTestState() *testState {
	return &testState{
		users:   make(map[keybase1.UID](*testUser)),
		now:     time.Unix(100, 0),
		evictCh: make(chan keybase1.UID, 1),
		pokeCh:  make(chan struct{}),
	}
}

type userNotFoundError struct {
}

func (e userNotFoundError) Error() string {
	return "user not found"
}

func (ts *testState) GetUser(_ context.Context, uid keybase1.UID) (
	un libkb.NormalizedUsername, sibkeys, subkeys []keybase1.KID, err error) {
	ts.Lock()
	defer ts.Unlock()
	u := ts.users[uid]
	if u == nil {
		return libkb.NormalizedUsername(""), nil, nil, userNotFoundError{}
	}
	ts.numGets++
	return u.username, u.sibkeys, u.subkeys, nil
}

func (ts *testState) PollForChanges(_ context.Context) ([]keybase1.UID, error) {
	ts.Lock()
	defer ts.Unlock()
	ret := ts.changes
	ts.changes = nil
	return ret, nil
}

var _ UserKeyAPIer = (*testState)(nil)
var _ engine = (*testState)(nil)

func (ts *testState) tick(d time.Duration) {
	ts.pokeCh <- struct{}{}
	ts.Lock()
	ts.now = ts.now.Add(d)
	ts.Unlock()
	ts.pokeCh <- struct{}{}
}

func (ts *testState) Now() time.Time {
	ts.Lock()
	ret := ts.now
	ts.Unlock()
	return ret
}

func (ts *testState) GetPokeCh() <-chan struct{} { return ts.pokeCh }

func (ts *testState) Evicted(uid keybase1.UID) {
	ts.evictCh <- uid
}

func newTestSetup() (*testState, *CredentialAuthority) {
	s := newTestState()
	c := newCredentialAuthorityWithEngine(logger.New("test", libkb.ErrorWriter()), s, s)
	return s, c
}

func TestSimple(t *testing.T) {
	S, C := newTestSetup()
	u0 := S.newTestUser(4)

	key0 := u0.sibkeys[0]
	key1 := u0.sibkeys[1]

	if S.numGets != 0 {
		t.Fatal("expected 0 gets")
	}

	err := C.CheckUserKey(context.TODO(), u0.uid, &u0.username, &key0)
	if err != nil {
		t.Fatal(err)
	}
	if S.numGets != 1 {
		t.Fatal("expected 1 get")
	}
	err = C.CheckUserKey(context.TODO(), u0.uid, &u0.username, &key0)
	if err != nil {
		t.Fatal(err)
	}
	if S.numGets != 1 {
		t.Fatal("expected 1 get")
	}

	S.mutateUser(u0.uid, func(u *testUser) {
		u.sibkeys = u.sibkeys[1:]
	})

	// Advance just an iota, so that our polling of the server
	// has a chance to complete.
	S.tick(pollWait)

	// wait for the first eviction
	uid := <-S.evictCh
	if uid != u0.uid {
		t.Fatalf("Wrong UID on eviction: %s != %s\n", uid, u0.uid)
	}

	err = C.CheckUserKey(context.TODO(), u0.uid, &u0.username, &key0)
	if err == nil {
		t.Fatal("Expected an error")
	} else if bke, ok := err.(BadKeyError); !ok {
		t.Fatal("Expected a bad key error")
	} else if bke.uid != u0.uid {
		t.Fatalf("Expected a bad key error on %s (not %s)", u0.uid, bke.uid)
	} else if bke.kid != key0 {
		t.Fatalf("Expected a bad key error on key %s (not %s)", key0, bke.kid)
	}

	err = C.CheckUserKey(context.TODO(), u0.uid, &u0.username, &key1)
	if err != nil {
		t.Fatal(err)
	}
	if S.numGets != 2 {
		t.Fatal("expected 2 gets")
	}
	S.tick(userTimeout + time.Millisecond)
	err = C.CheckUserKey(context.TODO(), u0.uid, &u0.username, &key1)
	if err != nil {
		t.Fatal(err)
	}
	if S.numGets != 3 {
		t.Fatal("expected 3 gets")
	}
	S.tick(cacheTimeout + time.Millisecond)

	// u0 should now be gone since we haven't touched him in over cacheTimeout
	// duration.
	uid = <-S.evictCh
	if uid != u0.uid {
		t.Fatalf("Wrong UID on eviction: %s != %s\n", uid, u0.uid)
	}

	// Make a new user -- u1!
	u1 := S.newTestUser(4)

	ng := 3
	for i := 0; i < 10; i++ {
		err = C.CheckUserKey(context.TODO(), u1.uid, &u1.username, &u1.sibkeys[0])
		if err != nil {
			t.Fatal(err)
		}
		ng++
		if S.numGets != ng {
			t.Fatalf("expected %d gets, got %d", ng, S.numGets)
		}
		S.tick(userTimeout + time.Millisecond)

		select {
		case uid = <-S.evictCh:
			t.Fatalf("Got unwanted eviction for %s", uid)
		default:
		}
	}

	S.tick(cacheTimeout - userTimeout + 3*time.Millisecond)
	uid = <-S.evictCh
	if uid != u1.uid {
		t.Fatalf("Got wrong eviction: wanted %s but got %s\n", u1.uid, uid)
	}

	// Make a new user -- u2!
	u2 := S.newTestUser(4)
	err = C.CheckUserKey(context.TODO(), u2.uid, &u2.username, &u2.sibkeys[0])
	if err != nil {
		t.Fatal(err)
	}
	ng++
	if S.numGets != ng {
		t.Fatalf("expected %d gets, got %d", ng, S.numGets)
	}

	// Check that u2 is evicted properly after we shutdown the CA.
	C.Shutdown()
	uid = <-S.evictCh
	if uid != u2.uid {
		t.Fatalf("Got wrong eviction: wanted %s but got %s\n", u2.uid, uid)
	}

}

func TestCheckUsers(t *testing.T) {
	S, C := newTestSetup()

	var users, usersWithDud []keybase1.UID
	for i := 0; i < 10; i++ {
		u := S.newTestUser(2)
		users = append(users, u.uid)
		usersWithDud = append(usersWithDud, u.uid)
	}
	usersWithDud = append(usersWithDud, libkb.UsernameToUID(genUsername()))

	if S.numGets != 0 {
		t.Fatal("expected 0 gets")
	}

	err := C.CheckUsers(context.TODO(), users)
	if err != nil {
		t.Fatal(err)
	}
	if S.numGets != 10 {
		t.Fatal("expected 10 gets")
	}
	err = C.CheckUsers(context.TODO(), users)
	if err != nil {
		t.Fatal(err)
	}
	if S.numGets != 10 {
		t.Fatal("expected 10 gets")
	}

	err = C.CheckUsers(context.TODO(), usersWithDud)
	if err == nil {
		t.Fatal("Expected an error")
	} else if _, ok := err.(userNotFoundError); !ok {
		t.Fatal("Expected a user not found error")
	}
	C.Shutdown()
}

func TestCompareKeys(t *testing.T) {
	S, C := newTestSetup()
	u := S.newTestUser(10)

	err := C.CompareUserKeys(context.TODO(), u.uid, u.sibkeys, u.subkeys)
	if err != nil {
		t.Fatal(err)
	}

	err = C.CompareUserKeys(context.TODO(), u.uid, nil, u.subkeys)
	if err != nil {
		t.Fatal(err)
	}

	err = C.CompareUserKeys(context.TODO(), u.uid, u.sibkeys, nil)
	if err != nil {
		t.Fatal(err)
	}

	missingSibkey := u.sibkeys[1:]
	err = C.CompareUserKeys(context.TODO(), u.uid, missingSibkey, u.subkeys)
	if err == nil {
		t.Fatal("Expected an error")
	} else if _, ok := err.(KeysNotEqualError); !ok {
		t.Fatal("Expected keys not equal error")
	}

	missingSubkey := u.subkeys[1:]
	err = C.CompareUserKeys(context.TODO(), u.uid, u.sibkeys, missingSubkey)
	if err == nil {
		t.Fatal("Expected an error")
	} else if _, ok := err.(KeysNotEqualError); !ok {
		t.Fatal("Expected keys not equal error")
	}
	C.Shutdown()
}
<commit_msg>fix Go 1.6 vet errors<commit_after>package auth

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"
	libkb "github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol"
	context "golang.org/x/net/context"
	"sync"
	"testing"
	"time"
)

type testUser struct {
	uid      keybase1.UID
	username libkb.NormalizedUsername
	sibkeys  []keybase1.KID
	subkeys  []keybase1.KID
}

type testState struct {
	sync.Mutex

	users     map[keybase1.UID](*testUser)
	changes   []keybase1.UID
	now       time.Time
	evictCh   chan keybase1.UID
	pokeCh    chan struct{}
	startOnce sync.Once
	numGets   int
}

var seq uint32

func genKID() keybase1.KID {
	var kid [35]byte
	kid[0] = 0x1
	kid[1] = 0x20
	binary.BigEndian.PutUint32(kid[30:34], seq)
	seq++
	kid[34] = 0xa0
	return keybase1.KIDFromSlice(kid[:])
}

func genUsername() string {
	w, _ := libkb.SecWordList(1)
	var buf [4]byte
	rand.Read(buf[:])
	return fmt.Sprintf("%s%x", w[0], buf)
}

func newTestUser(nKeys int) *testUser {
	un := genUsername()
	ret := testUser{
		username: libkb.NewNormalizedUsername(un),
		uid:      libkb.UsernameToUID(un),
		sibkeys:  make([]keybase1.KID, nKeys),
		subkeys:  make([]keybase1.KID, nKeys),
	}
	for i := 0; i < nKeys; i++ {
		ret.sibkeys[i] = genKID()
		ret.subkeys[i] = genKID()
	}
	return &ret
}

func (ts *testState) newTestUser(nKeys int) *testUser {
	ts.Lock()
	defer ts.Unlock()
	ret := newTestUser(nKeys)
	ts.users[ret.uid] = ret
	return ret
}

func (ts *testState) mutateUser(uid keybase1.UID, mutator func(u *testUser)) bool {
	ts.Lock()
	defer ts.Unlock()
	u := ts.users[uid]
	if u == nil {
		return false
	}
	mutator(u)
	ts.changes = append(ts.changes, uid)
	return true
}

func newTestState() *testState {
	return &testState{
		users:   make(map[keybase1.UID](*testUser)),
		now:     time.Unix(100, 0),
		evictCh: make(chan keybase1.UID, 1),
		pokeCh:  make(chan struct{}),
	}
}

type userNotFoundError struct {
}

func (e userNotFoundError) Error() string {
	return "user not found"
}

func (ts *testState) GetUser(_ context.Context, uid keybase1.UID) (
	un libkb.NormalizedUsername, sibkeys, subkeys []keybase1.KID, err error) {
	ts.Lock()
	defer ts.Unlock()
	u := ts.users[uid]
	if u == nil {
		return libkb.NormalizedUsername(""), nil, nil, userNotFoundError{}
	}
	ts.numGets++
	return u.username, u.sibkeys, u.subkeys, nil
}

func (ts *testState) PollForChanges(_ context.Context) ([]keybase1.UID, error) {
	ts.Lock()
	defer ts.Unlock()
	ret := ts.changes
	ts.changes = nil
	return ret, nil
}

var _ UserKeyAPIer = (*testState)(nil)
var _ engine = (*testState)(nil)

func (ts *testState) tick(d time.Duration) {
	ts.pokeCh <- struct{}{}
	ts.Lock()
	ts.now = ts.now.Add(d)
	ts.Unlock()
	ts.pokeCh <- struct{}{}
}

func (ts *testState) Now() time.Time {
	ts.Lock()
	ret := ts.now
	ts.Unlock()
	return ret
}

func (ts *testState) GetPokeCh() <-chan struct{} { return ts.pokeCh }

func (ts *testState) Evicted(uid keybase1.UID) {
	ts.evictCh <- uid
}

func newTestSetup() (*testState, *CredentialAuthority) {
	s := newTestState()
	c := newCredentialAuthorityWithEngine(logger.New("test", libkb.ErrorWriter()), s, s)
	return s, c
}

func TestSimple(t *testing.T) {
	state, credentialAuthority := newTestSetup()
	u0 := state.newTestUser(4)

	key0 := u0.sibkeys[0]
	key1 := u0.sibkeys[1]

	if state.numGets != 0 {
		t.Fatal("expected 0 gets")
	}

	err := credentialAuthority.CheckUserKey(context.TODO(), u0.uid, &u0.username, &key0)
	if err != nil {
		t.Fatal(err)
	}
	if state.numGets != 1 {
		t.Fatal("expected 1 get")
	}
	err = credentialAuthority.CheckUserKey(context.TODO(), u0.uid, &u0.username, &key0)
	if err != nil {
		t.Fatal(err)
	}
	if state.numGets != 1 {
		t.Fatal("expected 1 get")
	}

	state.mutateUser(u0.uid, func(u *testUser) {
		u.sibkeys = u.sibkeys[1:]
	})

	// Advance just an iota, so that our polling of the server
	// has a chance to complete.
	state.tick(pollWait)

	// wait for the first eviction
	uid := <-state.evictCh
	if uid != u0.uid {
		t.Fatalf("Wrong UID on eviction: %s != %s\n", uid, u0.uid)
	}

	err = credentialAuthority.CheckUserKey(context.TODO(), u0.uid, &u0.username, &key0)
	if err == nil {
		t.Fatal("Expected an error")
	} else if bke, ok := err.(BadKeyError); !ok {
		t.Fatal("Expected a bad key error")
	} else if bke.uid != u0.uid {
		t.Fatalf("Expected a bad key error on %s (not %s)", u0.uid, bke.uid)
	} else if bke.kid != key0 {
		t.Fatalf("Expected a bad key error on key %s (not %s)", key0, bke.kid)
	}

	err = credentialAuthority.CheckUserKey(context.TODO(), u0.uid, &u0.username, &key1)
	if err != nil {
		t.Fatal(err)
	}
	if state.numGets != 2 {
		t.Fatal("expected 2 gets")
	}
	state.tick(userTimeout + time.Millisecond)
	err = credentialAuthority.CheckUserKey(context.TODO(), u0.uid, &u0.username, &key1)
	if err != nil {
		t.Fatal(err)
	}
	if state.numGets != 3 {
		t.Fatal("expected 3 gets")
	}
	state.tick(cacheTimeout + time.Millisecond)

	// u0 should now be gone since we haven't touched him in over cacheTimeout
	// duration.
	uid = <-state.evictCh
	if uid != u0.uid {
		t.Fatalf("Wrong UID on eviction: %s != %s\n", uid, u0.uid)
	}

	// Make a new user -- u1!
	u1 := state.newTestUser(4)

	ng := 3
	for i := 0; i < 10; i++ {
		err = credentialAuthority.CheckUserKey(context.TODO(), u1.uid, &u1.username, &u1.sibkeys[0])
		if err != nil {
			t.Fatal(err)
		}
		ng++
		if state.numGets != ng {
			t.Fatalf("expected %d gets, got %d", ng, state.numGets)
		}
		state.tick(userTimeout + time.Millisecond)

		select {
		case uid = <-state.evictCh:
			t.Fatalf("Got unwanted eviction for %s", uid)
		default:
		}
	}

	state.tick(cacheTimeout - userTimeout + 3*time.Millisecond)
	uid = <-state.evictCh
	if uid != u1.uid {
		t.Fatalf("Got wrong eviction: wanted %s but got %s\n", u1.uid, uid)
	}

	// Make a new user -- u2!
	u2 := state.newTestUser(4)
	err = credentialAuthority.CheckUserKey(context.TODO(), u2.uid, &u2.username, &u2.sibkeys[0])
	if err != nil {
		t.Fatal(err)
	}
	ng++
	if state.numGets != ng {
		t.Fatalf("expected %d gets, got %d", ng, state.numGets)
	}

	// Check that u2 is evicted properly after we shutdown the CA.
	credentialAuthority.Shutdown()
	uid = <-state.evictCh
	if uid != u2.uid {
		t.Fatalf("Got wrong eviction: wanted %s but got %s\n", u2.uid, uid)
	}

}

func TestCheckUsers(t *testing.T) {
	state, credentialAuthority := newTestSetup()

	var users, usersWithDud []keybase1.UID
	for i := 0; i < 10; i++ {
		u := state.newTestUser(2)
		users = append(users, u.uid)
		usersWithDud = append(usersWithDud, u.uid)
	}
	usersWithDud = append(usersWithDud, libkb.UsernameToUID(genUsername()))

	if state.numGets != 0 {
		t.Fatal("expected 0 gets")
	}

	err := credentialAuthority.CheckUsers(context.TODO(), users)
	if err != nil {
		t.Fatal(err)
	}
	if state.numGets != 10 {
		t.Fatal("expected 10 gets")
	}
	err = credentialAuthority.CheckUsers(context.TODO(), users)
	if err != nil {
		t.Fatal(err)
	}
	if state.numGets != 10 {
		t.Fatal("expected 10 gets")
	}

	err = credentialAuthority.CheckUsers(context.TODO(), usersWithDud)
	if err == nil {
		t.Fatal("Expected an error")
	} else if _, ok := err.(userNotFoundError); !ok {
		t.Fatal("Expected a user not found error")
	}
	credentialAuthority.Shutdown()
}

func TestCompareKeys(t *testing.T) {
	state, credentialAuthority := newTestSetup()
	u := state.newTestUser(10)

	err := credentialAuthority.CompareUserKeys(context.TODO(), u.uid, u.sibkeys, u.subkeys)
	if err != nil {
		t.Fatal(err)
	}

	err = credentialAuthority.CompareUserKeys(context.TODO(), u.uid, nil, u.subkeys)
	if err != nil {
		t.Fatal(err)
	}

	err = credentialAuthority.CompareUserKeys(context.TODO(), u.uid, u.sibkeys, nil)
	if err != nil {
		t.Fatal(err)
	}

	missingSibkey := u.sibkeys[1:]
	err = credentialAuthority.CompareUserKeys(context.TODO(), u.uid, missingSibkey, u.subkeys)
	if err == nil {
		t.Fatal("Expected an error")
	} else if _, ok := err.(KeysNotEqualError); !ok {
		t.Fatal("Expected keys not equal error")
	}

	missingSubkey := u.subkeys[1:]
	err = credentialAuthority.CompareUserKeys(context.TODO(), u.uid, u.sibkeys, missingSubkey)
	if err == nil {
		t.Fatal("Expected an error")
	} else if _, ok := err.(KeysNotEqualError); !ok {
		t.Fatal("Expected keys not equal error")
	}
	credentialAuthority.Shutdown()
}
<|endoftext|>


<commit_before>package main

import (
	// _ "expvar"

	"fmt"
	"koding/db/mongodb/modelhelper"
	"net/http"
	// _ "net/http/pprof" // Imported for side-effect of handling /debug/pprof.

	"socialapi/config"
	"socialapi/models"
	"socialapi/workers/api/handlers"
	"socialapi/workers/common/runner"
	"socialapi/workers/helper"
	notificationapi "socialapi/workers/notification/api"
	paymentapi "socialapi/workers/payment/api"
	"socialapi/workers/payment/stripe"
	sitemapapi "socialapi/workers/sitemap/api"
	trollmodeapi "socialapi/workers/trollmode/api"

	"github.com/rcrowley/go-tigertonic"
)

var (
	mux, nsMux *tigertonic.TrieServeMux

	Name = "SocialAPI"
)

func init() {
	mux = tigertonic.NewTrieServeMux()
	mux = notificationapi.InitHandlers(mux)
	mux = trollmodeapi.InitHandlers(mux)
	mux = sitemapapi.InitHandlers(mux)

	// add namespace support into
	// all handlers
	nsMux = tigertonic.NewTrieServeMux()
	nsMux.HandleNamespace("", mux)
	nsMux.HandleNamespace("/1.0", mux)
	tigertonic.SnakeCaseHTTPEquivErrors = true

}

func main() {
	r := runner.New(Name)
	if err := r.Init(); err != nil {
		fmt.Println(err)
		return
	}

	server := newServer(r)
	// shutdown server
	defer server.Close()

	mux = handlers.Inject(mux, r.Metrics)

	// init payment handlers, this done here instead of in `init()`
	// like others so we can've access to `metrics`
	mux = paymentapi.InitHandlers(mux, r.Metrics)

	mux.HandleFunc("GET", "/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "Hello from socialapi")
	})

	// init redis
	redisConn := helper.MustInitRedisConn(r.Conf)
	defer redisConn.Close()

	// init mongo connection
	modelhelper.Initialize(r.Conf.Mongo)

	// init stripe client
	stripe.InitializeClientKey(config.MustGet().Stripe.SecretToken)

	go func() {
		err := stripe.CreateDefaultPlans()
		if err != nil {
			fmt.Println(err)
			panic(err)
		}
	}()

	r.Listen()
	r.Wait()
}

func newServer(r *runner.Runner) *tigertonic.Server {
	// go metrics.Log(
	// 	metrics.DefaultRegistry,
	// 	60e9,
	// 	stdlog.New(os.Stderr, "metrics ", stdlog.Lmicroseconds),
	// )

	conf := r.Conf

	var handler http.Handler
	handler = tigertonic.WithContext(nsMux, models.Context{})
	if conf.Debug {
		h := tigertonic.Logged(handler, nil)
		h.Logger = NewTigerTonicLogger(r.Log)
		handler = h
	}

	addr := conf.Host + ":" + conf.Port
	server := tigertonic.NewServer(addr, handler)

	go listener(server)
	return server
}

func listener(server *tigertonic.Server) {
	if err := server.ListenAndServe(); nil != err {
		panic(err)
	}
}
<commit_msg>Social: create default values for socialapi in a go routine, do not block server start<commit_after>package main

import (
	// _ "expvar"

	"fmt"
	"koding/db/mongodb/modelhelper"
	"net/http"
	// _ "net/http/pprof" // Imported for side-effect of handling /debug/pprof.

	"socialapi/config"
	"socialapi/models"
	"socialapi/workers/api/handlers"
	"socialapi/workers/common/runner"
	"socialapi/workers/helper"
	notificationapi "socialapi/workers/notification/api"
	paymentapi "socialapi/workers/payment/api"
	"socialapi/workers/payment/stripe"
	sitemapapi "socialapi/workers/sitemap/api"
	trollmodeapi "socialapi/workers/trollmode/api"

	"github.com/rcrowley/go-tigertonic"
)

var (
	mux, nsMux *tigertonic.TrieServeMux

	Name = "SocialAPI"
)

func init() {
	mux = tigertonic.NewTrieServeMux()
	mux = notificationapi.InitHandlers(mux)
	mux = trollmodeapi.InitHandlers(mux)
	mux = sitemapapi.InitHandlers(mux)

	// add namespace support into
	// all handlers
	nsMux = tigertonic.NewTrieServeMux()
	nsMux.HandleNamespace("", mux)
	nsMux.HandleNamespace("/1.0", mux)
	tigertonic.SnakeCaseHTTPEquivErrors = true

}

func main() {
	r := runner.New(Name)
	if err := r.Init(); err != nil {
		fmt.Println(err)
		return
	}

	server := newServer(r)
	// shutdown server
	defer server.Close()

	mux = handlers.Inject(mux, r.Metrics)

	// init payment handlers, this done here instead of in `init()`
	// like others so we can've access to `metrics`
	mux = paymentapi.InitHandlers(mux, r.Metrics)

	mux.HandleFunc("GET", "/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "Hello from socialapi")
	})

	// init redis
	redisConn := helper.MustInitRedisConn(r.Conf)
	defer redisConn.Close()

	// init mongo connection
	modelhelper.Initialize(r.Conf.Mongo)

	// init stripe client
	stripe.InitializeClientKey(config.MustGet().Stripe.SecretToken)

	go func() {
		err := stripe.CreateDefaultPlans()
		if err != nil {
			fmt.Println(err)
			panic(err)
		}
	}()

	go setDefaults(r.Log)

	r.Listen()
	r.Wait()
}

func newServer(r *runner.Runner) *tigertonic.Server {
	// go metrics.Log(
	// 	metrics.DefaultRegistry,
	// 	60e9,
	// 	stdlog.New(os.Stderr, "metrics ", stdlog.Lmicroseconds),
	// )

	conf := r.Conf

	var handler http.Handler
	handler = tigertonic.WithContext(nsMux, models.Context{})
	if conf.Debug {
		h := tigertonic.Logged(handler, nil)
		h.Logger = NewTigerTonicLogger(r.Log)
		handler = h
	}

	addr := conf.Host + ":" + conf.Port
	server := tigertonic.NewServer(addr, handler)

	go listener(server)
	return server
}

func listener(server *tigertonic.Server) {
	if err := server.ListenAndServe(); nil != err {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package parser

import (
	"github.com/Jeffail/benthos/v3/internal/bloblang/query"
)

func matchCaseParser() Func {
	whitespace := SpacesAndTabs()

	p := Sequence(
		OneOf(
			Sequence(
				Expect(
					Char('_'),
					"match case",
				),
				Optional(whitespace),
				Term("=>"),
			),
			Sequence(
				Expect(
					ParseQuery,
					"match case",
				),
				Optional(whitespace),
				Term("=>"),
			),
		),
		Optional(whitespace),
		ParseQuery,
	)

	return func(input []rune) Result {
		res := p(input)
		if res.Err != nil {
			return res
		}

		seqSlice := res.Payload.([]interface{})

		var caseFn query.Function
		switch t := seqSlice[0].([]interface{})[0].(type) {
		case query.Function:
			if lit, isLiteral := t.(*query.Literal); isLiteral {
				caseFn = query.ClosureFunction(func(ctx query.FunctionContext) (interface{}, error) {
					v := ctx.Value()
					if v == nil {
						return false, nil
					}
					return *v == lit.Value, nil
				}, nil)
			} else {
				caseFn = t
			}
		case string:
			caseFn = query.NewLiteralFunction(true)
		}

		return Success(
			query.NewMatchCase(caseFn, seqSlice[2].(query.Function)),
			res.Remaining,
		)
	}
}

func matchExpressionParser() Func {
	whitespace := DiscardAll(
		OneOf(
			SpacesAndTabs(),
			NewlineAllowComment(),
		),
	)

	return func(input []rune) Result {
		res := Sequence(
			Term("match"),
			SpacesAndTabs(),
			Optional(ParseQuery),
			whitespace,
			MustBe(
				DelimitedPattern(
					Sequence(
						Char('{'),
						whitespace,
					),
					matchCaseParser(),
					Sequence(
						Discard(SpacesAndTabs()),
						OneOf(
							Char(','),
							NewlineAllowComment(),
						),
						whitespace,
					),
					Sequence(
						whitespace,
						Char('}'),
					),
					true,
				),
			),
		)(input)
		if res.Err != nil {
			return res
		}

		seqSlice := res.Payload.([]interface{})
		contextFn, _ := seqSlice[2].(query.Function)

		cases := []query.MatchCase{}
		for _, caseVal := range seqSlice[4].([]interface{}) {
			cases = append(cases, caseVal.(query.MatchCase))
		}

		res.Payload = query.NewMatchFunction(contextFn, cases...)
		return res
	}
}

func ifExpressionParser() Func {
	optionalWhitespace := DiscardAll(
		OneOf(
			SpacesAndTabs(),
			NewlineAllowComment(),
		),
	)

	return func(input []rune) Result {
		res := Sequence(
			Term("if"),
			SpacesAndTabs(),
			ParseQuery,
			optionalWhitespace,
			Char('{'),
			optionalWhitespace,
			ParseQuery,
			optionalWhitespace,
			Char('}'),
			Optional(
				Sequence(
					optionalWhitespace,
					Term("else"),
					optionalWhitespace,
					Char('{'),
					optionalWhitespace,
					ParseQuery,
					optionalWhitespace,
					Char('}'),
				),
			),
		)(input)
		if res.Err != nil {
			return res
		}

		seqSlice := res.Payload.([]interface{})
		queryFn := seqSlice[2].(query.Function)
		ifFn := seqSlice[6].(query.Function)

		var elseFn query.Function
		elseSlice, _ := seqSlice[9].([]interface{})
		if len(elseSlice) > 0 {
			elseFn, _ = elseSlice[5].(query.Function)
		}

		res.Payload = query.NewIfFunction(queryFn, ifFn, elseFn)
		return res
	}
}

func bracketsExpressionParser() Func {
	whitespace := DiscardAll(
		OneOf(
			SpacesAndTabs(),
			NewlineAllowComment(),
		),
	)
	return func(input []rune) Result {
		res := Sequence(
			Expect(
				Char('('),
				"function",
			),
			whitespace,
			ParseQuery,
			whitespace,
			MustBe(Expect(Char(')'), "closing bracket")),
		)(input)
		if res.Err != nil {
			return res
		}
		res.Payload = res.Payload.([]interface{})[2]
		return res
	}
}
<commit_msg>Enforce partial if statement body errors<commit_after>package parser

import (
	"github.com/Jeffail/benthos/v3/internal/bloblang/query"
)

func matchCaseParser() Func {
	whitespace := SpacesAndTabs()

	p := Sequence(
		OneOf(
			Sequence(
				Expect(
					Char('_'),
					"match case",
				),
				Optional(whitespace),
				Term("=>"),
			),
			Sequence(
				Expect(
					ParseQuery,
					"match case",
				),
				Optional(whitespace),
				Term("=>"),
			),
		),
		Optional(whitespace),
		ParseQuery,
	)

	return func(input []rune) Result {
		res := p(input)
		if res.Err != nil {
			return res
		}

		seqSlice := res.Payload.([]interface{})

		var caseFn query.Function
		switch t := seqSlice[0].([]interface{})[0].(type) {
		case query.Function:
			if lit, isLiteral := t.(*query.Literal); isLiteral {
				caseFn = query.ClosureFunction(func(ctx query.FunctionContext) (interface{}, error) {
					v := ctx.Value()
					if v == nil {
						return false, nil
					}
					return *v == lit.Value, nil
				}, nil)
			} else {
				caseFn = t
			}
		case string:
			caseFn = query.NewLiteralFunction(true)
		}

		return Success(
			query.NewMatchCase(caseFn, seqSlice[2].(query.Function)),
			res.Remaining,
		)
	}
}

func matchExpressionParser() Func {
	whitespace := DiscardAll(
		OneOf(
			SpacesAndTabs(),
			NewlineAllowComment(),
		),
	)

	return func(input []rune) Result {
		res := Sequence(
			Term("match"),
			SpacesAndTabs(),
			Optional(ParseQuery),
			whitespace,
			MustBe(
				DelimitedPattern(
					Sequence(
						Char('{'),
						whitespace,
					),
					matchCaseParser(),
					Sequence(
						Discard(SpacesAndTabs()),
						OneOf(
							Char(','),
							NewlineAllowComment(),
						),
						whitespace,
					),
					Sequence(
						whitespace,
						Char('}'),
					),
					true,
				),
			),
		)(input)
		if res.Err != nil {
			return res
		}

		seqSlice := res.Payload.([]interface{})
		contextFn, _ := seqSlice[2].(query.Function)

		cases := []query.MatchCase{}
		for _, caseVal := range seqSlice[4].([]interface{}) {
			cases = append(cases, caseVal.(query.MatchCase))
		}

		res.Payload = query.NewMatchFunction(contextFn, cases...)
		return res
	}
}

func ifExpressionParser() Func {
	optionalWhitespace := DiscardAll(
		OneOf(
			SpacesAndTabs(),
			NewlineAllowComment(),
		),
	)

	return func(input []rune) Result {
		res := Sequence(
			Term("if"),
			SpacesAndTabs(),
			MustBe(ParseQuery),
			optionalWhitespace,
			MustBe(Char('{')),
			optionalWhitespace,
			MustBe(ParseQuery),
			optionalWhitespace,
			MustBe(Char('}')),
			Optional(
				Sequence(
					optionalWhitespace,
					Term("else"),
					optionalWhitespace,
					MustBe(Char('{')),
					optionalWhitespace,
					MustBe(ParseQuery),
					optionalWhitespace,
					MustBe(Char('}')),
				),
			),
		)(input)
		if res.Err != nil {
			return res
		}

		seqSlice := res.Payload.([]interface{})
		queryFn := seqSlice[2].(query.Function)
		ifFn := seqSlice[6].(query.Function)

		var elseFn query.Function
		elseSlice, _ := seqSlice[9].([]interface{})
		if len(elseSlice) > 0 {
			elseFn, _ = elseSlice[5].(query.Function)
		}

		res.Payload = query.NewIfFunction(queryFn, ifFn, elseFn)
		return res
	}
}

func bracketsExpressionParser() Func {
	whitespace := DiscardAll(
		OneOf(
			SpacesAndTabs(),
			NewlineAllowComment(),
		),
	)
	return func(input []rune) Result {
		res := Sequence(
			Expect(
				Char('('),
				"function",
			),
			whitespace,
			ParseQuery,
			whitespace,
			MustBe(Expect(Char(')'), "closing bracket")),
		)(input)
		if res.Err != nil {
			return res
		}
		res.Payload = res.Payload.([]interface{})[2]
		return res
	}
}
<|endoftext|>


<commit_before>package provider

import (
	"fmt"
	"os"
	"regexp"
	"testing"

	"github.com/hashicorp/terraform-plugin-sdk/helper/acctest"
	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
	"github.com/hashicorp/terraform-plugin-sdk/terraform"
	"github.com/mackerelio/mackerel-client-go"
)

func TestAccMackerelChannelEmail_Basic(t *testing.T) {
	resourceName := "mackerel_channel.email"
	rName := acctest.RandomWithPrefix("TerraformTestChannelEmail-")

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckMackerelChannelDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccMackerelChannelEmailConfigBasic(rName),
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceName, "name", rName),
					resource.TestCheckResourceAttr(resourceName, "type", "email"),
					resource.TestCheckResourceAttrSet(resourceName, "emails.0"),
					resource.TestCheckResourceAttrSet(resourceName, "events.0"),
					resource.TestCheckResourceAttrSet(resourceName, "user_ids.0"),
				),
			},
		},
	})
}

func TestAccMackerelChannelEmail_Invalid(t *testing.T) {
	rName := acctest.RandomWithPrefix("TerraformTestChannelEmail-")

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckMackerelChannelDestroy,
		Steps: []resource.TestStep{
			{
				Config:      testAccMackerelChannelEmailConfigInvalid(rName),
				ExpectError: regexp.MustCompile("API request failed: invalid userIds"),
			},
		},
	})
}

func TestAccMackerelChannelSlack_Basic(t *testing.T) {
	resourceName := "mackerel_channel.slack"
	rName := acctest.RandomWithPrefix("TerraformTestChannelSlack-")

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckMackerelChannelDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccMackerelChannelSlackConfigBasic(rName),
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceName, "name", rName),
					resource.TestCheckResourceAttr(resourceName, "type", "slack"),
					resource.TestCheckResourceAttrSet(resourceName, "events.0"),
					resource.TestCheckResourceAttr(resourceName, "mentions.%", "2"),
					resource.TestCheckResourceAttr(resourceName, "mentions.ok", "ok"),
					resource.TestCheckResourceAttr(resourceName, "mentions.critical", "critical"),
					resource.TestCheckResourceAttr(resourceName, "enabled_graph_image", "true"),
				),
			},
		},
	})
}

func TestAccMackerelChannelWebhook_Basic(t *testing.T) {
	resourceName := "mackerel_channel.webhook"
	rName := acctest.RandomWithPrefix("TerraformTestChannelWebhook-")

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckMackerelChannelDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccMackerelChannelWebhookConfigBasic(rName),
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceName, "name", rName),
					resource.TestCheckResourceAttr(resourceName, "type", "webhook"),
					resource.TestCheckResourceAttrSet(resourceName, "events.0"),
					resource.TestCheckResourceAttrSet(resourceName, "events.1"),
					resource.TestCheckResourceAttr(resourceName, "url", "https://hogehoge.com"),
					resource.TestCheckResourceAttr(resourceName, "enabled_graph_image", "true"),
				),
			},
		},
	})
}

func testAccCheckMackerelChannelDestroy(s *terraform.State) error {
	client := testAccProvider.Meta().(*mackerel.Client)

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "mackerel_channel" {
			continue
		}

		channels, err := client.FindChannels()
		if err != nil {
			return fmt.Errorf("find channel failed")
		}
		for _, chn := range channels {
			if rs.Primary.ID == chn.ID {
				return fmt.Errorf("channel still exists")
			}
		}
	}

	return nil
}

func testAccCheckMackerelChannelExists(resouceName string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		client := testAccProvider.Meta().(*mackerel.Client)

		for _, rs := range s.RootModule().Resources {
			if rs.Type != "mackerel_channel" {
				continue
			}

			channels, err := client.FindChannels()
			if err != nil {
				return fmt.Errorf("find channel failed")
			}
			for _, chn := range channels {
				if rs.Primary.ID == chn.ID {
					return nil
				}
			}
		}
		return fmt.Errorf("channel (%s) not found", resouceName)
	}
}

func testAccMackerelChannelEmailConfigBasic(rName string) string {
	return fmt.Sprintf(`
resource "mackerel_channel" "email" {
  name    = "%s"
  type    = "email"
  emails  = ["foo@exapmle.com","bar@exapmle.com"]
  events  = ["alert"]
  user_ids = ["%s"]
}
`, rName, os.Getenv("USER_ID"))
}

func testAccMackerelChannelEmailConfigInvalid(rName string) string {
	return fmt.Sprintf(`
resource "mackerel_channel" "email" {
  name    = "%s"
  type    = "email"
  emails  = ["foo@exapmle.com","bar@exapmle.com"]
  events  = ["alert"]
  user_ids = ["hoge"]
}
`, rName)
}

func testAccMackerelChannelSlackConfigBasic(rName string) string {
	return fmt.Sprintf(`
resource "mackerel_channel" "slack" {
  name    = "%s"
  type    = "slack"
  events  = ["alert"]
  url = "https://hooks.slack.com/services/"
  mentions = {
    "ok": "ok",
    "critical": "critical",
  }
  enabled_graph_image = true
}
`, rName)
}

func testAccMackerelChannelWebhookConfigBasic(rName string) string {
	return fmt.Sprintf(`
resource "mackerel_channel" "webhook" {
  name    = "%s"
  type    = "webhook"
  events  = ["alert", "monitor"]
  url = "https://hogehoge.com"
  enabled_graph_image = true
}
`, rName)
}
<commit_msg>Fix linter<commit_after>package provider

import (
	"fmt"
	"os"
	"regexp"
	"testing"

	"github.com/hashicorp/terraform-plugin-sdk/helper/acctest"
	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
	"github.com/hashicorp/terraform-plugin-sdk/terraform"
	"github.com/mackerelio/mackerel-client-go"
)

func TestAccMackerelChannelEmail_Basic(t *testing.T) {
	resourceName := "mackerel_channel.email"
	rName := acctest.RandomWithPrefix("TerraformTestChannelEmail-")

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckMackerelChannelDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccMackerelChannelEmailConfigBasic(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckMackerelChannelExists(resourceName),
					resource.TestCheckResourceAttr(resourceName, "name", rName),
					resource.TestCheckResourceAttr(resourceName, "type", "email"),
					resource.TestCheckResourceAttrSet(resourceName, "emails.0"),
					resource.TestCheckResourceAttrSet(resourceName, "events.0"),
					resource.TestCheckResourceAttrSet(resourceName, "user_ids.0"),
				),
			},
		},
	})
}

func TestAccMackerelChannelEmail_Invalid(t *testing.T) {
	rName := acctest.RandomWithPrefix("TerraformTestChannelEmail-")

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckMackerelChannelDestroy,
		Steps: []resource.TestStep{
			{
				Config:      testAccMackerelChannelEmailConfigInvalid(rName),
				ExpectError: regexp.MustCompile("API request failed: invalid userIds"),
			},
		},
	})
}

func TestAccMackerelChannelSlack_Basic(t *testing.T) {
	resourceName := "mackerel_channel.slack"
	rName := acctest.RandomWithPrefix("TerraformTestChannelSlack-")

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckMackerelChannelDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccMackerelChannelSlackConfigBasic(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckMackerelChannelExists(resourceName),
					resource.TestCheckResourceAttr(resourceName, "name", rName),
					resource.TestCheckResourceAttr(resourceName, "type", "slack"),
					resource.TestCheckResourceAttrSet(resourceName, "events.0"),
					resource.TestCheckResourceAttr(resourceName, "mentions.%", "2"),
					resource.TestCheckResourceAttr(resourceName, "mentions.ok", "ok"),
					resource.TestCheckResourceAttr(resourceName, "mentions.critical", "critical"),
					resource.TestCheckResourceAttr(resourceName, "enabled_graph_image", "true"),
				),
			},
		},
	})
}

func TestAccMackerelChannelWebhook_Basic(t *testing.T) {
	resourceName := "mackerel_channel.webhook"
	rName := acctest.RandomWithPrefix("TerraformTestChannelWebhook-")

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckMackerelChannelDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccMackerelChannelWebhookConfigBasic(rName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckMackerelChannelExists(resourceName),
					resource.TestCheckResourceAttr(resourceName, "name", rName),
					resource.TestCheckResourceAttr(resourceName, "type", "webhook"),
					resource.TestCheckResourceAttrSet(resourceName, "events.0"),
					resource.TestCheckResourceAttrSet(resourceName, "events.1"),
					resource.TestCheckResourceAttr(resourceName, "url", "https://hogehoge.com"),
					resource.TestCheckResourceAttr(resourceName, "enabled_graph_image", "true"),
				),
			},
		},
	})
}

func testAccCheckMackerelChannelDestroy(s *terraform.State) error {
	client := testAccProvider.Meta().(*mackerel.Client)

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "mackerel_channel" {
			continue
		}

		channels, err := client.FindChannels()
		if err != nil {
			return fmt.Errorf("find channel failed")
		}
		for _, chn := range channels {
			if rs.Primary.ID == chn.ID {
				return fmt.Errorf("channel still exists")
			}
		}
	}

	return nil
}

func testAccCheckMackerelChannelExists(resouceName string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		client := testAccProvider.Meta().(*mackerel.Client)

		for _, rs := range s.RootModule().Resources {
			if rs.Type != "mackerel_channel" {
				continue
			}

			channels, err := client.FindChannels()
			if err != nil {
				return fmt.Errorf("find channel failed")
			}
			for _, chn := range channels {
				if rs.Primary.ID == chn.ID {
					return nil
				}
			}
		}
		return fmt.Errorf("channel (%s) not found", resouceName)
	}
}

func testAccMackerelChannelEmailConfigBasic(rName string) string {
	return fmt.Sprintf(`
resource "mackerel_channel" "email" {
  name    = "%s"
  type    = "email"
  emails  = ["foo@exapmle.com","bar@exapmle.com"]
  events  = ["alert"]
  user_ids = ["%s"]
}
`, rName, os.Getenv("USER_ID"))
}

func testAccMackerelChannelEmailConfigInvalid(rName string) string {
	return fmt.Sprintf(`
resource "mackerel_channel" "email" {
  name    = "%s"
  type    = "email"
  emails  = ["foo@exapmle.com","bar@exapmle.com"]
  events  = ["alert"]
  user_ids = ["hoge"]
}
`, rName)
}

func testAccMackerelChannelSlackConfigBasic(rName string) string {
	return fmt.Sprintf(`
resource "mackerel_channel" "slack" {
  name    = "%s"
  type    = "slack"
  events  = ["alert"]
  url = "https://hooks.slack.com/services/"
  mentions = {
    "ok": "ok",
    "critical": "critical",
  }
  enabled_graph_image = true
}
`, rName)
}

func testAccMackerelChannelWebhookConfigBasic(rName string) string {
	return fmt.Sprintf(`
resource "mackerel_channel" "webhook" {
  name    = "%s"
  type    = "webhook"
  events  = ["alert", "monitor"]
  url = "https://hogehoge.com"
  enabled_graph_image = true
}
`, rName)
}
<|endoftext|>


<commit_before>package keylogger

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"syscall"

	"github.com/sirupsen/logrus"
)

// KeyLogger wrapper around file descriptior
type KeyLogger struct {
	fd *os.File
}

// New creates a new keylogger for a device path
func New(devPath string) (*KeyLogger, error) {
	k := &KeyLogger{}
	if !k.IsRoot() {
		return nil, errors.New("Must be run as root")
	}
	fd, err := os.Open(devPath)
	k.fd = fd
	return k, err
}

// FindKeyboardDevice by going through each device registered on OS
// Mostly it will contain keyword - keyboard
// Returns the file path which contains events
func FindKeyboardDevice() string {
	path := "/sys/class/input/event%d/device/name"
	resolved := "/dev/input/event%d"

	for i := 0; i < 255; i++ {
		buff, err := ioutil.ReadFile(fmt.Sprintf(path, i))
		if err != nil {
			logrus.Error(err)
		}

		// check if mouse is contained in the input event
		// if that is the case just skip.
		// We do this check as it seems that some mouses like the logitech MX mouse is also recognized as a mouse/keyboard
		if strings.Contains(strings.ToLower(string(buff)), "mouse") {
			continue
		}

		if strings.Contains(strings.ToLower(string(buff)), "keyboard") {
			return fmt.Sprintf(resolved, i)
		}
	}
	return ""
}

// Like FindKeyboardDevice, but finds all devices which contain keyword 'keyboard'
// Returns an array of file paths which contain keyboard events
func FindAllKeyboardDevices() []string {
	path := "/sys/class/input/event%d/device/name"
	resolved := "/dev/input/event%d"

	valid := make([]string, 0)

	for i := 0; i < 255; i++ {
		buff, err := ioutil.ReadFile(fmt.Sprintf(path, i))

		// prevent from checking non-existant files
		if os.IsNotExist(err) {
			break
		}
		if err != nil {
			logrus.Error(err)
		}

		// check if mouse is contained in the input event
		// if that is the case just skip.
		// We do this check as it seems that some mouses like the logitech MX mouse is also recognized as a mouse/keyboard
		if strings.Contains(strings.ToLower(string(buff)), "mouse") {
			continue
		}

		if strings.Contains(strings.ToLower(string(buff)), "keyboard") {
			valid = append(valid, fmt.Sprintf(resolved, i))
		}
	}
	return valid
}

// IsRoot checks if the process is run with root permission
func (k *KeyLogger) IsRoot() bool {
	return syscall.Getuid() == 0 && syscall.Geteuid() == 0
}

// Read from file descriptor
// Blocking call, returns channel
// Make sure to close channel when finish
func (k *KeyLogger) Read() chan InputEvent {
	event := make(chan InputEvent)
	go func(event chan InputEvent) {
		for {
			e, err := k.read()
			if err != nil {
				logrus.Error(err)
				close(event)
				break
			}

			if e != nil {
				event <- *e
			}
		}
	}(event)
	return event
}

// read from file description and parse binary into go struct
func (k *KeyLogger) read() (*InputEvent, error) {
	buffer := make([]byte, eventsize)
	n, err := k.fd.Read(buffer)
	if err != nil {
		return nil, err
	}
	// no input, dont send error
	if n <= 0 {
		return nil, nil
	}
	return k.eventFromBuffer(buffer)
}

// eventFromBuffer parser bytes into InputEvent struct
func (k *KeyLogger) eventFromBuffer(buffer []byte) (*InputEvent, error) {
	event := &InputEvent{}
	err := binary.Read(bytes.NewBuffer(buffer), binary.LittleEndian, event)
	return event, err
}

// Close file descriptor
func (k *KeyLogger) Close() error {
	if k.fd == nil {
		return nil
	}
	return k.fd.Close()
}
<commit_msg>Add Support for logitech mx keys<commit_after>package keylogger

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"syscall"

	"github.com/sirupsen/logrus"
)

// KeyLogger wrapper around file descriptior
type KeyLogger struct {
	fd *os.File
}

// New creates a new keylogger for a device path
func New(devPath string) (*KeyLogger, error) {
	k := &KeyLogger{}
	if !k.IsRoot() {
		return nil, errors.New("Must be run as root")
	}
	fd, err := os.Open(devPath)
	k.fd = fd
	return k, err
}

// FindKeyboardDevice by going through each device registered on OS
// Mostly it will contain keyword - keyboard
// Returns the file path which contains events
func FindKeyboardDevice() string {
	path := "/sys/class/input/event%d/device/name"
	resolved := "/dev/input/event%d"

	for i := 0; i < 255; i++ {
		buff, err := ioutil.ReadFile(fmt.Sprintf(path, i))
		if err != nil {
			logrus.Error(err)
		}

		stringifiedBuffer := string(buff)

		// check if mouse is contained in the input event
		// if that is the case just skip.
		// We do this check as it seems that some mouses like the logitech MX mouse is also recognized as a mouse/keyboard
		if strings.Contains(strings.ToLower(stringifiedBuffer), "mouse") {
			continue
		}

		if strings.Contains(strings.ToLower(stringifiedBuffer), "keyboard") || strings.Contains(stringifiedBuffer, "Logitech MX Keys") {
			return fmt.Sprintf(resolved, i)
		}
	}
	return ""
}

// Like FindKeyboardDevice, but finds all devices which contain keyword 'keyboard'
// Returns an array of file paths which contain keyboard events
func FindAllKeyboardDevices() []string {
	path := "/sys/class/input/event%d/device/name"
	resolved := "/dev/input/event%d"

	valid := make([]string, 0)

	for i := 0; i < 255; i++ {
		buff, err := ioutil.ReadFile(fmt.Sprintf(path, i))

		// prevent from checking non-existant files
		if os.IsNotExist(err) {
			break
		}
		if err != nil {
			logrus.Error(err)
		}

		// check if mouse is contained in the input event
		// if that is the case just skip.
		// We do this check as it seems that some mouses like the logitech MX mouse is also recognized as a mouse/keyboard
		if strings.Contains(strings.ToLower(string(buff)), "mouse") {
			continue
		}

		if strings.Contains(strings.ToLower(string(buff)), "keyboard") {
			valid = append(valid, fmt.Sprintf(resolved, i))
		}
	}
	return valid
}

// IsRoot checks if the process is run with root permission
func (k *KeyLogger) IsRoot() bool {
	return syscall.Getuid() == 0 && syscall.Geteuid() == 0
}

// Read from file descriptor
// Blocking call, returns channel
// Make sure to close channel when finish
func (k *KeyLogger) Read() chan InputEvent {
	event := make(chan InputEvent)
	go func(event chan InputEvent) {
		for {
			e, err := k.read()
			if err != nil {
				logrus.Error(err)
				close(event)
				break
			}

			if e != nil {
				event <- *e
			}
		}
	}(event)
	return event
}

// read from file description and parse binary into go struct
func (k *KeyLogger) read() (*InputEvent, error) {
	buffer := make([]byte, eventsize)
	n, err := k.fd.Read(buffer)
	if err != nil {
		return nil, err
	}
	// no input, dont send error
	if n <= 0 {
		return nil, nil
	}
	return k.eventFromBuffer(buffer)
}

// eventFromBuffer parser bytes into InputEvent struct
func (k *KeyLogger) eventFromBuffer(buffer []byte) (*InputEvent, error) {
	event := &InputEvent{}
	err := binary.Read(bytes.NewBuffer(buffer), binary.LittleEndian, event)
	return event, err
}

// Close file descriptor
func (k *KeyLogger) Close() error {
	if k.fd == nil {
		return nil
	}
	return k.fd.Close()
}
<|endoftext|>


<commit_before>package kafka

import (
	"errors"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"

	"github.com/Shopify/sarama"
	"github.com/samuel/go-zookeeper/zk"
)

var ErrNoBrokers = errors.New("No kafka brokers found")
var ErrAlreadyListening = errors.New("Already listening to stream")

type Message struct {
	Key       []byte
	Val       []byte
	Topic     string
	Partition int32
	Offset    int64
}

type Stream struct {
	Topic     string
	Partition int32
}

type Kafka struct {
	isclosed    int32
	logger      *log.Logger
	zkpeers     []string
	zooConn     *zk.Conn
	kfkConn     sarama.Client
	kfkConsumer sarama.Consumer
	kfkProducer sarama.SyncProducer
	lock        sync.Mutex
	consumers   map[Stream]*partConsumer
	messagebus  chan Message
}

func New(logger *log.Logger, zkpeers []string) (*Kafka, error) {
	k := new(Kafka)
	k.consumers = make(map[Stream]*partConsumer)
	k.messagebus = make(chan Message, 50)
	k.logger = logger
	k.zkpeers = zkpeers
	err := k.connect()
	if err != nil {
		return nil, err
	}
	return k, nil
}

func (k *Kafka) connect() error {
	zooConn, _, err := zk.Connect(k.zkpeers, time.Second)
	if err != nil {
		return fmt.Errorf("Failed to connect to zookeeper: %v", err)
	}
	k.zooConn = zooConn

	kafkaBrokers, err := k.GetKafkaBrokersFromZookeeper()
	if err != nil {
		k.Close()
		return fmt.Errorf("Failed to feth brokers from zookeeper: %v", err)
	}
	if len(kafkaBrokers) == 0 {
		k.Close()
		return ErrNoBrokers
	}

	config := sarama.NewConfig()
	config.Producer.Partitioner = sarama.NewRandomPartitioner
	config.Producer.RequiredAcks = sarama.WaitForAll
	config.ClientID = "kafka-go"

	kfkConn, err := sarama.NewClient(brokerStrings(kafkaBrokers), config)
	if err != nil {
		k.Close()
		return fmt.Errorf("Failed to connect to Kafka: %v", err)
	}
	k.kfkConn = kfkConn

	consumer, err := sarama.NewConsumerFromClient(kfkConn)
	if err != nil {
		k.Close()
		return fmt.Errorf("Failed to create Kafka consumer: %v", err)
	}
	k.kfkConsumer = consumer

	producer, err := sarama.NewSyncProducerFromClient(kfkConn)
	if err != nil {
		consumer.Close()
		return fmt.Errorf("Failed to create sarama syncproducer: %v", err)
	}
	k.kfkProducer = producer

	return nil
}

func (k *Kafka) Close() error {
	if atomic.CompareAndSwapInt32(&k.isclosed, 0, 1) {
		if k.kfkProducer != nil {
			k.kfkProducer.Close()
		}
		if k.kfkConsumer != nil {
			k.kfkConsumer.Close()
		}
		if k.kfkConn != nil {
			k.kfkConn.Close()
		}
		if k.zooConn != nil {
			k.zooConn.Close()
		}
	}
	return nil
}

type partConsumer struct {
	conn   sarama.PartitionConsumer
	offset int64
	stream Stream
	closed chan struct{}
}

func (pc *partConsumer) Close() {
	//TODO: Do we want asynch close, or close to stop the messages?
	pc.conn.Close()
}

func (k *Kafka) Listen(topic string, partition int32, offset int64) error {
	stream := Stream{Topic: topic, Partition: partition}

	// Check if the stream was already added
	k.lock.Lock()
	_, ok := k.consumers[stream]
	k.lock.Unlock()
	if ok {
		return ErrAlreadyListening
	}

	// Create parititon consumer
	kfkPartConsumer, err := k.kfkConsumer.ConsumePartition(topic, partition, offset)
	if err != nil {
		return err
	}

	pc := &partConsumer{
		conn:   kfkPartConsumer,
		offset: offset,
		stream: stream,
		closed: make(chan struct{}),
	}

	// Re-check if stream was added in the mean time
	k.lock.Lock()
	_, ok = k.consumers[stream]
	if !ok {
		k.consumers[stream] = pc
	}
	k.lock.Unlock()

	// It was added in the mean time; close the partition consumer
	if ok {
		err = kfkPartConsumer.Close()
		if err != nil {
			//TODO: Is this the right thing to do?
			return err
		}
		return ErrAlreadyListening
	}

	go k.run(pc)

	return nil
}

func (k *Kafka) run(pc *partConsumer) {
	defer close(pc.closed)
	for {
		//TODO: Is a killchan necessary here?
		select {
		case sMessage := <-pc.conn.Messages():
			k.messagebus <- Message{
				Key:       sMessage.Key,
				Val:       sMessage.Value,
				Topic:     pc.stream.Topic,
				Partition: pc.stream.Partition,
				Offset:    sMessage.Offset,
			}
		}
	}
}
<commit_msg>add unlisten<commit_after>package kafka

import (
	"errors"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"

	"github.com/PieterD/kafka-processor/killchan"
	"github.com/Shopify/sarama"
	"github.com/samuel/go-zookeeper/zk"
)

var ErrNoBrokers = errors.New("No kafka brokers found")
var ErrAlreadyListening = errors.New("Already listening to stream")
var ErrNotListening = errors.New("Not listening to stream")

type Message struct {
	Key       []byte
	Val       []byte
	Topic     string
	Partition int32
	Offset    int64
}

type Stream struct {
	Topic     string
	Partition int32
}

type Kafka struct {
	isclosed    int32
	logger      *log.Logger
	zkpeers     []string
	zooConn     *zk.Conn
	kfkConn     sarama.Client
	kfkConsumer sarama.Consumer
	kfkProducer sarama.SyncProducer
	lock        sync.Mutex
	consumers   map[Stream]*partConsumer
	messagebus  chan Message
}

func New(logger *log.Logger, zkpeers []string) (*Kafka, error) {
	k := new(Kafka)
	k.consumers = make(map[Stream]*partConsumer)
	k.messagebus = make(chan Message)
	k.logger = logger
	k.zkpeers = zkpeers
	err := k.connect()
	if err != nil {
		return nil, err
	}
	return k, nil
}

func (k *Kafka) connect() error {
	zooConn, _, err := zk.Connect(k.zkpeers, time.Second)
	if err != nil {
		return fmt.Errorf("Failed to connect to zookeeper: %v", err)
	}
	k.zooConn = zooConn

	kafkaBrokers, err := k.GetKafkaBrokersFromZookeeper()
	if err != nil {
		k.Close()
		return fmt.Errorf("Failed to feth brokers from zookeeper: %v", err)
	}
	if len(kafkaBrokers) == 0 {
		k.Close()
		return ErrNoBrokers
	}

	config := sarama.NewConfig()
	config.Producer.Partitioner = sarama.NewRandomPartitioner
	config.Producer.RequiredAcks = sarama.WaitForAll
	config.ClientID = "kafka-go"

	kfkConn, err := sarama.NewClient(brokerStrings(kafkaBrokers), config)
	if err != nil {
		k.Close()
		return fmt.Errorf("Failed to connect to Kafka: %v", err)
	}
	k.kfkConn = kfkConn

	consumer, err := sarama.NewConsumerFromClient(kfkConn)
	if err != nil {
		k.Close()
		return fmt.Errorf("Failed to create Kafka consumer: %v", err)
	}
	k.kfkConsumer = consumer

	producer, err := sarama.NewSyncProducerFromClient(kfkConn)
	if err != nil {
		consumer.Close()
		return fmt.Errorf("Failed to create sarama syncproducer: %v", err)
	}
	k.kfkProducer = producer

	return nil
}

func (k *Kafka) Close() error {
	if atomic.CompareAndSwapInt32(&k.isclosed, 0, 1) {
		if k.kfkProducer != nil {
			k.kfkProducer.Close()
		}
		if k.kfkConsumer != nil {
			k.kfkConsumer.Close()
		}
		if k.kfkConn != nil {
			k.kfkConn.Close()
		}
		if k.zooConn != nil {
			k.zooConn.Close()
		}
	}
	return nil
}

type partConsumer struct {
	conn   sarama.PartitionConsumer
	offset int64
	stream Stream
	kill   killchan.Killchan
	killed killchan.Killchan
}

func (k *Kafka) Listen(topic string, partition int32, offset int64) error {
	stream := Stream{Topic: topic, Partition: partition}

	// Check if the stream was already added
	k.lock.Lock()
	_, ok := k.consumers[stream]
	k.lock.Unlock()
	if ok {
		return ErrAlreadyListening
	}

	// Create parititon consumer
	kfkPartConsumer, err := k.kfkConsumer.ConsumePartition(topic, partition, offset)
	if err != nil {
		return err
	}

	pc := &partConsumer{
		conn:   kfkPartConsumer,
		offset: offset,
		stream: stream,
		kill:   killchan.New(),
		killed: killchan.New(),
	}

	// Re-check if stream was added in the mean time
	k.lock.Lock()
	_, ok = k.consumers[stream]
	if !ok {
		k.consumers[stream] = pc
	}
	k.lock.Unlock()

	// It was added in the mean time; close the partition consumer
	if ok {
		err = kfkPartConsumer.Close()
		if err != nil {
			//TODO: Is this the right thing to do?
			return err
		}
		return ErrAlreadyListening
	}

	go pc.run(k)

	return nil
}

func (k *Kafka) Unlisten(topic string, partition int32) error {
	k.lock.Lock()
	pc, ok := k.consumers[stream]
	k.lock.Unlock()
	if !ok {
		return ErrNotListening
	}

	pc.close()

	k.lock.Lock()
	delete(k.consumers, stream)
	k.lock.Unlock()

	return nil
}

func (pc *partConsumer) close() bool {
	if pc.kill.Kill() {
		pc.killed.Wait()
		return true
	}
	return false
}

func (pc *partConsumer) run(k *Kafka) {
	defer pc.killed.Kill()
	defer pc.conn.Close()
	for {
		var msg Message

		// Receive a message
		select {
		case sMessage := <-pc.conn.Messages():
			msg = Message{
				Key:       sMessage.Key,
				Val:       sMessage.Value,
				Topic:     pc.stream.Topic,
				Partition: pc.stream.Partition,
				Offset:    sMessage.Offset,
			}
		case <-pc.kill.Chan():
			return
		}

		// Send the message on
		select {
		case k.messagebus <- msg:
		case <-pc.kill.Chan():
			return
		}
	}
}
<|endoftext|>