<commit_before>package mc

// Handles SASL authentication.

import (
	"fmt"
	"strings"
)

func (cn *Conn) Auth(user, pass string) *MCError {
	s, err := cn.authList()
	if err != nil {
		return err
	}

	switch {
	case strings.Index(s, "PLAIN") != -1:
		return cn.authPlain(user, pass)
	}

	return &MCError{0xffff, fmt.Sprintf("mc: unknown auth types %q", s)}
}

func (cn *Conn) authList() (s string, err *MCError) {
	m := &msg{
		header: header{
			Op: OpAuthList,
		},
	}

	err = cn.sendRecv(m)
	return m.val, err
}

func (cn *Conn) authPlain(user, pass string) *MCError {
	m := &msg{
		header: header{
			Op: OpAuthStart,
		},

		key: "PLAIN",
		val: fmt.Sprintf("\x00%s\x00%s", user, pass),
	}

	return cn.sendRecv(m)
}
<commit_msg>document functions in auth.go<commit_after>package mc

// Handles SASL authentication.

import (
	"fmt"
	"strings"
)

// Auth performs SASL authentication (using the PLAIN method) with the server.
func (cn *Conn) Auth(user, pass string) *MCError {
	s, err := cn.authList()
	if err != nil {
		return err
	}

	switch {
	case strings.Index(s, "PLAIN") != -1:
		return cn.authPlain(user, pass)
	}

	return &MCError{0xffff, fmt.Sprintf("mc: unknown auth types %q", s)}
}

// authList runs the SASL authentication list command with the server to
// retrieve the list of support authentication mechansims.
func (cn *Conn) authList() (s string, err *MCError) {
	m := &msg{
		header: header{
			Op: OpAuthList,
		},
	}

	err = cn.sendRecv(m)
	return m.val, err
}

// authPlain performs SASL authentication using the PLAIN method.
func (cn *Conn) authPlain(user, pass string) *MCError {
	m := &msg{
		header: header{
			Op: OpAuthStart,
		},

		key: "PLAIN",
		val: fmt.Sprintf("\x00%s\x00%s", user, pass),
	}

	return cn.sendRecv(m)
}
<|endoftext|>


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

package ssh

import (
	"errors"
	"io"
	"sync"
)

// extendedDataTypeCode identifies an OpenSSL extended data type. See RFC 4254,
// section 5.2.
type extendedDataTypeCode uint32

// extendedDataStderr is the extended data type that is used for stderr.
const extendedDataStderr extendedDataTypeCode = 1

// A Channel is an ordered, reliable, duplex stream that is multiplexed over an
// SSH connection. Channel.Read can return a ChannelRequest as an error.
type Channel interface {
	// Accept accepts the channel creation request.
	Accept() error
	// Reject rejects the channel creation request. After calling this, no
	// other methods on the Channel may be called. If they are then the
	// peer is likely to signal a protocol error and drop the connection.
	Reject(reason RejectionReason, message string) error

	// Read may return a ChannelRequest as an error.
	Read(data []byte) (int, error)
	Write(data []byte) (int, error)
	Close() error

	// Stderr returns an io.Writer that writes to this channel with the
	// extended data type set to stderr.
	Stderr() io.Writer

	// AckRequest either sends an ack or nack to the channel request.
	AckRequest(ok bool) error

	// ChannelType returns the type of the channel, as supplied by the
	// client.
	ChannelType() string
	// ExtraData returns the arbitary payload for this channel, as supplied
	// by the client. This data is specific to the channel type.
	ExtraData() []byte
}

// ChannelRequest represents a request sent on a channel, outside of the normal
// stream of bytes. It may result from calling Read on a Channel.
type ChannelRequest struct {
	Request   string
	WantReply bool
	Payload   []byte
}

func (c ChannelRequest) Error() string {
	return "ssh: channel request received"
}

// RejectionReason is an enumeration used when rejecting channel creation
// requests. See RFC 4254, section 5.1.
type RejectionReason int

const (
	Prohibited RejectionReason = iota + 1
	ConnectionFailed
	UnknownChannelType
	ResourceShortage
)

type channel struct {
	// immutable once created
	chanType  string
	extraData []byte

	theyClosed  bool
	theySentEOF bool
	weClosed    bool
	dead        bool

	serverConn            *ServerConn
	myId, theirId         uint32
	myWindow, theirWindow uint32
	maxPacketSize         uint32
	err                   error

	pendingRequests []ChannelRequest
	pendingData     []byte
	head, length    int

	// This lock is inferior to serverConn.lock
	lock sync.Mutex
	cond *sync.Cond
}

func (c *channel) Accept() error {
	c.serverConn.lock.Lock()
	defer c.serverConn.lock.Unlock()

	if c.serverConn.err != nil {
		return c.serverConn.err
	}

	confirm := channelOpenConfirmMsg{
		PeersId:       c.theirId,
		MyId:          c.myId,
		MyWindow:      c.myWindow,
		MaxPacketSize: c.maxPacketSize,
	}
	return c.serverConn.writePacket(marshal(msgChannelOpenConfirm, confirm))
}

func (c *channel) Reject(reason RejectionReason, message string) error {
	c.serverConn.lock.Lock()
	defer c.serverConn.lock.Unlock()

	if c.serverConn.err != nil {
		return c.serverConn.err
	}

	reject := channelOpenFailureMsg{
		PeersId:  c.theirId,
		Reason:   reason,
		Message:  message,
		Language: "en",
	}
	return c.serverConn.writePacket(marshal(msgChannelOpenFailure, reject))
}

func (c *channel) handlePacket(packet interface{}) {
	c.lock.Lock()
	defer c.lock.Unlock()

	switch packet := packet.(type) {
	case *channelRequestMsg:
		req := ChannelRequest{
			Request:   packet.Request,
			WantReply: packet.WantReply,
			Payload:   packet.RequestSpecificData,
		}

		c.pendingRequests = append(c.pendingRequests, req)
		c.cond.Signal()
	case *channelCloseMsg:
		c.theyClosed = true
		c.cond.Signal()
	case *channelEOFMsg:
		c.theySentEOF = true
		c.cond.Signal()
	case *windowAdjustMsg:
		c.theirWindow += packet.AdditionalBytes
		c.cond.Signal()
	default:
		panic("unknown packet type")
	}
}

func (c *channel) handleData(data []byte) {
	c.lock.Lock()
	defer c.lock.Unlock()

	// The other side should never send us more than our window.
	if len(data)+c.length > len(c.pendingData) {
		// TODO(agl): we should tear down the channel with a protocol
		// error.
		return
	}

	c.myWindow -= uint32(len(data))
	for i := 0; i < 2; i++ {
		tail := c.head + c.length
		if tail >= len(c.pendingData) {
			tail -= len(c.pendingData)
		}
		n := copy(c.pendingData[tail:], data)
		data = data[n:]
		c.length += n
	}

	c.cond.Signal()
}

func (c *channel) Stderr() io.Writer {
	return extendedDataChannel{c: c, t: extendedDataStderr}
}

// extendedDataChannel is an io.Writer that writes any data to c as extended
// data of the given type.
type extendedDataChannel struct {
	t extendedDataTypeCode
	c *channel
}

func (edc extendedDataChannel) Write(data []byte) (n int, err error) {
	c := edc.c
	for len(data) > 0 {
		var space uint32
		if space, err = c.getWindowSpace(uint32(len(data))); err != nil {
			return 0, err
		}

		todo := data
		if uint32(len(todo)) > space {
			todo = todo[:space]
		}

		packet := make([]byte, 1+4+4+4+len(todo))
		packet[0] = msgChannelExtendedData
		marshalUint32(packet[1:], c.theirId)
		marshalUint32(packet[5:], uint32(edc.t))
		marshalUint32(packet[9:], uint32(len(todo)))
		copy(packet[13:], todo)

		c.serverConn.lock.Lock()
		err = c.serverConn.writePacket(packet)
		c.serverConn.lock.Unlock()

		if err != nil {
			return
		}

		n += len(todo)
		data = data[len(todo):]
	}

	return
}

func (c *channel) Read(data []byte) (n int, err error) {
	c.lock.Lock()
	defer c.lock.Unlock()

	if c.err != nil {
		return 0, c.err
	}

	for {
		if c.theySentEOF || c.theyClosed || c.dead {
			return 0, io.EOF
		}

		if len(c.pendingRequests) > 0 {
			req := c.pendingRequests[0]
			if len(c.pendingRequests) == 1 {
				c.pendingRequests = nil
			} else {
				oldPendingRequests := c.pendingRequests
				c.pendingRequests = make([]ChannelRequest, len(oldPendingRequests)-1)
				copy(c.pendingRequests, oldPendingRequests[1:])
			}

			return 0, req
		}

		if c.length > 0 {
			tail := min(c.head+c.length, len(c.pendingData))
			n = copy(data, c.pendingData[c.head:tail])
			c.head += n
			c.length -= n
			if c.head == len(c.pendingData) {
				c.head = 0
			}

			windowAdjustment := uint32(len(c.pendingData)-c.length) - c.myWindow
			if windowAdjustment >= uint32(len(c.pendingData)/2) {
				packet := marshal(msgChannelWindowAdjust, windowAdjustMsg{
					PeersId:         c.theirId,
					AdditionalBytes: windowAdjustment,
				})
				c.serverConn.lock.Lock()
				err = c.serverConn.writePacket(packet)
				c.serverConn.lock.Unlock()
				if err != nil {
					return
				}
				c.myWindow += windowAdjustment
			}

			return
		}

		c.cond.Wait()
	}

	panic("unreachable")
}

// getWindowSpace takes, at most, max bytes of space from the peer's window. It
// returns the number of bytes actually reserved.
func (c *channel) getWindowSpace(max uint32) (uint32, error) {
	c.lock.Lock()
	defer c.lock.Unlock()

	for {
		if c.dead || c.weClosed {
			return 0, io.EOF
		}

		if c.theirWindow > 0 {
			break
		}

		c.cond.Wait()
	}

	taken := c.theirWindow
	if taken > max {
		taken = max
	}

	c.theirWindow -= taken
	return taken, nil
}

func (c *channel) Write(data []byte) (n int, err error) {
	for len(data) > 0 {
		var space uint32
		if space, err = c.getWindowSpace(uint32(len(data))); err != nil {
			return 0, err
		}

		todo := data
		if uint32(len(todo)) > space {
			todo = todo[:space]
		}

		packet := make([]byte, 1+4+4+len(todo))
		packet[0] = msgChannelData
		marshalUint32(packet[1:], c.theirId)
		marshalUint32(packet[5:], uint32(len(todo)))
		copy(packet[9:], todo)

		c.serverConn.lock.Lock()
		err = c.serverConn.writePacket(packet)
		c.serverConn.lock.Unlock()

		if err != nil {
			return
		}

		n += len(todo)
		data = data[len(todo):]
	}

	return
}

func (c *channel) Close() error {
	c.serverConn.lock.Lock()
	defer c.serverConn.lock.Unlock()

	if c.serverConn.err != nil {
		return c.serverConn.err
	}

	if c.weClosed {
		return errors.New("ssh: channel already closed")
	}
	c.weClosed = true

	closeMsg := channelCloseMsg{
		PeersId: c.theirId,
	}
	return c.serverConn.writePacket(marshal(msgChannelClose, closeMsg))
}

func (c *channel) AckRequest(ok bool) error {
	c.serverConn.lock.Lock()
	defer c.serverConn.lock.Unlock()

	if c.serverConn.err != nil {
		return c.serverConn.err
	}

	if !ok {
		ack := channelRequestFailureMsg{
			PeersId: c.theirId,
		}
		return c.serverConn.writePacket(marshal(msgChannelFailure, ack))
	}

	ack := channelRequestSuccessMsg{
		PeersId: c.theirId,
	}
	return c.serverConn.writePacket(marshal(msgChannelSuccess, ack))
}

func (c *channel) ChannelType() string {
	return c.chanType
}

func (c *channel) ExtraData() []byte {
	return c.extraData
}
<commit_msg>ssh: fix deadlock<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ssh

import (
	"errors"
	"io"
	"sync"
)

// extendedDataTypeCode identifies an OpenSSL extended data type. See RFC 4254,
// section 5.2.
type extendedDataTypeCode uint32

// extendedDataStderr is the extended data type that is used for stderr.
const extendedDataStderr extendedDataTypeCode = 1

// A Channel is an ordered, reliable, duplex stream that is multiplexed over an
// SSH connection. Channel.Read can return a ChannelRequest as an error.
type Channel interface {
	// Accept accepts the channel creation request.
	Accept() error
	// Reject rejects the channel creation request. After calling this, no
	// other methods on the Channel may be called. If they are then the
	// peer is likely to signal a protocol error and drop the connection.
	Reject(reason RejectionReason, message string) error

	// Read may return a ChannelRequest as an error.
	Read(data []byte) (int, error)
	Write(data []byte) (int, error)
	Close() error

	// Stderr returns an io.Writer that writes to this channel with the
	// extended data type set to stderr.
	Stderr() io.Writer

	// AckRequest either sends an ack or nack to the channel request.
	AckRequest(ok bool) error

	// ChannelType returns the type of the channel, as supplied by the
	// client.
	ChannelType() string
	// ExtraData returns the arbitary payload for this channel, as supplied
	// by the client. This data is specific to the channel type.
	ExtraData() []byte
}

// ChannelRequest represents a request sent on a channel, outside of the normal
// stream of bytes. It may result from calling Read on a Channel.
type ChannelRequest struct {
	Request   string
	WantReply bool
	Payload   []byte
}

func (c ChannelRequest) Error() string {
	return "ssh: channel request received"
}

// RejectionReason is an enumeration used when rejecting channel creation
// requests. See RFC 4254, section 5.1.
type RejectionReason int

const (
	Prohibited RejectionReason = iota + 1
	ConnectionFailed
	UnknownChannelType
	ResourceShortage
)

type channel struct {
	// immutable once created
	chanType  string
	extraData []byte

	theyClosed  bool
	theySentEOF bool
	weClosed    bool
	dead        bool

	serverConn            *ServerConn
	myId, theirId         uint32
	myWindow, theirWindow uint32
	maxPacketSize         uint32
	err                   error

	pendingRequests []ChannelRequest
	pendingData     []byte
	head, length    int

	// This lock is inferior to serverConn.lock
	lock sync.Mutex
	cond *sync.Cond
}

func (c *channel) Accept() error {
	c.serverConn.lock.Lock()
	defer c.serverConn.lock.Unlock()

	if c.serverConn.err != nil {
		return c.serverConn.err
	}

	confirm := channelOpenConfirmMsg{
		PeersId:       c.theirId,
		MyId:          c.myId,
		MyWindow:      c.myWindow,
		MaxPacketSize: c.maxPacketSize,
	}
	return c.serverConn.writePacket(marshal(msgChannelOpenConfirm, confirm))
}

func (c *channel) Reject(reason RejectionReason, message string) error {
	c.serverConn.lock.Lock()
	defer c.serverConn.lock.Unlock()

	if c.serverConn.err != nil {
		return c.serverConn.err
	}

	reject := channelOpenFailureMsg{
		PeersId:  c.theirId,
		Reason:   reason,
		Message:  message,
		Language: "en",
	}
	return c.serverConn.writePacket(marshal(msgChannelOpenFailure, reject))
}

func (c *channel) handlePacket(packet interface{}) {
	c.lock.Lock()
	defer c.lock.Unlock()

	switch packet := packet.(type) {
	case *channelRequestMsg:
		req := ChannelRequest{
			Request:   packet.Request,
			WantReply: packet.WantReply,
			Payload:   packet.RequestSpecificData,
		}

		c.pendingRequests = append(c.pendingRequests, req)
		c.cond.Signal()
	case *channelCloseMsg:
		c.theyClosed = true
		c.cond.Signal()
	case *channelEOFMsg:
		c.theySentEOF = true
		c.cond.Signal()
	case *windowAdjustMsg:
		c.theirWindow += packet.AdditionalBytes
		c.cond.Signal()
	default:
		panic("unknown packet type")
	}
}

func (c *channel) handleData(data []byte) {
	c.lock.Lock()
	defer c.lock.Unlock()

	// The other side should never send us more than our window.
	if len(data)+c.length > len(c.pendingData) {
		// TODO(agl): we should tear down the channel with a protocol
		// error.
		return
	}

	c.myWindow -= uint32(len(data))
	for i := 0; i < 2; i++ {
		tail := c.head + c.length
		if tail >= len(c.pendingData) {
			tail -= len(c.pendingData)
		}
		n := copy(c.pendingData[tail:], data)
		data = data[n:]
		c.length += n
	}

	c.cond.Signal()
}

func (c *channel) Stderr() io.Writer {
	return extendedDataChannel{c: c, t: extendedDataStderr}
}

// extendedDataChannel is an io.Writer that writes any data to c as extended
// data of the given type.
type extendedDataChannel struct {
	t extendedDataTypeCode
	c *channel
}

func (edc extendedDataChannel) Write(data []byte) (n int, err error) {
	c := edc.c
	for len(data) > 0 {
		var space uint32
		if space, err = c.getWindowSpace(uint32(len(data))); err != nil {
			return 0, err
		}

		todo := data
		if uint32(len(todo)) > space {
			todo = todo[:space]
		}

		packet := make([]byte, 1+4+4+4+len(todo))
		packet[0] = msgChannelExtendedData
		marshalUint32(packet[1:], c.theirId)
		marshalUint32(packet[5:], uint32(edc.t))
		marshalUint32(packet[9:], uint32(len(todo)))
		copy(packet[13:], todo)

		c.serverConn.lock.Lock()
		err = c.serverConn.writePacket(packet)
		c.serverConn.lock.Unlock()

		if err != nil {
			return
		}

		n += len(todo)
		data = data[len(todo):]
	}

	return
}

func (c *channel) Read(data []byte) (n int, err error) {
	n, err, windowAdjustment := c.read(data)

	if windowAdjustment > 0 {
		packet := marshal(msgChannelWindowAdjust, windowAdjustMsg{
			PeersId:         c.theirId,
			AdditionalBytes: windowAdjustment,
		})
		c.serverConn.lock.Lock()
		err = c.serverConn.writePacket(packet)
		c.serverConn.lock.Unlock()
		if err != nil {
			return
		}
	}

	return
}

func (c *channel) read(data []byte) (n int, err error, windowAdjustment uint32) {
	c.lock.Lock()
	defer c.lock.Unlock()

	if c.err != nil {
		return 0, c.err, 0
	}

	for {
		if c.theySentEOF || c.theyClosed || c.dead {
			return 0, io.EOF, 0
		}

		if len(c.pendingRequests) > 0 {
			req := c.pendingRequests[0]
			if len(c.pendingRequests) == 1 {
				c.pendingRequests = nil
			} else {
				oldPendingRequests := c.pendingRequests
				c.pendingRequests = make([]ChannelRequest, len(oldPendingRequests)-1)
				copy(c.pendingRequests, oldPendingRequests[1:])
			}

			return 0, req, 0
		}

		if c.length > 0 {
			tail := min(c.head+c.length, len(c.pendingData))
			n = copy(data, c.pendingData[c.head:tail])
			c.head += n
			c.length -= n
			if c.head == len(c.pendingData) {
				c.head = 0
			}

			windowAdjustment = uint32(len(c.pendingData)-c.length) - c.myWindow
			if windowAdjustment < uint32(len(c.pendingData)/2) {
				windowAdjustment = 0
			}
			c.myWindow += windowAdjustment

			return
		}

		c.cond.Wait()
	}

	panic("unreachable")
}

// getWindowSpace takes, at most, max bytes of space from the peer's window. It
// returns the number of bytes actually reserved.
func (c *channel) getWindowSpace(max uint32) (uint32, error) {
	c.lock.Lock()
	defer c.lock.Unlock()

	for {
		if c.dead || c.weClosed {
			return 0, io.EOF
		}

		if c.theirWindow > 0 {
			break
		}

		c.cond.Wait()
	}

	taken := c.theirWindow
	if taken > max {
		taken = max
	}

	c.theirWindow -= taken
	return taken, nil
}

func (c *channel) Write(data []byte) (n int, err error) {
	for len(data) > 0 {
		var space uint32
		if space, err = c.getWindowSpace(uint32(len(data))); err != nil {
			return 0, err
		}

		todo := data
		if uint32(len(todo)) > space {
			todo = todo[:space]
		}

		packet := make([]byte, 1+4+4+len(todo))
		packet[0] = msgChannelData
		marshalUint32(packet[1:], c.theirId)
		marshalUint32(packet[5:], uint32(len(todo)))
		copy(packet[9:], todo)

		c.serverConn.lock.Lock()
		err = c.serverConn.writePacket(packet)
		c.serverConn.lock.Unlock()

		if err != nil {
			return
		}

		n += len(todo)
		data = data[len(todo):]
	}

	return
}

func (c *channel) Close() error {
	c.serverConn.lock.Lock()
	defer c.serverConn.lock.Unlock()

	if c.serverConn.err != nil {
		return c.serverConn.err
	}

	if c.weClosed {
		return errors.New("ssh: channel already closed")
	}
	c.weClosed = true

	closeMsg := channelCloseMsg{
		PeersId: c.theirId,
	}
	return c.serverConn.writePacket(marshal(msgChannelClose, closeMsg))
}

func (c *channel) AckRequest(ok bool) error {
	c.serverConn.lock.Lock()
	defer c.serverConn.lock.Unlock()

	if c.serverConn.err != nil {
		return c.serverConn.err
	}

	if !ok {
		ack := channelRequestFailureMsg{
			PeersId: c.theirId,
		}
		return c.serverConn.writePacket(marshal(msgChannelFailure, ack))
	}

	ack := channelRequestSuccessMsg{
		PeersId: c.theirId,
	}
	return c.serverConn.writePacket(marshal(msgChannelSuccess, ack))
}

func (c *channel) ChannelType() string {
	return c.chanType
}

func (c *channel) ExtraData() []byte {
	return c.extraData
}
<|endoftext|>


<commit_before>package stacksignal

import (
	"os"
	"os/signal"
	"runtime"
	"syscall"
)

func init() {
	go func() {
		c := make(chan os.Signal)
		signal.Notify(c, syscall.SIGUSR1)
		buf := make([]byte, 1024*16)
		for _ = range c {
			runtime.Stack(buf, true)
			os.Stderr.Write(buf)
		}
	}()
}
<commit_msg>Add documentation<commit_after>/*
Package stacksignal registers a signal handler for SIGUSR1.
Whenever the importing program receives SIGUSR1, a full
stacktrace of all goroutines will be printed to StdErr.

Since there are no functions to use, the import line should
be

	import _ "github.com/surma/stacksignal"
*/
package stacksignal

import (
	"os"
	"os/signal"
	"runtime"
	"syscall"
)

func init() {
	go func() {
		c := make(chan os.Signal)
		signal.Notify(c, syscall.SIGUSR1)
		buf := make([]byte, 1024*16)
		for _ = range c {
			runtime.Stack(buf, true)
			os.Stderr.Write(buf)
		}
	}()
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/hmac"
	"crypto/md4"
	"crypto/md5"
	"crypto/ripemd160"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"flag"
	"fmt"
	"hash"
	"hash/adler32"
	"hash/crc32"
	"hash/crc64"
	"hash/fnv"
	"io"
	"log"
	"os"
	"sort"
	"strings"
)

const (
	version   = "1.0"
	chunkSize = 32768 // Read files in 32KB chunks
)

var (
	alg  *string = flag.String("a", "md5", "algorithm")
	str  *string = flag.String("s", "", "string to hash instead of file")
	key  *string = flag.String("k", "", "key (for hashes that use a key, e.g. HMAC)")
	salt *string = flag.String("salt", "", "salt")
)

var (
	algDescs = map[string]string{
		"adler32":         "Adler-32 checksum (RFC 1950)",
		"crc32":           "32-bit cyclic redundancy check (CRC-32) checksum (defaults to IEEE polynomial)",
		"crc32ieee":       "CRC-32 using the IEEE polynomial (0xedb88320)",
		"crc32castagnoli": "CRC-32 using the Castagnoli polynomial (0x82f63b78)",
		"crc32koopman":    "CRC-32 using the Koopman polynomial (0xeb31d82e)",
		"crc64":           "64-bit cyclic redundancy check (CRC-64) checksum (defaults to ISO polynomial)",
		"crc64iso":        "CRC-64 using the ISO polynomial (0xD800000000000000)",
		"crc64ecma":       "CRC-64 using the ECMA polynomial (0xC96C5795D7870F42)",
		"fnv":             "FNV-1 non-cryptographic hash (defaults to fnv32)",
		"fnv32":           "32-bit FNV-1",
		"fnv32a":          "32-bit FNV-1a",
		"fnv64":           "64-bit FNV-1",
		"fnv64a":          "64-bit FNV-1a",
		"hmac":            "Keyed-Hash Message Authentication Code (HMAC) (requires -k <key>) (defaults to SHA-256)",
		"hmacmd5":         "HMAC using MD5 (requires -k <key>)",
		"hmacsha1":        "HMAC using SHA-1 (requires -k <key>)",
		"hmacsha256":      "HMAC using SHA-256 (requires -k <key>)",
		"hmacsha512":      "HMAC using SHA-512 (requires -k <key>)",
		"md4":             "MD4 hash (RFC 1320)",
		"md5":             "MD5 hash (RFC 1321)",
		"ripemd160":       "RIPEMD-160 hash",
		"sha1":            "SHA-1 hash (RFC 3174)",
		"sha224":          "SHA-224 hash (FIPS 180-2)",
		"sha256":          "SHA-256 hash (FIPS 180-2)",
		"sha384":          "SHA-384 hash (FIPS 180-2)",
		"sha512":          "SHA-512 hash (FIPS 180-2)",
	}
)

func GetGenerator(a string) (hash.Hash, os.Error) {
	var g hash.Hash
	switch a {
	default:
		return md5.New(), os.NewError("Invalid algorithm")
	case "adler32":
		g = adler32.New()
	case "crc32", "crc32ieee":
		g = crc32.New(crc32.MakeTable(crc32.IEEE))
	case "crc32castagnoli":
		g = crc32.New(crc32.MakeTable(crc32.Castagnoli))
	case "crc32koopman":
		g = crc32.New(crc32.MakeTable(crc32.Koopman))
	case "crc64", "crc64iso":
		g = crc64.New(crc64.MakeTable(crc64.ISO))
	case "crc64ecma":
		g = crc64.New(crc64.MakeTable(crc64.ECMA))
	case "fnv", "fnv32":
		g = fnv.New32()
	case "fnv32a":
		g = fnv.New32a()
	case "fnv64":
		g = fnv.New64()
	case "fnv64a":
		g = fnv.New64a()
	case "hmac", "hmacsha256":
		g = hmac.NewSHA256([]byte(*key))
	case "hmacmd5":
		g = hmac.NewMD5([]byte(*key))
	case "hmacsha1":
		g = hmac.NewSHA1([]byte(*key))
	case "hmacsha512":
		g = hmac.New(sha512.New, []byte(*key))
	case "md4":
		g = md4.New()
	case "md5":
		g = md5.New()
	case "ripemd160":
		g = ripemd160.New()
	case "sha1":
		g = sha1.New()
	case "sha224":
		g = sha256.New224()
	case "sha256":
		g = sha256.New()
	case "sha384":
		g = sha512.New384()
	case "sha512":
		g = sha512.New()
	}
	return g, nil
}

func HashString(gen hash.Hash, s string) string {
	gen.Write([]byte(s))
	return fmt.Sprintf("%x", gen.Sum())
}

func HashFile(gen hash.Hash, f io.Reader) (string, os.Error) {
	buf := make([]byte, chunkSize)
	gen.Write([]byte(*salt))
	for {
		bytesRead, err := f.Read(buf)
		if err != nil {
			if err == os.EOF && bytesRead == 0 { // Empty file
				gen.Write([]byte(""))
			} else {
				return "", err
			}
		} else {
			gen.Write(buf[:bytesRead])
		}
		if bytesRead < chunkSize { // EOF
			break
		}
	}
	return fmt.Sprintf("%x", gen.Sum()), nil
}

func Usage() {
	fmt.Println("Usage:", os.Args[0], "[-a <algorithm>] [-salt <salt>] -s <string to hash>/-f <file to hash>\n")
	fmt.Println("Examples:")
	fmt.Println(" ", os.Args[0], "-a md5 document.txt               ", "Generate MD5 hash of a file")
	fmt.Println(" ", os.Args[0], "-a md5 *                          ", "Generate MD5 hash of all files in folder")
	fmt.Println(" ", os.Args[0], "-a sha1 -s hello world            ", "Generate SHA-1 hash of a string")
	fmt.Println(" ", os.Args[0], "-a sha1 -salt s4lt -s hello world ", "Generate salted SHA-1 hash of a string")
	fmt.Println("")
	fmt.Println("Available algorithms (default is MD5):")
	mk := make([]string, len(algDescs))
	i := 0
	for k, _ := range algDescs {
		mk[i] = k
		i++
	}
	sort.Strings(mk)
	for _, v := range mk {
		fmt.Printf("  %-15s %s\n", v, algDescs[v])
	}
}

func main() {
	flag.Parse()
	if flag.NArg() == 0 && *str == "" {
		Usage()
		return
	}
	*alg = strings.ToLower(*alg)
	gen, err := GetGenerator(*alg)
	if err != nil {
		log.Fatalln(err)
	}
	if *str == "" { // Hash file(s)
		for _, path := range flag.Args() {
			var res string
			f, err := os.Open(path)
			defer f.Close()
			if err != nil {
				res = err.String()
			} else {
				h, err := HashFile(gen, f)
				if err != nil {
					res = err.String()
				} else {
					res = h
				}
			}
			fmt.Println(res, "", path)
			gen.Reset()
		}
	} else { // Hash string
		fmt.Println(HashString(gen, *salt+*str))
	}
}
<commit_msg>Remove log pkg dependency<commit_after>package main

import (
	"crypto/hmac"
	"crypto/md4"
	"crypto/md5"
	"crypto/ripemd160"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"flag"
	"fmt"
	"hash"
	"hash/adler32"
	"hash/crc32"
	"hash/crc64"
	"hash/fnv"
	"io"
	"os"
	"sort"
	"strings"
)

const (
	version   = "1.0"
	chunkSize = 32768 // Read files in 32KB chunks
)

var (
	alg  *string = flag.String("a", "md5", "algorithm")
	str  *string = flag.String("s", "", "string to hash instead of file")
	key  *string = flag.String("k", "", "key (for hashes that use a key, e.g. HMAC)")
	salt *string = flag.String("salt", "", "salt")
)

var (
	algDescs = map[string]string{
		"adler32":         "Adler-32 checksum (RFC 1950)",
		"crc32":           "32-bit cyclic redundancy check (CRC-32) checksum (defaults to IEEE polynomial)",
		"crc32ieee":       "CRC-32 using the IEEE polynomial (0xedb88320)",
		"crc32castagnoli": "CRC-32 using the Castagnoli polynomial (0x82f63b78)",
		"crc32koopman":    "CRC-32 using the Koopman polynomial (0xeb31d82e)",
		"crc64":           "64-bit cyclic redundancy check (CRC-64) checksum (defaults to ISO polynomial)",
		"crc64iso":        "CRC-64 using the ISO polynomial (0xD800000000000000)",
		"crc64ecma":       "CRC-64 using the ECMA polynomial (0xC96C5795D7870F42)",
		"fnv":             "FNV-1 non-cryptographic hash (defaults to fnv32)",
		"fnv32":           "32-bit FNV-1",
		"fnv32a":          "32-bit FNV-1a",
		"fnv64":           "64-bit FNV-1",
		"fnv64a":          "64-bit FNV-1a",
		"hmac":            "Keyed-Hash Message Authentication Code (HMAC) (requires -k <key>) (defaults to SHA-256)",
		"hmacmd5":         "HMAC using MD5 (requires -k <key>)",
		"hmacsha1":        "HMAC using SHA-1 (requires -k <key>)",
		"hmacsha256":      "HMAC using SHA-256 (requires -k <key>)",
		"hmacsha512":      "HMAC using SHA-512 (requires -k <key>)",
		"md4":             "MD4 hash (RFC 1320)",
		"md5":             "MD5 hash (RFC 1321)",
		"ripemd160":       "RIPEMD-160 hash",
		"sha1":            "SHA-1 hash (RFC 3174)",
		"sha224":          "SHA-224 hash (FIPS 180-2)",
		"sha256":          "SHA-256 hash (FIPS 180-2)",
		"sha384":          "SHA-384 hash (FIPS 180-2)",
		"sha512":          "SHA-512 hash (FIPS 180-2)",
	}
)

func GetGenerator(a string) (hash.Hash, os.Error) {
	var g hash.Hash
	switch a {
	default:
		return md5.New(), os.NewError("Invalid algorithm")
	case "adler32":
		g = adler32.New()
	case "crc32", "crc32ieee":
		g = crc32.New(crc32.MakeTable(crc32.IEEE))
	case "crc32castagnoli":
		g = crc32.New(crc32.MakeTable(crc32.Castagnoli))
	case "crc32koopman":
		g = crc32.New(crc32.MakeTable(crc32.Koopman))
	case "crc64", "crc64iso":
		g = crc64.New(crc64.MakeTable(crc64.ISO))
	case "crc64ecma":
		g = crc64.New(crc64.MakeTable(crc64.ECMA))
	case "fnv", "fnv32":
		g = fnv.New32()
	case "fnv32a":
		g = fnv.New32a()
	case "fnv64":
		g = fnv.New64()
	case "fnv64a":
		g = fnv.New64a()
	case "hmac", "hmacsha256":
		g = hmac.NewSHA256([]byte(*key))
	case "hmacmd5":
		g = hmac.NewMD5([]byte(*key))
	case "hmacsha1":
		g = hmac.NewSHA1([]byte(*key))
	case "hmacsha512":
		g = hmac.New(sha512.New, []byte(*key))
	case "md4":
		g = md4.New()
	case "md5":
		g = md5.New()
	case "ripemd160":
		g = ripemd160.New()
	case "sha1":
		g = sha1.New()
	case "sha224":
		g = sha256.New224()
	case "sha256":
		g = sha256.New()
	case "sha384":
		g = sha512.New384()
	case "sha512":
		g = sha512.New()
	}
	return g, nil
}

func HashString(gen hash.Hash, s string) string {
	gen.Write([]byte(s))
	return fmt.Sprintf("%x", gen.Sum())
}

func HashFile(gen hash.Hash, f io.Reader) (string, os.Error) {
	buf := make([]byte, chunkSize)
	gen.Write([]byte(*salt))
	for {
		bytesRead, err := f.Read(buf)
		if err != nil {
			if err == os.EOF && bytesRead == 0 { // Empty file
				gen.Write([]byte(""))
			} else {
				return "", err
			}
		} else {
			gen.Write(buf[:bytesRead])
		}
		if bytesRead < chunkSize { // EOF
			break
		}
	}
	return fmt.Sprintf("%x", gen.Sum()), nil
}

func Usage() {
	fmt.Println("Usage:", os.Args[0], "[-a <algorithm>] [-salt <salt>] -s <string to hash>/-f <file to hash>\n")
	fmt.Println("Examples:")
	fmt.Println(" ", os.Args[0], "-a md5 document.txt               ", "Generate MD5 hash of a file")
	fmt.Println(" ", os.Args[0], "-a md5 *                          ", "Generate MD5 hash of all files in folder")
	fmt.Println(" ", os.Args[0], "-a sha1 -s hello world            ", "Generate SHA-1 hash of a string")
	fmt.Println(" ", os.Args[0], "-a sha1 -salt s4lt -s hello world ", "Generate salted SHA-1 hash of a string")
	fmt.Println("")
	fmt.Println("Available algorithms (default is MD5):")
	mk := make([]string, len(algDescs))
	i := 0
	for k, _ := range algDescs {
		mk[i] = k
		i++
	}
	sort.Strings(mk)
	for _, v := range mk {
		fmt.Printf("  %-15s %s\n", v, algDescs[v])
	}
}

func main() {
	flag.Parse()
	if flag.NArg() == 0 && *str == "" {
		Usage()
		return
	}
	*alg = strings.ToLower(*alg)
	gen, err := GetGenerator(*alg)
	if err != nil {
		fmt.Println("Error:", err)
		return
	}
	if *str == "" { // Hash file(s)
		for _, path := range flag.Args() {
			var res string
			f, err := os.Open(path)
			defer f.Close()
			if err != nil {
				res = err.String()
			} else {
				h, err := HashFile(gen, f)
				if err != nil {
					res = err.String()
				} else {
					res = h
				}
			}
			fmt.Println(res, "", path)
			gen.Reset()
		}
	} else { // Hash string
		fmt.Println(HashString(gen, *salt+*str))
	}
}
<|endoftext|>


<commit_before>package state

import (
	"fmt"

	abci "github.com/tendermint/tendermint/abci/types"
	cmn "github.com/tendermint/tendermint/libs/common"
	dbm "github.com/tendermint/tendermint/libs/db"
	"github.com/tendermint/tendermint/types"
)

//------------------------------------------------------------------------

func calcValidatorsKey(height int64) []byte {
	return []byte(cmn.Fmt("validatorsKey:%v", height))
}

func calcConsensusParamsKey(height int64) []byte {
	return []byte(cmn.Fmt("consensusParamsKey:%v", height))
}

func calcABCIResponsesKey(height int64) []byte {
	return []byte(cmn.Fmt("abciResponsesKey:%v", height))
}

// LoadStateFromDBOrGenesisFile loads the most recent state from the database,
// or creates a new one from the given genesisFilePath and persists the result
// to the database.
func LoadStateFromDBOrGenesisFile(stateDB dbm.DB, genesisFilePath string) (State, error) {
	state := LoadState(stateDB)
	if state.IsEmpty() {
		var err error
		state, err = MakeGenesisStateFromFile(genesisFilePath)
		if err != nil {
			return state, err
		}
		SaveState(stateDB, state)
	}

	return state, nil
}

// LoadStateFromDBOrGenesisDoc loads the most recent state from the database,
// or creates a new one from the given genesisDoc and persists the result
// to the database.
func LoadStateFromDBOrGenesisDoc(stateDB dbm.DB, genesisDoc *types.GenesisDoc) (State, error) {
	state := LoadState(stateDB)
	if state.IsEmpty() {
		var err error
		state, err = MakeGenesisState(genesisDoc)
		if err != nil {
			return state, err
		}
		SaveState(stateDB, state)
	}

	return state, nil
}

// LoadState loads the State from the database.
func LoadState(db dbm.DB) State {
	return loadState(db, stateKey)
}

func loadState(db dbm.DB, key []byte) (state State) {
	buf := db.Get(key)
	if len(buf) == 0 {
		return state
	}

	err := cdc.UnmarshalBinaryBare(buf, &state)
	if err != nil {
		// DATA HAS BEEN CORRUPTED OR THE SPEC HAS CHANGED
		cmn.Exit(cmn.Fmt(`LoadState: Data has been corrupted or its spec has changed:
                %v\n`, err))
	}
	// TODO: ensure that buf is completely read.

	return state
}

// SaveState persists the State, the ValidatorsInfo, and the ConsensusParamsInfo to the database.
func SaveState(db dbm.DB, state State) {
	saveState(db, state, stateKey)
}

func saveState(db dbm.DB, state State, key []byte) {
	nextHeight := state.LastBlockHeight + 1
	saveValidatorsInfo(db, nextHeight, state.LastHeightValidatorsChanged, state.Validators)
	saveConsensusParamsInfo(db, nextHeight, state.LastHeightConsensusParamsChanged, state.ConsensusParams)
	db.SetSync(stateKey, state.Bytes())
}

//------------------------------------------------------------------------

// ABCIResponses retains the responses
// of the various ABCI calls during block processing.
// It is persisted to disk for each height before calling Commit.
type ABCIResponses struct {
	DeliverTx []*abci.ResponseDeliverTx
	EndBlock  *abci.ResponseEndBlock
}

// NewABCIResponses returns a new ABCIResponses
func NewABCIResponses(block *types.Block) *ABCIResponses {
	resDeliverTxs := make([]*abci.ResponseDeliverTx, block.NumTxs)
	if block.NumTxs == 0 {
		// This makes Amino encoding/decoding consistent.
		resDeliverTxs = nil
	}
	return &ABCIResponses{
		DeliverTx: resDeliverTxs,
	}
}

// Bytes serializes the ABCIResponse using go-amino.
func (arz *ABCIResponses) Bytes() []byte {
	return cdc.MustMarshalBinaryBare(arz)
}

func (arz *ABCIResponses) ResultsHash() []byte {
	results := types.NewResults(arz.DeliverTx)
	return results.Hash()
}

// LoadABCIResponses loads the ABCIResponses for the given height from the database.
// This is useful for recovering from crashes where we called app.Commit and before we called
// s.Save(). It can also be used to produce Merkle proofs of the result of txs.
func LoadABCIResponses(db dbm.DB, height int64) (*ABCIResponses, error) {
	buf := db.Get(calcABCIResponsesKey(height))
	if len(buf) == 0 {
		return nil, ErrNoABCIResponsesForHeight{height}
	}

	abciResponses := new(ABCIResponses)
	err := cdc.UnmarshalBinaryBare(buf, abciResponses)
	if err != nil {
		// DATA HAS BEEN CORRUPTED OR THE SPEC HAS CHANGED
		cmn.Exit(cmn.Fmt(`LoadABCIResponses: Data has been corrupted or its spec has
                changed: %v\n`, err))
	}
	// TODO: ensure that buf is completely read.

	return abciResponses, nil
}

// SaveABCIResponses persists the ABCIResponses to the database.
// This is useful in case we crash after app.Commit and before s.Save().
// Responses are indexed by height so they can also be loaded later to produce Merkle proofs.
func saveABCIResponses(db dbm.DB, height int64, abciResponses *ABCIResponses) {
	db.SetSync(calcABCIResponsesKey(height), abciResponses.Bytes())
}

//-----------------------------------------------------------------------------

// ValidatorsInfo represents the latest validator set, or the last height it changed
type ValidatorsInfo struct {
	ValidatorSet      *types.ValidatorSet
	LastHeightChanged int64
}

// Bytes serializes the ValidatorsInfo using go-amino.
func (valInfo *ValidatorsInfo) Bytes() []byte {
	return cdc.MustMarshalBinaryBare(valInfo)
}

// LoadValidators loads the ValidatorSet for a given height.
// Returns ErrNoValSetForHeight if the validator set can't be found for this height.
func LoadValidators(db dbm.DB, height int64) (*types.ValidatorSet, error) {
	valInfo := loadValidatorsInfo(db, height)
	if valInfo == nil {
		return nil, ErrNoValSetForHeight{height}
	}

	if valInfo.ValidatorSet == nil {
		valInfo2 := loadValidatorsInfo(db, valInfo.LastHeightChanged)
		if valInfo2 == nil {
			panic(fmt.Sprintf(`Couldn't find validators at height %d as
                        last changed from height %d`, valInfo.LastHeightChanged, height))
		}
		valInfo = valInfo2
	}

	return valInfo.ValidatorSet, nil
}

func loadValidatorsInfo(db dbm.DB, height int64) *ValidatorsInfo {
	buf := db.Get(calcValidatorsKey(height))
	if len(buf) == 0 {
		return nil
	}

	v := new(ValidatorsInfo)
	err := cdc.UnmarshalBinaryBare(buf, v)
	if err != nil {
		// DATA HAS BEEN CORRUPTED OR THE SPEC HAS CHANGED
		cmn.Exit(cmn.Fmt(`LoadValidators: Data has been corrupted or its spec has changed:
                %v\n`, err))
	}
	// TODO: ensure that buf is completely read.

	return v
}

// saveValidatorsInfo persists the validator set for the next block to disk.
// It should be called from s.Save(), right before the state itself is persisted.
// If the validator set did not change after processing the latest block,
// only the last height for which the validators changed is persisted.
func saveValidatorsInfo(db dbm.DB, nextHeight, changeHeight int64, valSet *types.ValidatorSet) {
	valInfo := &ValidatorsInfo{
		LastHeightChanged: changeHeight,
	}
	if changeHeight == nextHeight {
		valInfo.ValidatorSet = valSet
	}
	db.SetSync(calcValidatorsKey(nextHeight), valInfo.Bytes())
}

//-----------------------------------------------------------------------------

// ConsensusParamsInfo represents the latest consensus params, or the last height it changed
type ConsensusParamsInfo struct {
	ConsensusParams   types.ConsensusParams
	LastHeightChanged int64
}

// Bytes serializes the ConsensusParamsInfo using go-amino.
func (params ConsensusParamsInfo) Bytes() []byte {
	return cdc.MustMarshalBinaryBare(params)
}

// LoadConsensusParams loads the ConsensusParams for a given height.
func LoadConsensusParams(db dbm.DB, height int64) (types.ConsensusParams, error) {
	empty := types.ConsensusParams{}

	paramsInfo := loadConsensusParamsInfo(db, height)
	if paramsInfo == nil {
		return empty, ErrNoConsensusParamsForHeight{height}
	}

	if paramsInfo.ConsensusParams == empty {
		paramsInfo2 := loadConsensusParamsInfo(db, paramsInfo.LastHeightChanged)
		if paramsInfo2 == nil {
			panic(fmt.Sprintf(`Couldn't find consensus params at height %d as
                        last changed from height %d`, paramsInfo.LastHeightChanged, height))
		}
		paramsInfo = paramsInfo2
	}

	return paramsInfo.ConsensusParams, nil
}

func loadConsensusParamsInfo(db dbm.DB, height int64) *ConsensusParamsInfo {
	buf := db.Get(calcConsensusParamsKey(height))
	if len(buf) == 0 {
		return nil
	}

	paramsInfo := new(ConsensusParamsInfo)
	err := cdc.UnmarshalBinaryBare(buf, paramsInfo)
	if err != nil {
		// DATA HAS BEEN CORRUPTED OR THE SPEC HAS CHANGED
		cmn.Exit(cmn.Fmt(`LoadConsensusParams: Data has been corrupted or its spec has changed:
                %v\n`, err))
	}
	// TODO: ensure that buf is completely read.

	return paramsInfo
}

// saveConsensusParamsInfo persists the consensus params for the next block to disk.
// It should be called from s.Save(), right before the state itself is persisted.
// If the consensus params did not change after processing the latest block,
// only the last height for which they changed is persisted.
func saveConsensusParamsInfo(db dbm.DB, nextHeight, changeHeight int64, params types.ConsensusParams) {
	paramsInfo := &ConsensusParamsInfo{
		LastHeightChanged: changeHeight,
	}
	if changeHeight == nextHeight {
		paramsInfo.ConsensusParams = params
	}
	db.SetSync(calcConsensusParamsKey(nextHeight), paramsInfo.Bytes())
}
<commit_msg>state: Format panics<commit_after>package state

import (
	"fmt"

	abci "github.com/tendermint/tendermint/abci/types"
	cmn "github.com/tendermint/tendermint/libs/common"
	dbm "github.com/tendermint/tendermint/libs/db"
	"github.com/tendermint/tendermint/types"
)

//------------------------------------------------------------------------

func calcValidatorsKey(height int64) []byte {
	return []byte(cmn.Fmt("validatorsKey:%v", height))
}

func calcConsensusParamsKey(height int64) []byte {
	return []byte(cmn.Fmt("consensusParamsKey:%v", height))
}

func calcABCIResponsesKey(height int64) []byte {
	return []byte(cmn.Fmt("abciResponsesKey:%v", height))
}

// LoadStateFromDBOrGenesisFile loads the most recent state from the database,
// or creates a new one from the given genesisFilePath and persists the result
// to the database.
func LoadStateFromDBOrGenesisFile(stateDB dbm.DB, genesisFilePath string) (State, error) {
	state := LoadState(stateDB)
	if state.IsEmpty() {
		var err error
		state, err = MakeGenesisStateFromFile(genesisFilePath)
		if err != nil {
			return state, err
		}
		SaveState(stateDB, state)
	}

	return state, nil
}

// LoadStateFromDBOrGenesisDoc loads the most recent state from the database,
// or creates a new one from the given genesisDoc and persists the result
// to the database.
func LoadStateFromDBOrGenesisDoc(stateDB dbm.DB, genesisDoc *types.GenesisDoc) (State, error) {
	state := LoadState(stateDB)
	if state.IsEmpty() {
		var err error
		state, err = MakeGenesisState(genesisDoc)
		if err != nil {
			return state, err
		}
		SaveState(stateDB, state)
	}

	return state, nil
}

// LoadState loads the State from the database.
func LoadState(db dbm.DB) State {
	return loadState(db, stateKey)
}

func loadState(db dbm.DB, key []byte) (state State) {
	buf := db.Get(key)
	if len(buf) == 0 {
		return state
	}

	err := cdc.UnmarshalBinaryBare(buf, &state)
	if err != nil {
		// DATA HAS BEEN CORRUPTED OR THE SPEC HAS CHANGED
		cmn.Exit(cmn.Fmt(`LoadState: Data has been corrupted or its spec has changed:
                %v\n`, err))
	}
	// TODO: ensure that buf is completely read.

	return state
}

// SaveState persists the State, the ValidatorsInfo, and the ConsensusParamsInfo to the database.
func SaveState(db dbm.DB, state State) {
	saveState(db, state, stateKey)
}

func saveState(db dbm.DB, state State, key []byte) {
	nextHeight := state.LastBlockHeight + 1
	saveValidatorsInfo(db, nextHeight, state.LastHeightValidatorsChanged, state.Validators)
	saveConsensusParamsInfo(db, nextHeight, state.LastHeightConsensusParamsChanged, state.ConsensusParams)
	db.SetSync(stateKey, state.Bytes())
}

//------------------------------------------------------------------------

// ABCIResponses retains the responses
// of the various ABCI calls during block processing.
// It is persisted to disk for each height before calling Commit.
type ABCIResponses struct {
	DeliverTx []*abci.ResponseDeliverTx
	EndBlock  *abci.ResponseEndBlock
}

// NewABCIResponses returns a new ABCIResponses
func NewABCIResponses(block *types.Block) *ABCIResponses {
	resDeliverTxs := make([]*abci.ResponseDeliverTx, block.NumTxs)
	if block.NumTxs == 0 {
		// This makes Amino encoding/decoding consistent.
		resDeliverTxs = nil
	}
	return &ABCIResponses{
		DeliverTx: resDeliverTxs,
	}
}

// Bytes serializes the ABCIResponse using go-amino.
func (arz *ABCIResponses) Bytes() []byte {
	return cdc.MustMarshalBinaryBare(arz)
}

func (arz *ABCIResponses) ResultsHash() []byte {
	results := types.NewResults(arz.DeliverTx)
	return results.Hash()
}

// LoadABCIResponses loads the ABCIResponses for the given height from the database.
// This is useful for recovering from crashes where we called app.Commit and before we called
// s.Save(). It can also be used to produce Merkle proofs of the result of txs.
func LoadABCIResponses(db dbm.DB, height int64) (*ABCIResponses, error) {
	buf := db.Get(calcABCIResponsesKey(height))
	if len(buf) == 0 {
		return nil, ErrNoABCIResponsesForHeight{height}
	}

	abciResponses := new(ABCIResponses)
	err := cdc.UnmarshalBinaryBare(buf, abciResponses)
	if err != nil {
		// DATA HAS BEEN CORRUPTED OR THE SPEC HAS CHANGED
		cmn.Exit(cmn.Fmt(`LoadABCIResponses: Data has been corrupted or its spec has
                changed: %v\n`, err))
	}
	// TODO: ensure that buf is completely read.

	return abciResponses, nil
}

// SaveABCIResponses persists the ABCIResponses to the database.
// This is useful in case we crash after app.Commit and before s.Save().
// Responses are indexed by height so they can also be loaded later to produce Merkle proofs.
func saveABCIResponses(db dbm.DB, height int64, abciResponses *ABCIResponses) {
	db.SetSync(calcABCIResponsesKey(height), abciResponses.Bytes())
}

//-----------------------------------------------------------------------------

// ValidatorsInfo represents the latest validator set, or the last height it changed
type ValidatorsInfo struct {
	ValidatorSet      *types.ValidatorSet
	LastHeightChanged int64
}

// Bytes serializes the ValidatorsInfo using go-amino.
func (valInfo *ValidatorsInfo) Bytes() []byte {
	return cdc.MustMarshalBinaryBare(valInfo)
}

// LoadValidators loads the ValidatorSet for a given height.
// Returns ErrNoValSetForHeight if the validator set can't be found for this height.
func LoadValidators(db dbm.DB, height int64) (*types.ValidatorSet, error) {
	valInfo := loadValidatorsInfo(db, height)
	if valInfo == nil {
		return nil, ErrNoValSetForHeight{height}
	}

	if valInfo.ValidatorSet == nil {
		valInfo2 := loadValidatorsInfo(db, valInfo.LastHeightChanged)
		if valInfo2 == nil {
			panic(
				fmt.Sprintf(
					"Couldn't find validators at height %d as last changed from height %d",
					valInfo.LastHeightChanged,
					height,
				),
			)
		}
		valInfo = valInfo2
	}

	return valInfo.ValidatorSet, nil
}

func loadValidatorsInfo(db dbm.DB, height int64) *ValidatorsInfo {
	buf := db.Get(calcValidatorsKey(height))
	if len(buf) == 0 {
		return nil
	}

	v := new(ValidatorsInfo)
	err := cdc.UnmarshalBinaryBare(buf, v)
	if err != nil {
		// DATA HAS BEEN CORRUPTED OR THE SPEC HAS CHANGED
		cmn.Exit(cmn.Fmt(`LoadValidators: Data has been corrupted or its spec has changed:
                %v\n`, err))
	}
	// TODO: ensure that buf is completely read.

	return v
}

// saveValidatorsInfo persists the validator set for the next block to disk.
// It should be called from s.Save(), right before the state itself is persisted.
// If the validator set did not change after processing the latest block,
// only the last height for which the validators changed is persisted.
func saveValidatorsInfo(db dbm.DB, nextHeight, changeHeight int64, valSet *types.ValidatorSet) {
	valInfo := &ValidatorsInfo{
		LastHeightChanged: changeHeight,
	}
	if changeHeight == nextHeight {
		valInfo.ValidatorSet = valSet
	}
	db.SetSync(calcValidatorsKey(nextHeight), valInfo.Bytes())
}

//-----------------------------------------------------------------------------

// ConsensusParamsInfo represents the latest consensus params, or the last height it changed
type ConsensusParamsInfo struct {
	ConsensusParams   types.ConsensusParams
	LastHeightChanged int64
}

// Bytes serializes the ConsensusParamsInfo using go-amino.
func (params ConsensusParamsInfo) Bytes() []byte {
	return cdc.MustMarshalBinaryBare(params)
}

// LoadConsensusParams loads the ConsensusParams for a given height.
func LoadConsensusParams(db dbm.DB, height int64) (types.ConsensusParams, error) {
	empty := types.ConsensusParams{}

	paramsInfo := loadConsensusParamsInfo(db, height)
	if paramsInfo == nil {
		return empty, ErrNoConsensusParamsForHeight{height}
	}

	if paramsInfo.ConsensusParams == empty {
		paramsInfo2 := loadConsensusParamsInfo(db, paramsInfo.LastHeightChanged)
		if paramsInfo2 == nil {
			panic(
				fmt.Sprintf(
					"Couldn't find consensus params at height %d as last changed from height %d",
					paramsInfo.LastHeightChanged,
					height,
				),
			)
		}
		paramsInfo = paramsInfo2
	}

	return paramsInfo.ConsensusParams, nil
}

func loadConsensusParamsInfo(db dbm.DB, height int64) *ConsensusParamsInfo {
	buf := db.Get(calcConsensusParamsKey(height))
	if len(buf) == 0 {
		return nil
	}

	paramsInfo := new(ConsensusParamsInfo)
	err := cdc.UnmarshalBinaryBare(buf, paramsInfo)
	if err != nil {
		// DATA HAS BEEN CORRUPTED OR THE SPEC HAS CHANGED
		cmn.Exit(cmn.Fmt(`LoadConsensusParams: Data has been corrupted or its spec has changed:
                %v\n`, err))
	}
	// TODO: ensure that buf is completely read.

	return paramsInfo
}

// saveConsensusParamsInfo persists the consensus params for the next block to disk.
// It should be called from s.Save(), right before the state itself is persisted.
// If the consensus params did not change after processing the latest block,
// only the last height for which they changed is persisted.
func saveConsensusParamsInfo(db dbm.DB, nextHeight, changeHeight int64, params types.ConsensusParams) {
	paramsInfo := &ConsensusParamsInfo{
		LastHeightChanged: changeHeight,
	}
	if changeHeight == nextHeight {
		paramsInfo.ConsensusParams = params
	}
	db.SetSync(calcConsensusParamsKey(nextHeight), paramsInfo.Bytes())
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"log"
	"math"
	"net"
	"os"
	"os/signal"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"
)

const (
	VERSION                 = "0.5.2-alpha"
	MAX_UNPROCESSED_PACKETS = 1000
	MAX_UDP_PACKET_SIZE     = 512
)

var signalchan chan os.Signal

type Packet struct {
	Bucket   string
	Value    interface{}
	Modifier string
	Sampling float32
}

type Uint64Slice []uint64

func (s Uint64Slice) Len() int           { return len(s) }
func (s Uint64Slice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s Uint64Slice) Less(i, j int) bool { return s[i] < s[j] }

type Percentiles []*Percentile
type Percentile struct {
	float float64
	str   string
}

func (a *Percentiles) Set(s string) error {
	f, err := strconv.ParseFloat(s, 64)
	if err != nil {
		return err
	}
	*a = append(*a, &Percentile{f, strings.Replace(s, ".", "_", -1)})
	return nil
}
func (p *Percentile) String() string {
	return p.str
}
func (a *Percentiles) String() string {
	return fmt.Sprintf("%v", *a)
}

var (
	serviceAddress   = flag.String("address", ":8125", "UDP service address")
	graphiteAddress  = flag.String("graphite", "127.0.0.1:2003", "Graphite service address (or - to disable)")
	flushInterval    = flag.Int64("flush-interval", 10, "Flush interval (seconds)")
	debug            = flag.Bool("debug", false, "print statistics sent to graphite")
	showVersion      = flag.Bool("version", false, "print version string")
	persistCountKeys = flag.Int64("persist-count-keys", 60, "number of flush-interval's to persist count keys")
	receiveCounter   = flag.String("receive-counter", "", "Metric name for total metrics recevied per interval")
	percentThreshold = Percentiles{}
	prefix           = flag.String("prefix", "", "Prefix for all stats")
)

func init() {
	flag.Var(&percentThreshold, "percent-threshold", "Threshold percent (0-100, may be given multiple times)")
}

var (
	In       = make(chan *Packet, MAX_UNPROCESSED_PACKETS)
	counters = make(map[string]int64)
	gauges   = make(map[string]uint64)
	timers   = make(map[string]Uint64Slice)
)

func monitor() {
	period := time.Duration(*flushInterval) * time.Second
	ticker := time.NewTicker(period)
	for {
		select {
		case sig := <-signalchan:
			fmt.Printf("!! Caught signal %d... shutting down\n", sig)
			if err := submit(time.Now().Add(period)); err != nil {
				log.Printf("ERROR: %s", err)
			}
			return
		case <-ticker.C:
			if err := submit(time.Now().Add(period)); err != nil {
				log.Printf("ERROR: %s", err)
			}
		case s := <-In:
			if *receiveCounter != "" {
				v, ok := counters[*receiveCounter]
				if !ok || v < 0 {
					counters[*receiveCounter] = 0
				}
				counters[*receiveCounter] += 1
			}

			if s.Modifier == "ms" {
				_, ok := timers[s.Bucket]
				if !ok {
					var t Uint64Slice
					timers[s.Bucket] = t
				}
				timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
			} else if s.Modifier == "g" {
				gauges[s.Bucket] = s.Value.(uint64)
			} else {
				v, ok := counters[s.Bucket]
				if !ok || v < 0 {
					counters[s.Bucket] = 0
				}
				counters[s.Bucket] += int64(float64(s.Value.(int64)) * float64(1/s.Sampling))
			}
		}
	}
}

func submit(deadline time.Time) error {
	var buffer bytes.Buffer
	var num int64

	now := time.Now().Unix()

	client, err := net.Dial("tcp", *graphiteAddress)
	if err != nil {
		if *debug {
			log.Printf("WARNING: resetting counters when in debug mode")
			processCounters(&buffer, now)
			processGauges(&buffer, now)
			processTimers(&buffer, now, percentThreshold)
		}
		errmsg := fmt.Sprintf("dialing %s failed - %s", *graphiteAddress, err)
		return errors.New(errmsg)
	}
	defer client.Close()

	err = client.SetDeadline(deadline)
	if err != nil {
		errmsg := fmt.Sprintf("could not set deadline:", err)
		return errors.New(errmsg)
	}

	num += processCounters(&buffer, now)
	num += processGauges(&buffer, now)
	num += processTimers(&buffer, now, percentThreshold)
	if num == 0 {
		return nil
	}

	if *debug {
		for _, line := range bytes.Split(buffer.Bytes(), []byte("\n")) {
			if len(line) == 0 {
				continue
			}
			log.Printf("DEBUG: %s", line)
		}
	}

	_, err = client.Write(buffer.Bytes())
	if err != nil {
		errmsg := fmt.Sprintf("failed to write stats - %s", err)
		return errors.New(errmsg)
	}

	log.Printf("sent %d stats to %s", num, *graphiteAddress)

	return nil
}

func processCounters(buffer *bytes.Buffer, now int64) int64 {
	var num int64
	// continue sending zeros for counters for a short period of time even if we have no new data
	// note we use the same in-memory value to denote both the actual value of the counter (value >= 0)
	// as well as how many turns to keep the counter for (value < 0)
	// for more context see https://github.com/bitly/statsdaemon/pull/8
	for s, c := range counters {
		switch {
		case c <= *persistCountKeys:
			// consider this purgable
			delete(counters, s)
			continue
		case c < 0:
			counters[s] -= 1
			fmt.Fprintf(buffer, "%s %d %d\n", s, 0, now)
		case c >= 0:
			counters[s] = -1
			fmt.Fprintf(buffer, "%s %d %d\n", s, c, now)
		}
		num++
	}
	return num
}

func processGauges(buffer *bytes.Buffer, now int64) int64 {
	var num int64
	for g, c := range gauges {
		if c == math.MaxUint64 {
			continue
		}
		fmt.Fprintf(buffer, "%s %d %d\n", g, c, now)
		gauges[g] = math.MaxUint64
		num++
	}
	return num
}

func processTimers(buffer *bytes.Buffer, now int64, pctls Percentiles) int64 {
	var num int64
	for u, t := range timers {
		if len(t) == 0 {
			continue
		}

		num++

		sort.Sort(t)
		min := t[0]
		max := t[len(t)-1]
		maxAtThreshold := max
		count := len(t)

		sum := uint64(0)
		for _, value := range t {
			sum += value
		}
		mean := float64(sum) / float64(len(t))

		for _, pct := range pctls {
			if len(t) > 1 {
				var abs float64
				if pct.float >= 0 {
					abs = pct.float
				} else {
					abs = 100 + pct.float
				}
				// poor man's math.Round(x):
				// math.Floor(x + 0.5)
				indexOfPerc := int(math.Floor(((abs / 100.0) * float64(count)) + 0.5))
				if pct.float >= 0 {
					indexOfPerc -= 1 // index offset=0
				}
				maxAtThreshold = t[indexOfPerc]
			}

			var tmpl string
			var pctstr string
			if pct.float >= 0 {
				tmpl = "%s.upper_%s %d %d\n"
				pctstr = pct.str
			} else {
				tmpl = "%s.lower_%s %d %d\n"
				pctstr = pct.str[1:]
			}
			fmt.Fprintf(buffer, tmpl, u, pctstr, maxAtThreshold, now)
		}

		var z Uint64Slice
		timers[u] = z

		fmt.Fprintf(buffer, "%s.mean %f %d\n", u, mean, now)
		fmt.Fprintf(buffer, "%s.upper %d %d\n", u, max, now)
		fmt.Fprintf(buffer, "%s.lower %d %d\n", u, min, now)
		fmt.Fprintf(buffer, "%s.count %d %d\n", u, count, now)
	}
	return num
}

func parseMessage(data []byte) []*Packet {
	var (
		output []*Packet
		input  []byte
	)

	for _, line := range bytes.Split(data, []byte("\n")) {
		if len(line) == 0 {
			continue
		}
		input = line

		index := bytes.IndexByte(input, ':')
		if index < 0 {
			if *debug {
				log.Printf("ERROR: failed to parse line: %s\n", string(line))
			}
			continue
		}

		name := input[:index]

		index++
		input = input[index:]

		index = bytes.IndexByte(input, '|')
		if index < 0 {
			if *debug {
				log.Printf("ERROR: failed to parse line: %s\n", string(line))
			}
			continue
		}

		val := input[:index]
		index++

		var mtypeStr string

		if input[index] == 'm' {
			index++
			if index >= len(input) || input[index] != 's' {
				if *debug {
					log.Printf("ERROR: failed to parse line: %s\n", string(line))
				}
				continue
			}
			mtypeStr = "ms"
		} else {
			mtypeStr = string(input[index])
		}

		index++
		input = input[index:]

		var (
			value interface{}
			err   error
		)

		if mtypeStr[0] == 'c' {
			value, err = strconv.ParseInt(string(val), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseInt %s - %s", string(val), err)
				continue
			}
		} else {
			value, err = strconv.ParseUint(string(val), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseUint %s - %s", string(val), err)
				continue
			}
		}

		var sampleRate float32 = 1

		if len(input) > 0 && bytes.HasPrefix(input, []byte("|@")) {
			input = input[2:]
			rate, err := strconv.ParseFloat(string(input), 32)
			if err == nil {
				sampleRate = float32(rate)
			}
		}

		packet := &Packet{
			Bucket:   *prefix + string(name),
			Value:    value,
			Modifier: mtypeStr,
			Sampling: sampleRate,
		}
		output = append(output, packet)
	}
	return output
}

func udpListener() {
	address, _ := net.ResolveUDPAddr("udp", *serviceAddress)
	log.Printf("listening on %s", address)
	listener, err := net.ListenUDP("udp", address)
	if err != nil {
		log.Fatalf("ERROR: ListenUDP - %s", err)
	}
	defer listener.Close()

	message := make([]byte, MAX_UDP_PACKET_SIZE)
	for {
		n, remaddr, err := listener.ReadFromUDP(message)
		if err != nil {
			log.Printf("ERROR: reading UDP packet from %+v - %s", remaddr, err)
			continue
		}

		for _, p := range parseMessage(message[:n]) {
			In <- p
		}
	}
}

func main() {
	flag.Parse()

	if *showVersion {
		fmt.Printf("statsdaemon v%s (built w/%s)\n", VERSION, runtime.Version())
		return
	}

	signalchan = make(chan os.Signal, 1)
	signal.Notify(signalchan, syscall.SIGTERM)
	*persistCountKeys = -1 * (*persistCountKeys)

	go udpListener()
	monitor()
}
<commit_msg>bump to 0.6-alpha<commit_after>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"log"
	"math"
	"net"
	"os"
	"os/signal"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"
)

const (
	VERSION                 = "0.6-alpha"
	MAX_UNPROCESSED_PACKETS = 1000
	MAX_UDP_PACKET_SIZE     = 512
)

var signalchan chan os.Signal

type Packet struct {
	Bucket   string
	Value    interface{}
	Modifier string
	Sampling float32
}

type Uint64Slice []uint64

func (s Uint64Slice) Len() int           { return len(s) }
func (s Uint64Slice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s Uint64Slice) Less(i, j int) bool { return s[i] < s[j] }

type Percentiles []*Percentile
type Percentile struct {
	float float64
	str   string
}

func (a *Percentiles) Set(s string) error {
	f, err := strconv.ParseFloat(s, 64)
	if err != nil {
		return err
	}
	*a = append(*a, &Percentile{f, strings.Replace(s, ".", "_", -1)})
	return nil
}
func (p *Percentile) String() string {
	return p.str
}
func (a *Percentiles) String() string {
	return fmt.Sprintf("%v", *a)
}

var (
	serviceAddress   = flag.String("address", ":8125", "UDP service address")
	graphiteAddress  = flag.String("graphite", "127.0.0.1:2003", "Graphite service address (or - to disable)")
	flushInterval    = flag.Int64("flush-interval", 10, "Flush interval (seconds)")
	debug            = flag.Bool("debug", false, "print statistics sent to graphite")
	showVersion      = flag.Bool("version", false, "print version string")
	persistCountKeys = flag.Int64("persist-count-keys", 60, "number of flush-interval's to persist count keys")
	receiveCounter   = flag.String("receive-counter", "", "Metric name for total metrics recevied per interval")
	percentThreshold = Percentiles{}
	prefix           = flag.String("prefix", "", "Prefix for all stats")
)

func init() {
	flag.Var(&percentThreshold, "percent-threshold", "Threshold percent (0-100, may be given multiple times)")
}

var (
	In       = make(chan *Packet, MAX_UNPROCESSED_PACKETS)
	counters = make(map[string]int64)
	gauges   = make(map[string]uint64)
	timers   = make(map[string]Uint64Slice)
)

func monitor() {
	period := time.Duration(*flushInterval) * time.Second
	ticker := time.NewTicker(period)
	for {
		select {
		case sig := <-signalchan:
			fmt.Printf("!! Caught signal %d... shutting down\n", sig)
			if err := submit(time.Now().Add(period)); err != nil {
				log.Printf("ERROR: %s", err)
			}
			return
		case <-ticker.C:
			if err := submit(time.Now().Add(period)); err != nil {
				log.Printf("ERROR: %s", err)
			}
		case s := <-In:
			if *receiveCounter != "" {
				v, ok := counters[*receiveCounter]
				if !ok || v < 0 {
					counters[*receiveCounter] = 0
				}
				counters[*receiveCounter] += 1
			}

			if s.Modifier == "ms" {
				_, ok := timers[s.Bucket]
				if !ok {
					var t Uint64Slice
					timers[s.Bucket] = t
				}
				timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
			} else if s.Modifier == "g" {
				gauges[s.Bucket] = s.Value.(uint64)
			} else {
				v, ok := counters[s.Bucket]
				if !ok || v < 0 {
					counters[s.Bucket] = 0
				}
				counters[s.Bucket] += int64(float64(s.Value.(int64)) * float64(1/s.Sampling))
			}
		}
	}
}

func submit(deadline time.Time) error {
	var buffer bytes.Buffer
	var num int64

	now := time.Now().Unix()

	client, err := net.Dial("tcp", *graphiteAddress)
	if err != nil {
		if *debug {
			log.Printf("WARNING: resetting counters when in debug mode")
			processCounters(&buffer, now)
			processGauges(&buffer, now)
			processTimers(&buffer, now, percentThreshold)
		}
		errmsg := fmt.Sprintf("dialing %s failed - %s", *graphiteAddress, err)
		return errors.New(errmsg)
	}
	defer client.Close()

	err = client.SetDeadline(deadline)
	if err != nil {
		errmsg := fmt.Sprintf("could not set deadline:", err)
		return errors.New(errmsg)
	}

	num += processCounters(&buffer, now)
	num += processGauges(&buffer, now)
	num += processTimers(&buffer, now, percentThreshold)
	if num == 0 {
		return nil
	}

	if *debug {
		for _, line := range bytes.Split(buffer.Bytes(), []byte("\n")) {
			if len(line) == 0 {
				continue
			}
			log.Printf("DEBUG: %s", line)
		}
	}

	_, err = client.Write(buffer.Bytes())
	if err != nil {
		errmsg := fmt.Sprintf("failed to write stats - %s", err)
		return errors.New(errmsg)
	}

	log.Printf("sent %d stats to %s", num, *graphiteAddress)

	return nil
}

func processCounters(buffer *bytes.Buffer, now int64) int64 {
	var num int64
	// continue sending zeros for counters for a short period of time even if we have no new data
	// note we use the same in-memory value to denote both the actual value of the counter (value >= 0)
	// as well as how many turns to keep the counter for (value < 0)
	// for more context see https://github.com/bitly/statsdaemon/pull/8
	for s, c := range counters {
		switch {
		case c <= *persistCountKeys:
			// consider this purgable
			delete(counters, s)
			continue
		case c < 0:
			counters[s] -= 1
			fmt.Fprintf(buffer, "%s %d %d\n", s, 0, now)
		case c >= 0:
			counters[s] = -1
			fmt.Fprintf(buffer, "%s %d %d\n", s, c, now)
		}
		num++
	}
	return num
}

func processGauges(buffer *bytes.Buffer, now int64) int64 {
	var num int64
	for g, c := range gauges {
		if c == math.MaxUint64 {
			continue
		}
		fmt.Fprintf(buffer, "%s %d %d\n", g, c, now)
		gauges[g] = math.MaxUint64
		num++
	}
	return num
}

func processTimers(buffer *bytes.Buffer, now int64, pctls Percentiles) int64 {
	var num int64
	for u, t := range timers {
		if len(t) == 0 {
			continue
		}

		num++

		sort.Sort(t)
		min := t[0]
		max := t[len(t)-1]
		maxAtThreshold := max
		count := len(t)

		sum := uint64(0)
		for _, value := range t {
			sum += value
		}
		mean := float64(sum) / float64(len(t))

		for _, pct := range pctls {
			if len(t) > 1 {
				var abs float64
				if pct.float >= 0 {
					abs = pct.float
				} else {
					abs = 100 + pct.float
				}
				// poor man's math.Round(x):
				// math.Floor(x + 0.5)
				indexOfPerc := int(math.Floor(((abs / 100.0) * float64(count)) + 0.5))
				if pct.float >= 0 {
					indexOfPerc -= 1 // index offset=0
				}
				maxAtThreshold = t[indexOfPerc]
			}

			var tmpl string
			var pctstr string
			if pct.float >= 0 {
				tmpl = "%s.upper_%s %d %d\n"
				pctstr = pct.str
			} else {
				tmpl = "%s.lower_%s %d %d\n"
				pctstr = pct.str[1:]
			}
			fmt.Fprintf(buffer, tmpl, u, pctstr, maxAtThreshold, now)
		}

		var z Uint64Slice
		timers[u] = z

		fmt.Fprintf(buffer, "%s.mean %f %d\n", u, mean, now)
		fmt.Fprintf(buffer, "%s.upper %d %d\n", u, max, now)
		fmt.Fprintf(buffer, "%s.lower %d %d\n", u, min, now)
		fmt.Fprintf(buffer, "%s.count %d %d\n", u, count, now)
	}
	return num
}

func parseMessage(data []byte) []*Packet {
	var (
		output []*Packet
		input  []byte
	)

	for _, line := range bytes.Split(data, []byte("\n")) {
		if len(line) == 0 {
			continue
		}
		input = line

		index := bytes.IndexByte(input, ':')
		if index < 0 {
			if *debug {
				log.Printf("ERROR: failed to parse line: %s\n", string(line))
			}
			continue
		}

		name := input[:index]

		index++
		input = input[index:]

		index = bytes.IndexByte(input, '|')
		if index < 0 {
			if *debug {
				log.Printf("ERROR: failed to parse line: %s\n", string(line))
			}
			continue
		}

		val := input[:index]
		index++

		var mtypeStr string

		if input[index] == 'm' {
			index++
			if index >= len(input) || input[index] != 's' {
				if *debug {
					log.Printf("ERROR: failed to parse line: %s\n", string(line))
				}
				continue
			}
			mtypeStr = "ms"
		} else {
			mtypeStr = string(input[index])
		}

		index++
		input = input[index:]

		var (
			value interface{}
			err   error
		)

		if mtypeStr[0] == 'c' {
			value, err = strconv.ParseInt(string(val), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseInt %s - %s", string(val), err)
				continue
			}
		} else {
			value, err = strconv.ParseUint(string(val), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseUint %s - %s", string(val), err)
				continue
			}
		}

		var sampleRate float32 = 1

		if len(input) > 0 && bytes.HasPrefix(input, []byte("|@")) {
			input = input[2:]
			rate, err := strconv.ParseFloat(string(input), 32)
			if err == nil {
				sampleRate = float32(rate)
			}
		}

		packet := &Packet{
			Bucket:   *prefix + string(name),
			Value:    value,
			Modifier: mtypeStr,
			Sampling: sampleRate,
		}
		output = append(output, packet)
	}
	return output
}

func udpListener() {
	address, _ := net.ResolveUDPAddr("udp", *serviceAddress)
	log.Printf("listening on %s", address)
	listener, err := net.ListenUDP("udp", address)
	if err != nil {
		log.Fatalf("ERROR: ListenUDP - %s", err)
	}
	defer listener.Close()

	message := make([]byte, MAX_UDP_PACKET_SIZE)
	for {
		n, remaddr, err := listener.ReadFromUDP(message)
		if err != nil {
			log.Printf("ERROR: reading UDP packet from %+v - %s", remaddr, err)
			continue
		}

		for _, p := range parseMessage(message[:n]) {
			In <- p
		}
	}
}

func main() {
	flag.Parse()

	if *showVersion {
		fmt.Printf("statsdaemon v%s (built w/%s)\n", VERSION, runtime.Version())
		return
	}

	signalchan = make(chan os.Signal, 1)
	signal.Notify(signalchan, syscall.SIGTERM)
	*persistCountKeys = -1 * (*persistCountKeys)

	go udpListener()
	monitor()
}
<|endoftext|>


<commit_before>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package store

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"

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

// RWFileSystem is an interface which includes http.FileSystem and a Create method
// for creating files.
type RWFileSystem interface {
	FileSystem

	// Create a file with associated path, returns an io.WriteCloser.  Only when Close()
	// returns can it be assumed that the file has been written.
	Create(ctx context.Context, path string) (io.WriteCloser, error)

	// Wait blocks until any pending write calls have been completed.
	Wait() error
}

// dir implements RWFileSystem by extending the behaviour of http.Dir to include a Create
// method which creates files under the root.
type dir struct {
	FileSystem

	root string
}

// Dir creates a new RWFileSystem with the specified root (similar to http.Dir)
func Dir(root string) RWFileSystem {
	return &dir{
		NewFileSystem(http.Dir(root), fmt.Sprintf("local (%v)", root)),
		root,
	}
}

func (d *dir) absPath(path string) (string, error) {
	cleanPath := filepath.Clean(d.root + "/" + path)
	path, err := filepath.Abs(cleanPath)
	if err != nil {
		return "", fmt.Errorf("error finding absolute path for: '%v' ('%v'): %v", path, cleanPath, err)
	}

	absRoot, err := filepath.Abs(d.root)
	if err != nil {
		return "", fmt.Errorf("error finding absolute path for root '%v': %v", d.root, err)
	}

	if !strings.HasPrefix(filepath.Dir(path), absRoot) {
		return "", fmt.Errorf("invalid path ('%v' is outside '%v'): %v", filepath.Dir(path), absRoot, path)
	}
	return path, nil
}

// Create a file rooted in the Dir file system.
func (d *dir) Create(ctx context.Context, path string) (io.WriteCloser, error) {
	absPath, err := d.absPath(path)
	if err != nil {
		return nil, err
	}
	err = os.MkdirAll(filepath.Dir(absPath), os.ModePerm)
	if err != nil {
		return nil, err
	}
	return os.Create(absPath)
}

// Wait implements RWFileSystem.
func (d *dir) Wait() error { return nil }

// CachedFileSystem is an implemetation of http.FileServer which caches the results of
// calls to src in a RWFileSystem.
type CachedFileSystem struct {
	src   FileSystem
	cache RWFileSystem

	mu       sync.RWMutex // protects errCache
	errCache map[string]error

	errCh chan<- error
	wg    sync.WaitGroup
}

// setError sets an error for a path.  This is intended to prevent the src being continually
// queried after returning a non-temporary failure to retrieve a path.
// TODO: Better handle temporary errors!
func (c *CachedFileSystem) setError(path string, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.errCache[path] = err
}

func (c *CachedFileSystem) error(path string) (err error, exists bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	err, ok := c.errCache[path]
	return err, ok
}

// Open implements FileSystem.  If the required file isn't in the cache
// then the file is opened from the src, and then concurrently copied into the
// cache (with errors passed back on the filesystem error channel).
func (c *CachedFileSystem) Open(ctx context.Context, path string) (http.File, error) {
	if err, ok := c.error(path); ok {
		return nil, fmt.Errorf("cached error: %v", err)
	}

	f, err := c.cache.Open(ctx, path)
	if err == nil {
		return f, nil
	}

	f, err = c.src.Open(ctx, path)
	if err != nil {
		c.setError(path, err)
		return nil, err
	}

	go func() { // TODO: improve this so that we don't have to fetch the file again!
		c.wg.Add(1)
		defer c.wg.Done()

		src, err := c.src.Open(ctx, path)
		if err != nil {
			c.errCh <- fmt.Errorf("error opening file for second time: %v", err)
			return
		}
		defer func() {
			err := src.Close()
			if err != nil {
				c.errCh <- err
			}
		}()

		cache, err := c.cache.Create(ctx, path)
		if err != nil {
			c.errCh <- fmt.Errorf("error creating file in cache: %v", err)
			return
		}
		defer func() {
			err := cache.Close()
			if err != nil {
				c.errCh <- err
			}
		}()

		_, err = io.Copy(cache, src)
		if err != nil {
			c.errCh <- fmt.Errorf("error copying src file data into cache: %v", err)
		}
	}()

	return f, nil
}

// Wait implements RWFileSystem.
func (c *CachedFileSystem) Wait() error {
	c.wg.Wait()
	return nil
}

// NewCachedFileSystem implements http.FileSystem and caches every request made to
// src in cache.  The returned error channel passes back any errors which occur when
// files are being concurrently copied into the cache.
func NewCachedFileSystem(src FileSystem, cache RWFileSystem) (*CachedFileSystem, <-chan error) {
	errCh := make(chan error)
	return &CachedFileSystem{
		src:      src,
		cache:    cache,
		errCh:    errCh,
		errCache: make(map[string]error),
	}, errCh
}
<commit_msg>Refactored FileSystem error caching into ErrorCachedFileSystem.<commit_after>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package store

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"

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

// RWFileSystem is an interface which includes http.FileSystem and a Create method
// for creating files.
type RWFileSystem interface {
	FileSystem

	// Create a file with associated path, returns an io.WriteCloser.  Only when Close()
	// returns can it be assumed that the file has been written.
	Create(ctx context.Context, path string) (io.WriteCloser, error)

	// Wait blocks until any pending write calls have been completed.
	Wait() error
}

// dir implements RWFileSystem by extending the behaviour of http.Dir to include a Create
// method which creates files under the root.
type dir struct {
	FileSystem

	root string
}

// Dir creates a new RWFileSystem with the specified root (similar to http.Dir)
func Dir(root string) RWFileSystem {
	return &dir{
		NewFileSystem(http.Dir(root), fmt.Sprintf("local (%v)", root)),
		root,
	}
}

func (d *dir) absPath(path string) (string, error) {
	cleanPath := filepath.Clean(d.root + "/" + path)
	path, err := filepath.Abs(cleanPath)
	if err != nil {
		return "", fmt.Errorf("error finding absolute path for: '%v' ('%v'): %v", path, cleanPath, err)
	}

	absRoot, err := filepath.Abs(d.root)
	if err != nil {
		return "", fmt.Errorf("error finding absolute path for root '%v': %v", d.root, err)
	}

	if !strings.HasPrefix(filepath.Dir(path), absRoot) {
		return "", fmt.Errorf("invalid path ('%v' is outside '%v'): %v", filepath.Dir(path), absRoot, path)
	}
	return path, nil
}

// Create a file rooted in the Dir file system.
func (d *dir) Create(ctx context.Context, path string) (io.WriteCloser, error) {
	absPath, err := d.absPath(path)
	if err != nil {
		return nil, err
	}
	err = os.MkdirAll(filepath.Dir(absPath), os.ModePerm)
	if err != nil {
		return nil, err
	}
	return os.Create(absPath)
}

// Wait implements RWFileSystem.
func (d *dir) Wait() error { return nil }

// CachedError is an error returned by CachedErrorFileSystems when Open errors are cached
// rather than live.
type CachedError struct {
	Err error
}

// Error implements error.
func (c *CachedError) Error() string {
	return fmt.Sprintf("cached error: %v", c.Err)
}

// CachedErrorFileSystem provides an error cache to prevent erroring FileSystem requests
// from being repeated. See open for more details.
type CachedErrorFileSystem struct {
	FileSystem

	sync.RWMutex
	m map[string]error
}

func (c *CachedErrorFileSystem) setError(path string, err error) {
	c.Lock()
	defer c.Unlock()

	c.m[path] = err
}

func (c *CachedErrorFileSystem) getError(path string) (error, bool) {
	c.RLock()
	defer c.RUnlock()

	err, ok := c.m[path]
	return err, ok
}

// Open implements FileSystem, and caches errors from the underlying FileSystem.  The first time
// an error is encountered it is returned unchanged. Subsequent calls with an erroring path
// return a CachedError-wrapped version of the original error.
func (c CachedErrorFileSystem) Open(ctx context.Context, path string) (http.File, error) {
	err, ok := c.getError(path)
	if ok {
		return nil, &CachedError{
			Err:err,
		}
	}

	f, err := c.FileSystem.Open(ctx, path)
	if err != nil {
		c.setError(path, err)
		return nil, err
	}
	return f, nil
}

// CachedFileSystem is an implemetation of http.FileServer which caches the results of
// calls to src in a RWFileSystem.
type CachedFileSystem struct {
	src   FileSystem
	cache RWFileSystem

	errCh chan<- error
	wg    sync.WaitGroup
}

// Open implements FileSystem.  If the required file isn't in the cache
// then the file is opened from the src, and then concurrently copied into the
// cache (with errors passed back on the filesystem error channel).
func (c *CachedFileSystem) Open(ctx context.Context, path string) (http.File, error) {
	f, err := c.cache.Open(ctx, path)
	if err == nil {
		return f, nil
	}

	f, err = c.src.Open(ctx, path)
	if err != nil {
		return nil, err
	}

	go func() { // TODO: improve this so that we don't have to fetch the file again!
		c.wg.Add(1)
		defer c.wg.Done()

		src, err := c.src.Open(ctx, path)
		if err != nil {
			c.errCh <- fmt.Errorf("error opening file for second time: %v", err)
			return
		}
		defer func() {
			err := src.Close()
			if err != nil {
				c.errCh <- err
			}
		}()

		cache, err := c.cache.Create(ctx, path)
		if err != nil {
			c.errCh <- fmt.Errorf("error creating file in cache: %v", err)
			return
		}
		defer func() {
			err := cache.Close()
			if err != nil {
				c.errCh <- err
			}
		}()

		_, err = io.Copy(cache, src)
		if err != nil {
			c.errCh <- fmt.Errorf("error copying src file data into cache: %v", err)
		}
	}()

	return f, nil
}

// Wait implements RWFileSystem.
func (c *CachedFileSystem) Wait() error {
	c.wg.Wait()
	return nil
}

// NewCachedFileSystem implements http.FileSystem and caches every request made to
// src in cache.  The returned error channel passes back any errors which occur when
// files are being concurrently copied into the cache.
func NewCachedFileSystem(src FileSystem, cache RWFileSystem) (*CachedFileSystem, <-chan error) {
	errCh := make(chan error)
	return &CachedFileSystem{
		src:   src,
		cache: cache,
		errCh: errCh,
	}, errCh
}
<|endoftext|>


<commit_before>package store

import (
	"crypto/x509"
	"errors"
	"fmt"
	"strings"
)

// Store represents a certificate store (often called 'pool') and has
// operations on it which mutate the underlying state (e.g. a file or
// directory).
type Store interface {
	// List returns the currently trusted X509 certificates contained
	// within the cert store
	List() ([]*x509.Certificate, error)
}

// Platform returns a new instance of Store for the running os/platform
func Platform() Store {
	return platform()
}

// ForApp returns a `Store` instance for the given app
func ForApp(app string) (Store, error) {
	switch strings.ToLower(app) {
	case "chrome":
		return NssStore(), nil
	case "firefox":
		return NssStore(), nil
	case "java":
		return JavaStore(), nil
	default:
		return nil, fmt.Errorf("application '%s' not found", app)
	}
}
<commit_msg>store: fix compile<commit_after>package store

import (
	"crypto/x509"
	"fmt"
	"strings"
)

// Store represents a certificate store (often called 'pool') and has
// operations on it which mutate the underlying state (e.g. a file or
// directory).
type Store interface {
	// List returns the currently trusted X509 certificates contained
	// within the cert store
	List() ([]*x509.Certificate, error)
}

// Platform returns a new instance of Store for the running os/platform
func Platform() Store {
	return platform()
}

// ForApp returns a `Store` instance for the given app
func ForApp(app string) (Store, error) {
	switch strings.ToLower(app) {
	case "chrome":
		return NssStore(), nil
	case "firefox":
		return NssStore(), nil
	case "java":
		return JavaStore(), nil
	default:
		return nil, fmt.Errorf("application '%s' not found", app)
	}
}
<|endoftext|>


<commit_before>// Package store provides support to interact with the entries
package store

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha512"
	"encoding/json"
	"path/filepath"
	"time"

	"github.com/maxibanki/golang-url-shortener/util"
	"github.com/pborman/uuid"
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/bcrypt"

	"github.com/asaskevich/govalidator"
	"github.com/boltdb/bolt"
	"github.com/pkg/errors"
)

// Store holds internal funcs and vars about the store
type Store struct {
	db       *bolt.DB
	idLength int
}

// Entry is the data set which is stored in the DB as JSON
type Entry struct {
	OAuthProvider, OAuthID string
	RemoteAddr             string `json:",omitempty"`
	DeletionURL            string `json:",omitempty"`
	Password               []byte `json:",omitempty"`
	Public                 EntryPublicData
}

// Visitor is the entry which is stored in the visitors bucket
type Visitor struct {
	IP, Referer, UserAgent                                 string
	Timestamp                                              time.Time
	UTMSource, UTMMedium, UTMCampaign, UTMContent, UTMTerm string `json:",omitempty"`
}

// EntryPublicData is the public part of an entry
type EntryPublicData struct {
	CreatedOn             time.Time
	LastVisit, Expiration *time.Time `json:",omitempty"`
	VisitCount            int
	URL                   string
}

// ErrNoEntryFound is returned when no entry to a id is found
var ErrNoEntryFound = errors.New("no entry found with this ID")

// ErrNoValidURL is returned when the URL is not valid
var ErrNoValidURL = errors.New("the given URL is no valid URL")

// ErrGeneratingIDFailed is returned when the 10 tries to generate an id failed
var ErrGeneratingIDFailed = errors.New("could not generate unique id, all ten tries failed")

// ErrEntryIsExpired is returned when the entry is expired
var ErrEntryIsExpired = errors.New("entry is expired")

var (
	shortedURLsBucket      = []byte("shorted")
	shortedIDsToUserBucket = []byte("shorted2Users")
)

// New initializes the store with the db
func New() (*Store, error) {
	db, err := bolt.Open(filepath.Join(util.GetConfig().DataDir, "main.db"), 0644, &bolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		return nil, errors.Wrap(err, "could not open bolt DB database")
	}
	err = db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(shortedURLsBucket)
		return err
	})
	if err != nil {
		return nil, err
	}
	return &Store{
		db:       db,
		idLength: util.GetConfig().ShortedIDLength,
	}, nil
}

// GetEntryByID returns a unmarshalled entry of the db by a given ID
func (s *Store) GetEntryByID(id string) (*Entry, error) {
	if id == "" {
		return nil, ErrNoEntryFound
	}
	rawEntry, err := s.GetEntryByIDRaw(id)
	if err != nil {
		return nil, err
	}
	var entry *Entry
	return entry, json.Unmarshal(rawEntry, &entry)
}

// IncreaseVisitCounter increments the visit counter of an entry
func (s *Store) IncreaseVisitCounter(id string) error {
	entry, err := s.GetEntryByID(id)
	if err != nil {
		return errors.Wrap(err, "could not get entry by ID")
	}
	entry.Public.VisitCount++
	currentTime := time.Now()
	entry.Public.LastVisit = &currentTime
	raw, err := json.Marshal(entry)
	if err != nil {
		return err
	}
	return s.db.Update(func(tx *bolt.Tx) error {
		if err := tx.Bucket(shortedURLsBucket).Put([]byte(id), raw); err != nil {
			return errors.Wrap(err, "could not put updated visitor count JSON into the bucket")
		}
		return nil
	})
}

// GetEntryAndIncrease Increases the visitor count, checks
// if the URL is expired and returns the origin URL
func (s *Store) GetEntryAndIncrease(id string) (*Entry, error) {
	entry, err := s.GetEntryByID(id)
	if err != nil {
		return nil, err
	}
	if entry.Public.Expiration != nil && time.Now().After(*entry.Public.Expiration) {
		return nil, ErrEntryIsExpired
	}
	if err := s.IncreaseVisitCounter(id); err != nil {
		return nil, errors.Wrap(err, "could not increase visitor counter")
	}
	return entry, nil
}

// GetEntryByIDRaw returns the raw data (JSON) of a data set
func (s *Store) GetEntryByIDRaw(id string) ([]byte, error) {
	var raw []byte
	return raw, s.db.View(func(tx *bolt.Tx) error {
		raw = tx.Bucket(shortedURLsBucket).Get([]byte(id))
		if raw == nil {
			return ErrNoEntryFound
		}
		return nil
	})
}

// CreateEntry creates a new record and returns his short id
func (s *Store) CreateEntry(entry Entry, givenID, password string) (string, []byte, error) {
	if !govalidator.IsURL(entry.Public.URL) {
		return "", nil, ErrNoValidURL
	}
	if password != "" {
		var err error
		entry.Password, err = bcrypt.GenerateFromPassword([]byte(password), 10)
		if err != nil {
			return "", nil, errors.Wrap(err, "could not generate bcrypt from password")
		}
	}
	// try it 10 times to make a short URL
	for i := 1; i <= 10; i++ {
		id, delID, err := s.createEntry(entry, givenID)
		if err != nil && givenID != "" {
			return "", nil, err
		} else if err != nil {
			logrus.Debugf("Could not create entry: %v", err)
			continue
		}
		return id, delID, nil
	}
	return "", nil, ErrGeneratingIDFailed
}

// DeleteEntry deletes an Entry fully from the DB
func (s *Store) DeleteEntry(id string, givenHmac []byte) error {
	mac := hmac.New(sha512.New, util.GetPrivateKey())
	if _, err := mac.Write([]byte(id)); err != nil {
		return errors.Wrap(err, "could not write hmac")
	}
	if !hmac.Equal(mac.Sum(nil), givenHmac) {
		return errors.New("hmac verification failed")
	}
	return s.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(shortedURLsBucket)
		if bucket.Get([]byte(id)) == nil {
			return errors.New("entry already deleted")
		}
		if err := bucket.Delete([]byte(id)); err != nil {
			return errors.Wrap(err, "could not delete entry")
		}
		if err := tx.DeleteBucket([]byte(id)); err != nil && err != bolt.ErrBucketNotFound && err != bolt.ErrBucketExists {
			return errors.Wrap(err, "could not delte bucket")
		}
		uTsIDsBucket := tx.Bucket(shortedIDsToUserBucket)
		return uTsIDsBucket.ForEach(func(k, v []byte) error {
			if bytes.Equal(k, []byte(id)) {
				return uTsIDsBucket.Delete(k)
			}
			return nil
		})
	})
}

// RegisterVisit registers an new incoming request in the store
func (s *Store) RegisterVisit(id string, visitor Visitor) {
	requestID := uuid.New()
	logrus.WithFields(logrus.Fields{
		"ClientIP":  visitor.IP,
		"ID":        id,
		"RequestID": requestID,
	}).Info("New redirect was registered...")

	err := s.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte(id))
		if err != nil {
			return errors.Wrap(err, "could not create bucket")
		}
		data, err := json.Marshal(visitor)
		if err != nil {
			return errors.Wrap(err, "could not create json")
		}
		return bucket.Put([]byte(requestID), data)
	})
	if err != nil {
		logrus.Warningf("could not register visit: %v", err)
	}
}

// GetVisitors returns all the visits of a shorted URL
func (s *Store) GetVisitors(id string) ([]Visitor, error) {
	output := []Visitor{}
	return output, s.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte(id))
		if err != nil {
			return errors.Wrap(err, "could not create bucket")
		}
		return bucket.ForEach(func(k, v []byte) error {
			var value Visitor
			if err := json.Unmarshal(v, &value); err != nil {
				return errors.Wrap(err, "could not unmarshal json")
			}
			output = append(output, value)
			return nil
		})
	})
}

// GetUserEntries returns all the shorted URL entries of an user
func (s *Store) GetUserEntries(oAuthProvider, oAuthID string) (map[string]Entry, error) {
	entries := map[string]Entry{}
	return entries, s.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists(shortedIDsToUserBucket)
		if err != nil {
			return errors.Wrap(err, "could not create bucket")
		}
		return bucket.ForEach(func(k, v []byte) error {
			if bytes.Equal(v, []byte(oAuthProvider+oAuthID)) {
				entry, err := s.GetEntryByID(string(k))
				if err != nil {
					return errors.Wrap(err, "could not get entry")
				}
				entries[string(k)] = *entry
			}
			return nil
		})
	})
}

// Close closes the bolt db database
func (s *Store) Close() error {
	return s.db.Close()
}
<commit_msg>Fixed shareX uploading issue, when the filename contains space that an error was thrown<commit_after>// Package store provides support to interact with the entries
package store

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha512"
	"encoding/json"
	"path/filepath"
	"strings"
	"time"

	"github.com/maxibanki/golang-url-shortener/util"
	"github.com/pborman/uuid"
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/bcrypt"

	"github.com/asaskevich/govalidator"
	"github.com/boltdb/bolt"
	"github.com/pkg/errors"
)

// Store holds internal funcs and vars about the store
type Store struct {
	db       *bolt.DB
	idLength int
}

// Entry is the data set which is stored in the DB as JSON
type Entry struct {
	OAuthProvider, OAuthID string
	RemoteAddr             string `json:",omitempty"`
	DeletionURL            string `json:",omitempty"`
	Password               []byte `json:",omitempty"`
	Public                 EntryPublicData
}

// Visitor is the entry which is stored in the visitors bucket
type Visitor struct {
	IP, Referer, UserAgent                                 string
	Timestamp                                              time.Time
	UTMSource, UTMMedium, UTMCampaign, UTMContent, UTMTerm string `json:",omitempty"`
}

// EntryPublicData is the public part of an entry
type EntryPublicData struct {
	CreatedOn             time.Time
	LastVisit, Expiration *time.Time `json:",omitempty"`
	VisitCount            int
	URL                   string
}

// ErrNoEntryFound is returned when no entry to a id is found
var ErrNoEntryFound = errors.New("no entry found with this ID")

// ErrNoValidURL is returned when the URL is not valid
var ErrNoValidURL = errors.New("the given URL is no valid URL")

// ErrGeneratingIDFailed is returned when the 10 tries to generate an id failed
var ErrGeneratingIDFailed = errors.New("could not generate unique id, all ten tries failed")

// ErrEntryIsExpired is returned when the entry is expired
var ErrEntryIsExpired = errors.New("entry is expired")

var (
	shortedURLsBucket      = []byte("shorted")
	shortedIDsToUserBucket = []byte("shorted2Users")
)

// New initializes the store with the db
func New() (*Store, error) {
	db, err := bolt.Open(filepath.Join(util.GetConfig().DataDir, "main.db"), 0644, &bolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		return nil, errors.Wrap(err, "could not open bolt DB database")
	}
	err = db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(shortedURLsBucket)
		return err
	})
	if err != nil {
		return nil, err
	}
	return &Store{
		db:       db,
		idLength: util.GetConfig().ShortedIDLength,
	}, nil
}

// GetEntryByID returns a unmarshalled entry of the db by a given ID
func (s *Store) GetEntryByID(id string) (*Entry, error) {
	if id == "" {
		return nil, ErrNoEntryFound
	}
	rawEntry, err := s.GetEntryByIDRaw(id)
	if err != nil {
		return nil, err
	}
	var entry *Entry
	return entry, json.Unmarshal(rawEntry, &entry)
}

// IncreaseVisitCounter increments the visit counter of an entry
func (s *Store) IncreaseVisitCounter(id string) error {
	entry, err := s.GetEntryByID(id)
	if err != nil {
		return errors.Wrap(err, "could not get entry by ID")
	}
	entry.Public.VisitCount++
	currentTime := time.Now()
	entry.Public.LastVisit = &currentTime
	raw, err := json.Marshal(entry)
	if err != nil {
		return err
	}
	return s.db.Update(func(tx *bolt.Tx) error {
		if err := tx.Bucket(shortedURLsBucket).Put([]byte(id), raw); err != nil {
			return errors.Wrap(err, "could not put updated visitor count JSON into the bucket")
		}
		return nil
	})
}

// GetEntryAndIncrease Increases the visitor count, checks
// if the URL is expired and returns the origin URL
func (s *Store) GetEntryAndIncrease(id string) (*Entry, error) {
	entry, err := s.GetEntryByID(id)
	if err != nil {
		return nil, err
	}
	if entry.Public.Expiration != nil && time.Now().After(*entry.Public.Expiration) {
		return nil, ErrEntryIsExpired
	}
	if err := s.IncreaseVisitCounter(id); err != nil {
		return nil, errors.Wrap(err, "could not increase visitor counter")
	}
	return entry, nil
}

// GetEntryByIDRaw returns the raw data (JSON) of a data set
func (s *Store) GetEntryByIDRaw(id string) ([]byte, error) {
	var raw []byte
	return raw, s.db.View(func(tx *bolt.Tx) error {
		raw = tx.Bucket(shortedURLsBucket).Get([]byte(id))
		if raw == nil {
			return ErrNoEntryFound
		}
		return nil
	})
}

// CreateEntry creates a new record and returns his short id
func (s *Store) CreateEntry(entry Entry, givenID, password string) (string, []byte, error) {
	entry.Public.URL = strings.Replace(entry.Public.URL, " ", "%20", -1)
	if !govalidator.IsURL(entry.Public.URL) {
		return "", nil, ErrNoValidURL
	}
	if password != "" {
		var err error
		entry.Password, err = bcrypt.GenerateFromPassword([]byte(password), 10)
		if err != nil {
			return "", nil, errors.Wrap(err, "could not generate bcrypt from password")
		}
	}
	// try it 10 times to make a short URL
	for i := 1; i <= 10; i++ {
		id, delID, err := s.createEntry(entry, givenID)
		if err != nil && givenID != "" {
			return "", nil, err
		} else if err != nil {
			logrus.Debugf("Could not create entry: %v", err)
			continue
		}
		return id, delID, nil
	}
	return "", nil, ErrGeneratingIDFailed
}

// DeleteEntry deletes an Entry fully from the DB
func (s *Store) DeleteEntry(id string, givenHmac []byte) error {
	mac := hmac.New(sha512.New, util.GetPrivateKey())
	if _, err := mac.Write([]byte(id)); err != nil {
		return errors.Wrap(err, "could not write hmac")
	}
	if !hmac.Equal(mac.Sum(nil), givenHmac) {
		return errors.New("hmac verification failed")
	}
	return s.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(shortedURLsBucket)
		if bucket.Get([]byte(id)) == nil {
			return errors.New("entry already deleted")
		}
		if err := bucket.Delete([]byte(id)); err != nil {
			return errors.Wrap(err, "could not delete entry")
		}
		if err := tx.DeleteBucket([]byte(id)); err != nil && err != bolt.ErrBucketNotFound && err != bolt.ErrBucketExists {
			return errors.Wrap(err, "could not delte bucket")
		}
		uTsIDsBucket := tx.Bucket(shortedIDsToUserBucket)
		return uTsIDsBucket.ForEach(func(k, v []byte) error {
			if bytes.Equal(k, []byte(id)) {
				return uTsIDsBucket.Delete(k)
			}
			return nil
		})
	})
}

// RegisterVisit registers an new incoming request in the store
func (s *Store) RegisterVisit(id string, visitor Visitor) {
	requestID := uuid.New()
	logrus.WithFields(logrus.Fields{
		"ClientIP":  visitor.IP,
		"ID":        id,
		"RequestID": requestID,
	}).Info("New redirect was registered...")

	err := s.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte(id))
		if err != nil {
			return errors.Wrap(err, "could not create bucket")
		}
		data, err := json.Marshal(visitor)
		if err != nil {
			return errors.Wrap(err, "could not create json")
		}
		return bucket.Put([]byte(requestID), data)
	})
	if err != nil {
		logrus.Warningf("could not register visit: %v", err)
	}
}

// GetVisitors returns all the visits of a shorted URL
func (s *Store) GetVisitors(id string) ([]Visitor, error) {
	output := []Visitor{}
	return output, s.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte(id))
		if err != nil {
			return errors.Wrap(err, "could not create bucket")
		}
		return bucket.ForEach(func(k, v []byte) error {
			var value Visitor
			if err := json.Unmarshal(v, &value); err != nil {
				return errors.Wrap(err, "could not unmarshal json")
			}
			output = append(output, value)
			return nil
		})
	})
}

// GetUserEntries returns all the shorted URL entries of an user
func (s *Store) GetUserEntries(oAuthProvider, oAuthID string) (map[string]Entry, error) {
	entries := map[string]Entry{}
	return entries, s.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists(shortedIDsToUserBucket)
		if err != nil {
			return errors.Wrap(err, "could not create bucket")
		}
		return bucket.ForEach(func(k, v []byte) error {
			if bytes.Equal(v, []byte(oAuthProvider+oAuthID)) {
				entry, err := s.GetEntryByID(string(k))
				if err != nil {
					return errors.Wrap(err, "could not get entry")
				}
				entries[string(k)] = *entry
			}
			return nil
		})
	})
}

// Close closes the bolt db database
func (s *Store) Close() error {
	return s.db.Close()
}
<|endoftext|>