<commit_before>package lzma

import (
	"errors"
	"fmt"
	"io"
)

// opLenMargin provides an upper limit for the encoding of a single
// operation. The value assumes that all range-encoded bits require
// actually two bits. A typical operations will be shorter.
const opLenMargin = 10

// OpFinder enables the support of multiple different OpFinder
// algorithms.
type OpFinder interface {
	findOps(s *State, all bool) []operation
	fmt.Stringer
}

type CompressorParams struct {
	LC           int
	LP           int
	PB           int
	DictSize     int64
	ExtraBufSize int64
}

func (p *CompressorParams) Properties() Properties {
	props, err := NewProperties(p.LC, p.LP, p.PB)
	if err != nil {
		panic(err)
	}
	return props
}

// SetProperties sets the LC, LP and PB fields.
func (p *CompressorParams) SetProperties(props Properties) {
	p.LC, p.LP, p.PB = props.LC(), props.LP(), props.PB()
}

// Verify checks parameters for errors.
func (p *CompressorParams) Verify() error {
	if p == nil {
		return lzmaError{"parameters must be non-nil"}
	}
	if err := verifyProperties(p.LC, p.LP, p.PB); err != nil {
		return err
	}
	if !(MinDictSize <= p.DictSize && p.DictSize <= MaxDictSize) {
		return rangeError{"DictSize", p.DictSize}
	}
	if p.DictSize != int64(int(p.DictSize)) {
		return lzmaError{fmt.Sprintf("DictSize %d too large for int", p.DictSize)}
	}
	if p.ExtraBufSize < 0 {
		return negError{"ExtraBufSize", p.ExtraBufSize}
	}
	bufSize := p.DictSize + p.ExtraBufSize
	if bufSize != int64(int(bufSize)) {
		return lzmaError{"buffer size too large for int"}
	}
	return nil
}

type Compressor struct {
	properties Properties
	OpFinder   OpFinder
	state      *State
	re         *rangeEncoder
	dict       *hashDict
	closed     bool
	start      int64
}

func NewCompressor(lzma io.Writer, p CompressorParams) (c *Compressor, err error) {
	if lzma == nil {
		return nil, errors.New("NewCompressor: argument lzma is nil")
	}
	if err = p.Verify(); err != nil {
		return nil, err
	}
	buf, err := newBuffer(p.DictSize + p.ExtraBufSize)
	if err != nil {
		return nil, err
	}
	d, err := newHashDict(buf, buf.bottom, p.DictSize)
	if err != nil {
		return nil, err
	}
	d.sync()
	props := p.Properties()
	state := NewState(props, d)
	c = &Compressor{
		properties: props,
		OpFinder:   Greedy,
		state:      state,
		dict:       d,
		re:         newRangeEncoder(lzma),
		start:      d.head,
	}
	return c, nil
}

func (c *Compressor) Write(p []byte) (n int, err error) {
	if c.closed {
		return 0, errCompressorClosed
	}
	return c.dict.buf.Write(p)
}

func (c *Compressor) Compress(limit int64, all bool) (n int64, err error) {
	if c.closed {
		return 0, errCompressorClosed
	}
	if limit < 0 {
		return 0, errCompressLimit
	}
	ops := c.OpFinder.findOps(c.state, all)
	l := limit - opLenMargin
	start := c.re.Len()
	for _, op := range ops {
		if n > l {
			break
		}
		if err := c.writeOp(op); err != nil {
			return n, err
		}
		n = c.re.Len() - start
	}
	c.dict.sync()
	return n, nil
}

// writeLiteral writes a literal into the operation stream
func (c *Compressor) writeLiteral(l lit) error {
	var err error
	state, state2, _ := c.state.states()
	if err = c.state.isMatch[state2].Encode(c.re, 0); err != nil {
		return err
	}
	litState := c.state.litState()
	match := c.state.dict.byteAt(int64(c.state.rep[0]) + 1)
	err = c.state.litCodec.Encode(c.re, l.b, state, match, litState)
	if err != nil {
		return err
	}
	c.state.updateStateLiteral()
	return nil
}

// iverson implements the Iverson operator as proposed by Donald Knuth in his
// book Concrete Mathematics.
func iverson(ok bool) uint32 {
	if ok {
		return 1
	}
	return 0
}

// writeMatch writes a repetition operation into the operation stream
func (c *Compressor) writeMatch(m match) error {
	var err error
	if !(minDistance <= m.distance && m.distance <= maxDistance) {
		panic(rangeError{"match distance", m.distance})
	}
	dist := uint32(m.distance - minDistance)
	if !(MinLength <= m.n && m.n <= MaxLength) &&
		!(dist == c.state.rep[0] && m.n == 1) {
		panic(rangeError{"match length", m.n})
	}
	state, state2, posState := c.state.states()
	if err = c.state.isMatch[state2].Encode(c.re, 1); err != nil {
		return err
	}
	g := 0
	for ; g < 4; g++ {
		if c.state.rep[g] == dist {
			break
		}
	}
	b := iverson(g < 4)
	if err = c.state.isRep[state].Encode(c.re, b); err != nil {
		return err
	}
	n := uint32(m.n - MinLength)
	if b == 0 {
		// simple match
		c.state.rep[3], c.state.rep[2], c.state.rep[1], c.state.rep[0] =
			c.state.rep[2], c.state.rep[1], c.state.rep[0], dist
		c.state.updateStateMatch()
		if err = c.state.lenCodec.Encode(c.re, n, posState); err != nil {
			return err
		}
		return c.state.distCodec.Encode(c.re, dist, n)
	}
	b = iverson(g != 0)
	if err = c.state.isRepG0[state].Encode(c.re, b); err != nil {
		return err
	}
	if b == 0 {
		// g == 0
		b = iverson(m.n != 1)
		if err = c.state.isRepG0Long[state2].Encode(c.re, b); err != nil {
			return err
		}
		if b == 0 {
			c.state.updateStateShortRep()
			return nil
		}
	} else {
		// g in {1,2,3}
		b = iverson(g != 1)
		if err = c.state.isRepG1[state].Encode(c.re, b); err != nil {
			return err
		}
		if b == 1 {
			// g in {2,3}
			b = iverson(g != 2)
			err = c.state.isRepG2[state].Encode(c.re, b)
			if err != nil {
				return err
			}
			if b == 1 {
				c.state.rep[3] = c.state.rep[2]
			}
			c.state.rep[2] = c.state.rep[1]
		}
		c.state.rep[1] = c.state.rep[0]
		c.state.rep[0] = dist
	}
	c.state.updateStateRep()
	return c.state.repLenCodec.Encode(c.re, n, posState)
}

// discard processes an operation after it has been written into the
// compressed LZMA street by moving the dictionary head forward.
func (c *Compressor) discard(op operation) error {
	k := op.Len()
	n, err := c.dict.move(k)
	if err != nil {
		return fmt.Errorf("operation %s: move %d error %s", op, k, err)
	}
	if n < k {
		return fmt.Errorf("operation %s: move %d incomplete", op, k)
	}
	return nil
}

// writeOp writes an operation value into the stream. It checks whether there
// is still enough space available using an upper limit for the size required.
func (c *Compressor) writeOp(op operation) error {
	var err error
	switch x := op.(type) {
	case match:
		err = c.writeMatch(x)
	case lit:
		err = c.writeLiteral(x)
	}
	if err != nil {
		return err
	}
	err = c.discard(op)
	return err
}

func (c *Compressor) MarkEOS() error {
	panic("TODO")
}

func (c *Compressor) Close() error {
	panic("TODO")
}

func (c *Compressor) ResetState() {
	panic("TODO")
}

func (c *Compressor) SetProperties(p Properties) {
	panic("TODO")
}

func (c *Compressor) ResetDict(p Properties) {
	panic("TODO")
}

func (c *Compressor) SetWriter(w io.Writer) {
	panic("TODO")
}
<commit_msg>lzma: rewrote method Compress of Compressor to support limit argument<commit_after>package lzma

import (
	"errors"
	"fmt"
	"io"
)

// opLenMargin provides an upper limit for the encoding of a single
// operation. The value assumes that all range-encoded bits require
// actually two bits. A typical operations will be shorter.
const opLenMargin = 10

// OpFinder enables the support of multiple different OpFinder
// algorithms.
type OpFinder interface {
	findOps(s *State, all bool) []operation
	fmt.Stringer
}

type CompressorParams struct {
	LC           int
	LP           int
	PB           int
	DictSize     int64
	ExtraBufSize int64
}

func (p *CompressorParams) Properties() Properties {
	props, err := NewProperties(p.LC, p.LP, p.PB)
	if err != nil {
		panic(err)
	}
	return props
}

// SetProperties sets the LC, LP and PB fields.
func (p *CompressorParams) SetProperties(props Properties) {
	p.LC, p.LP, p.PB = props.LC(), props.LP(), props.PB()
}

// Verify checks parameters for errors.
func (p *CompressorParams) Verify() error {
	if p == nil {
		return lzmaError{"parameters must be non-nil"}
	}
	if err := verifyProperties(p.LC, p.LP, p.PB); err != nil {
		return err
	}
	if !(MinDictSize <= p.DictSize && p.DictSize <= MaxDictSize) {
		return rangeError{"DictSize", p.DictSize}
	}
	if p.DictSize != int64(int(p.DictSize)) {
		return lzmaError{fmt.Sprintf("DictSize %d too large for int", p.DictSize)}
	}
	if p.ExtraBufSize < 0 {
		return negError{"ExtraBufSize", p.ExtraBufSize}
	}
	bufSize := p.DictSize + p.ExtraBufSize
	if bufSize != int64(int(bufSize)) {
		return lzmaError{"buffer size too large for int"}
	}
	return nil
}

type Compressor struct {
	properties Properties
	OpFinder   OpFinder
	state      *State
	re         *rangeEncoder
	dict       *hashDict
	closed     bool
	start      int64
}

func NewCompressor(lzma io.Writer, p CompressorParams) (c *Compressor, err error) {
	if lzma == nil {
		return nil, errors.New("NewCompressor: argument lzma is nil")
	}
	if err = p.Verify(); err != nil {
		return nil, err
	}
	buf, err := newBuffer(p.DictSize + p.ExtraBufSize)
	if err != nil {
		return nil, err
	}
	d, err := newHashDict(buf, buf.bottom, p.DictSize)
	if err != nil {
		return nil, err
	}
	d.sync()
	props := p.Properties()
	state := NewState(props, d)
	c = &Compressor{
		properties: props,
		OpFinder:   Greedy,
		state:      state,
		dict:       d,
		re:         newRangeEncoder(lzma),
		start:      d.head,
	}
	return c, nil
}

func (c *Compressor) Write(p []byte) (n int, err error) {
	if c.closed {
		return 0, errCompressorClosed
	}
	return c.dict.buf.Write(p)
}

// writeLiteral writes a literal into the operation stream
func (c *Compressor) writeLiteral(l lit) error {
	var err error
	state, state2, _ := c.state.states()
	if err = c.state.isMatch[state2].Encode(c.re, 0); err != nil {
		return err
	}
	litState := c.state.litState()
	match := c.state.dict.byteAt(int64(c.state.rep[0]) + 1)
	err = c.state.litCodec.Encode(c.re, l.b, state, match, litState)
	if err != nil {
		return err
	}
	c.state.updateStateLiteral()
	return nil
}

// iverson implements the Iverson operator as proposed by Donald Knuth in his
// book Concrete Mathematics.
func iverson(ok bool) uint32 {
	if ok {
		return 1
	}
	return 0
}

// writeMatch writes a repetition operation into the operation stream
func (c *Compressor) writeMatch(m match) error {
	var err error
	if !(minDistance <= m.distance && m.distance <= maxDistance) {
		panic(rangeError{"match distance", m.distance})
	}
	dist := uint32(m.distance - minDistance)
	if !(MinLength <= m.n && m.n <= MaxLength) &&
		!(dist == c.state.rep[0] && m.n == 1) {
		panic(rangeError{"match length", m.n})
	}
	state, state2, posState := c.state.states()
	if err = c.state.isMatch[state2].Encode(c.re, 1); err != nil {
		return err
	}
	g := 0
	for ; g < 4; g++ {
		if c.state.rep[g] == dist {
			break
		}
	}
	b := iverson(g < 4)
	if err = c.state.isRep[state].Encode(c.re, b); err != nil {
		return err
	}
	n := uint32(m.n - MinLength)
	if b == 0 {
		// simple match
		c.state.rep[3], c.state.rep[2], c.state.rep[1], c.state.rep[0] =
			c.state.rep[2], c.state.rep[1], c.state.rep[0], dist
		c.state.updateStateMatch()
		if err = c.state.lenCodec.Encode(c.re, n, posState); err != nil {
			return err
		}
		return c.state.distCodec.Encode(c.re, dist, n)
	}
	b = iverson(g != 0)
	if err = c.state.isRepG0[state].Encode(c.re, b); err != nil {
		return err
	}
	if b == 0 {
		// g == 0
		b = iverson(m.n != 1)
		if err = c.state.isRepG0Long[state2].Encode(c.re, b); err != nil {
			return err
		}
		if b == 0 {
			c.state.updateStateShortRep()
			return nil
		}
	} else {
		// g in {1,2,3}
		b = iverson(g != 1)
		if err = c.state.isRepG1[state].Encode(c.re, b); err != nil {
			return err
		}
		if b == 1 {
			// g in {2,3}
			b = iverson(g != 2)
			err = c.state.isRepG2[state].Encode(c.re, b)
			if err != nil {
				return err
			}
			if b == 1 {
				c.state.rep[3] = c.state.rep[2]
			}
			c.state.rep[2] = c.state.rep[1]
		}
		c.state.rep[1] = c.state.rep[0]
		c.state.rep[0] = dist
	}
	c.state.updateStateRep()
	return c.state.repLenCodec.Encode(c.re, n, posState)
}

// discard processes an operation after it has been written into the
// compressed LZMA street by moving the dictionary head forward.
func (c *Compressor) discard(op operation) error {
	k := op.Len()
	n, err := c.dict.move(k)
	if err != nil {
		return fmt.Errorf("operation %s: move %d error %s", op, k, err)
	}
	if n < k {
		return fmt.Errorf("operation %s: move %d incomplete", op, k)
	}
	return nil
}

// writeOp writes an operation value into the stream. It checks whether there
// is still enough space available using an upper limit for the size required.
func (c *Compressor) writeOp(op operation) error {
	var err error
	switch x := op.(type) {
	case match:
		err = c.writeMatch(x)
	case lit:
		err = c.writeLiteral(x)
	}
	if err != nil {
		return err
	}
	err = c.discard(op)
	return err
}

func (c *Compressor) Compress(limit int64, all bool) (n int64, err error) {
	if c.closed {
		return 0, errCompressorClosed
	}
	if limit < opLenMargin {
		return 0, errCompressLimit
	}
	ops := c.OpFinder.findOps(c.state, all)
	start := c.re.Len()
	end := start + limit - opLenMargin
	for _, op := range ops {
		if err := c.writeOp(op); err != nil {
			return 0, err
		}
		if c.re.Len() >= end {
			break
		}
	}
	c.dict.sync()
	return c.re.Len() - start, nil
}

func (c *Compressor) MarkEOS() error {
	panic("TODO")
}

func (c *Compressor) Close() error {
	panic("TODO")
}

func (c *Compressor) ResetState() {
	panic("TODO")
}

func (c *Compressor) SetProperties(p Properties) {
	panic("TODO")
}

func (c *Compressor) ResetDict(p Properties) {
	panic("TODO")
}

func (c *Compressor) SetWriter(w io.Writer) {
	panic("TODO")
}
<|endoftext|>


<commit_before>package ninja

import (
	"encoding/json"
	"fmt"
	"log"
	"net/url"
	"strings"
	"sync"

	"git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"

	"github.com/ninjasphere/go-ninja/config"
	"github.com/ninjasphere/go-ninja/logger"
	"github.com/ninjasphere/go-ninja/model"
	"github.com/ninjasphere/go-ninja/rpc"
	"github.com/ninjasphere/go-ninja/rpc/json2"
)

// Connection Holds the connection to the Ninja MQTT bus, and provides all the methods needed to communicate with
// the other modules in Sphere.
type Connection struct {
	mqtt      *mqtt.MqttClient
	log       *logger.Logger
	rpc       *rpc.Client
	rpcServer *rpc.Server
}

type Driver interface {
	GetModuleInfo() *model.Module
	SetEventHandler(func(event string, payload interface{}) error)
}

type Device interface {
	GetDriver() Driver
	GetDeviceInfo() *model.Device
	SetEventHandler(func(event string, payload interface{}) error)
}

type Channel interface {
	GetProtocol() string
	SetEventHandler(func(event string, payload interface{}) error)
}

// Connect Builds a new ninja connection to the MQTT broker, using the given client ID
func Connect(clientID string) (*Connection, error) {

	log := logger.GetLogger("connection")

	conn := Connection{log: log}

	mqttURL := fmt.Sprintf("tcp://%s:%d", config.MustString("mqtt", "host"), config.MustInt("mqtt", "port"))

	opts := mqtt.NewClientOptions().AddBroker(mqttURL).SetClientId(clientID).SetCleanSession(true)
	conn.mqtt = mqtt.NewClient(opts)

	if _, err := conn.mqtt.Start(); err != nil {
		return nil, err
	}

	conn.rpc = rpc.NewClient(conn.mqtt, json2.NewClientCodec())
	conn.rpcServer = rpc.NewServer(conn.mqtt, json2.NewCodec())

	log.Infof("Connected to %s using cid:%s", mqttURL, clientID)

	job, err := CreateStatusJob(&conn, clientID)

	if err != nil {
		return nil, err
	}
	job.Start()

	return &conn, nil
}

// GetMqttClient will be removed in a later version. All communication should happen via methods on Connection
func (c *Connection) GetMqttClient() *mqtt.MqttClient {
	return c.mqtt
}

type rpcMessage struct {
	Params *json.RawMessage `json:"params"`
}

// Subscribe allows you to subscribe to an MQTT topic. Topics can contain variables of the form ":myvar" which will
// be returned in the values map in the callback. The callback must return "true" if it wants to receive more messages.
func (c *Connection) Subscribe(topic string, callback func(params *json.RawMessage, values map[string]string) bool) error {

	filter, err := mqtt.NewTopicFilter(GetSubscribeTopic(topic), 0)
	if err != nil {
		c.log.FatalError(err, "Failed to subscribe to "+topic)
	}

	finished := false
	mutex := &sync.Mutex{}

	receipt, err := c.mqtt.StartSubscription(func(_ *mqtt.MqttClient, message mqtt.Message) {
		// We lock so that the callback has a chance to return false,
		// to prevent any more messages arriving on this subscription
		mutex.Lock()

		// TODO: Implement unsubscribing. For now, it will just skip over any subscriptions that have finished
		if finished {
			return
		}

		values, ok := MatchTopicPattern(topic, message.Topic())
		if !ok {
			c.log.Warningf("Failed to read params from topic: %s using template: %s", message.Topic(), topic)
			mutex.Unlock()
		} else {

			msg := &rpcMessage{}
			err := json.Unmarshal(message.Payload(), msg)

			if err != nil {
				c.log.Warningf("Failed to read parameters in rpc call to %s", message.Topic())
				return
			}

			var params json.RawMessage

			json2.ReadRPCParams(msg.Params, &params)
			if err != nil {
				c.log.Warningf("Failed to read parameters in rpc call to %s", message.Topic())
				return
			}

			// The callback needs to be run in a goroutine as blocking this thread prevents any other messages arriving
			go func() {
				if !callback(&params, *values) {
					// The callback has returned false, indicating that it does not want to receive any more messages.
					finished = true
				}
				mutex.Unlock()
			}()

		}
	}, filter)

	if err != nil {
		return err
	}

	<-receipt
	return nil
}

// GetServiceClient returns an RPC client for the given service.
func (c *Connection) GetServiceClient(serviceTopic string) *ServiceClient {
	return &ServiceClient{c, serviceTopic}
}

// ExportDriver Exports a driver using the 'driver' protocol, and announces it
func (c *Connection) ExportDriver(driver Driver) error {
	topic := fmt.Sprintf("$node/%s/driver/%s", config.Serial(), driver.GetModuleInfo().ID)

	announcement := driver.GetModuleInfo()

	announcement.ServiceAnnouncement = model.ServiceAnnouncement{
		Schema: "http://schema.ninjablocks.com/service/driver",
	}

	_, err := c.exportService(driver, topic, announcement)

	if err != nil {
		return err
	}

	return nil
}

// ExportDevice Exports a device using the 'device' protocol, and announces it
func (c *Connection) ExportDevice(device Device) error {
	announcement := device.GetDeviceInfo()
	announcement.ID = getGUID(device.GetDeviceInfo().NaturalIDType, device.GetDeviceInfo().NaturalID)

	topic := fmt.Sprintf("$device/%s", announcement.ID)

	announcement.ServiceAnnouncement = model.ServiceAnnouncement{
		Schema: "http://schema.ninjablocks.com/service/device",
	}

	_, err := c.exportService(device, topic, announcement)

	if err != nil {
		return err
	}

	return nil
}

// ExportChannel Exports a device using the given protocol, and announces it
func (c *Connection) ExportChannel(device Device, channel Channel, id string) error {
	return c.ExportChannelWithSupported(device, channel, id, nil, nil)
}

// ExportChannelWithSupported is the same as ExportChannel, but any methods provided must actually be exported by the
// channel, or an error is returned
func (c *Connection) ExportChannelWithSupported(device Device, channel Channel, id string, supportedMethods *[]string, supportedEvents *[]string) error {
	if channel.GetProtocol() == "" {
		return fmt.Errorf("The channel must have a protocol. Channel ID: %s", id)
	}

	announcement := &model.Channel{
		ID:       id,
		Protocol: channel.GetProtocol(),
		Device:   device.GetDeviceInfo(),
	}

	topic := fmt.Sprintf("$device/%s/channel/%s", device.GetDeviceInfo().ID, id)

	announcement.ServiceAnnouncement = model.ServiceAnnouncement{
		Schema:           resolveProtocolURI(channel.GetProtocol()),
		SupportedMethods: supportedMethods,
		SupportedEvents:  supportedEvents,
	}

	_, err := c.exportService(channel, topic, announcement)

	if err != nil {
		return err
	}

	// <TEMPORARY> - Expose channels using the old topic (with the protocol)
	properAnnouncement := announcement.ServiceAnnouncement

	shortProtocol := strings.TrimPrefix(resolveProtocolURI(channel.GetProtocol()), protocolSchemaURL.String())
	oldTopic := fmt.Sprintf("$device/%s/channel/%s/%s", device.GetDeviceInfo().ID, id, shortProtocol)

	deprecated := true
	announcement.ServiceAnnouncement = model.ServiceAnnouncement{
		Schema:           resolveProtocolURI(channel.GetProtocol()),
		SupportedMethods: supportedMethods,
		SupportedEvents:  supportedEvents,
		Deprecated:       &deprecated,
	}

	_, err = c.exportService(channel, oldTopic, announcement)

	announcement.ServiceAnnouncement = properAnnouncement

	if err != nil {
		return err
	}
	// </TEMPORARY>

	return nil
}

type eventingService interface {
	SetEventHandler(func(event string, payload interface{}) error)
}

type serviceAnnouncement interface {
	GetServiceAnnouncement() *model.ServiceAnnouncement
}

// exportService Exports an RPC service, and announces it over TOPIC/event/announce
func (c *Connection) exportService(service interface{}, topic string, announcement serviceAnnouncement) (*rpc.ExportedService, error) {

	exportedService, err := c.rpcServer.RegisterService(service, topic, announcement.GetServiceAnnouncement().Schema)

	if err != nil {
		return nil, fmt.Errorf("Failed to register service on %s : %s", topic, err)
	}

	if announcement.GetServiceAnnouncement().SupportedMethods == nil {
		announcement.GetServiceAnnouncement().SupportedMethods = &exportedService.Methods
	} else {
		// TODO: Check that all strings in announcement.SupportedMethods exist in exportedService.Methods
		if len(*announcement.GetServiceAnnouncement().SupportedMethods) > len(exportedService.Methods) {
			return nil, fmt.Errorf("The number of actual exported methods is less than the number said to be exported. Check the method signatures of the service. topic:%s", topic)
		}
	}

	if announcement.GetServiceAnnouncement().SupportedEvents == nil {
		events := []string{}
		announcement.GetServiceAnnouncement().SupportedEvents = &events
	}

	// send out service announcement
	err = exportedService.SendEvent("announce", announcement)
	if err != nil {
		return nil, fmt.Errorf("Failed sending service announcement: %s", err)
	}

	c.log.Debugf("Exported service on topic: %s (schema: %s) with methods: %s", topic, announcement.GetServiceAnnouncement().Schema, strings.Join(*announcement.GetServiceAnnouncement().SupportedMethods, ", "))

	switch service := service.(type) {
	case eventingService:
		service.SetEventHandler(func(event string, payload interface{}) error {
			return exportedService.SendEvent(event, payload)
		})
	}

	return exportedService, nil
}

// SendNotification Sends a simple json-rpc notification to a topic
func (c *Connection) SendNotification(topic string, params ...interface{}) error {
	return c.rpcServer.SendNotification(topic, params...)
}

// Pull this out into the schema validation package when we have one
var rootSchemaURL, _ = url.Parse("http://schemas.ninjablocks.com")
var protocolSchemaURL, _ = url.Parse("http://schemas.ninjablocks.com/protocol/")

func resolveSchemaURI(uri string) string {
	return resolveSchemaURIWithBase(rootSchemaURL, uri)
}

func resolveProtocolURI(uri string) string {
	return resolveSchemaURIWithBase(protocolSchemaURL, uri)
}

func resolveSchemaURIWithBase(base *url.URL, uri string) string {

	u, err := url.Parse(uri)
	if err != nil {
		log.Fatalf("Expected URL to parse: %q, got error: %v", uri, err)
	}
	return base.ResolveReference(u).String()
}
<commit_msg>Add the topic to all service announcements<commit_after>package ninja

import (
	"encoding/json"
	"fmt"
	"log"
	"net/url"
	"strings"
	"sync"

	"git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"

	"github.com/ninjasphere/go-ninja/config"
	"github.com/ninjasphere/go-ninja/logger"
	"github.com/ninjasphere/go-ninja/model"
	"github.com/ninjasphere/go-ninja/rpc"
	"github.com/ninjasphere/go-ninja/rpc/json2"
)

// Connection Holds the connection to the Ninja MQTT bus, and provides all the methods needed to communicate with
// the other modules in Sphere.
type Connection struct {
	mqtt      *mqtt.MqttClient
	log       *logger.Logger
	rpc       *rpc.Client
	rpcServer *rpc.Server
}

type Driver interface {
	GetModuleInfo() *model.Module
	SetEventHandler(func(event string, payload interface{}) error)
}

type Device interface {
	GetDriver() Driver
	GetDeviceInfo() *model.Device
	SetEventHandler(func(event string, payload interface{}) error)
}

type Channel interface {
	GetProtocol() string
	SetEventHandler(func(event string, payload interface{}) error)
}

// Connect Builds a new ninja connection to the MQTT broker, using the given client ID
func Connect(clientID string) (*Connection, error) {

	log := logger.GetLogger("connection")

	conn := Connection{log: log}

	mqttURL := fmt.Sprintf("tcp://%s:%d", config.MustString("mqtt", "host"), config.MustInt("mqtt", "port"))

	opts := mqtt.NewClientOptions().AddBroker(mqttURL).SetClientId(clientID).SetCleanSession(true)
	conn.mqtt = mqtt.NewClient(opts)

	if _, err := conn.mqtt.Start(); err != nil {
		return nil, err
	}

	conn.rpc = rpc.NewClient(conn.mqtt, json2.NewClientCodec())
	conn.rpcServer = rpc.NewServer(conn.mqtt, json2.NewCodec())

	log.Infof("Connected to %s using cid:%s", mqttURL, clientID)

	job, err := CreateStatusJob(&conn, clientID)

	if err != nil {
		return nil, err
	}
	job.Start()

	return &conn, nil
}

// GetMqttClient will be removed in a later version. All communication should happen via methods on Connection
func (c *Connection) GetMqttClient() *mqtt.MqttClient {
	return c.mqtt
}

type rpcMessage struct {
	Params *json.RawMessage `json:"params"`
}

// Subscribe allows you to subscribe to an MQTT topic. Topics can contain variables of the form ":myvar" which will
// be returned in the values map in the callback. The callback must return "true" if it wants to receive more messages.
func (c *Connection) Subscribe(topic string, callback func(params *json.RawMessage, values map[string]string) bool) error {

	filter, err := mqtt.NewTopicFilter(GetSubscribeTopic(topic), 0)
	if err != nil {
		c.log.FatalError(err, "Failed to subscribe to "+topic)
	}

	finished := false
	mutex := &sync.Mutex{}

	receipt, err := c.mqtt.StartSubscription(func(_ *mqtt.MqttClient, message mqtt.Message) {
		// We lock so that the callback has a chance to return false,
		// to prevent any more messages arriving on this subscription
		mutex.Lock()

		// TODO: Implement unsubscribing. For now, it will just skip over any subscriptions that have finished
		if finished {
			return
		}

		values, ok := MatchTopicPattern(topic, message.Topic())
		if !ok {
			c.log.Warningf("Failed to read params from topic: %s using template: %s", message.Topic(), topic)
			mutex.Unlock()
		} else {

			msg := &rpcMessage{}
			err := json.Unmarshal(message.Payload(), msg)

			if err != nil {
				c.log.Warningf("Failed to read parameters in rpc call to %s", message.Topic())
				return
			}

			var params json.RawMessage

			json2.ReadRPCParams(msg.Params, &params)
			if err != nil {
				c.log.Warningf("Failed to read parameters in rpc call to %s", message.Topic())
				return
			}

			// The callback needs to be run in a goroutine as blocking this thread prevents any other messages arriving
			go func() {
				if !callback(&params, *values) {
					// The callback has returned false, indicating that it does not want to receive any more messages.
					finished = true
				}
				mutex.Unlock()
			}()

		}
	}, filter)

	if err != nil {
		return err
	}

	<-receipt
	return nil
}

// GetServiceClient returns an RPC client for the given service.
func (c *Connection) GetServiceClient(serviceTopic string) *ServiceClient {
	return &ServiceClient{c, serviceTopic}
}

// ExportDriver Exports a driver using the 'driver' protocol, and announces it
func (c *Connection) ExportDriver(driver Driver) error {
	topic := fmt.Sprintf("$node/%s/driver/%s", config.Serial(), driver.GetModuleInfo().ID)

	announcement := driver.GetModuleInfo()

	announcement.ServiceAnnouncement = model.ServiceAnnouncement{
		Schema: "http://schema.ninjablocks.com/service/driver",
	}

	_, err := c.exportService(driver, topic, announcement)

	if err != nil {
		return err
	}

	return nil
}

// ExportDevice Exports a device using the 'device' protocol, and announces it
func (c *Connection) ExportDevice(device Device) error {
	announcement := device.GetDeviceInfo()
	announcement.ID = getGUID(device.GetDeviceInfo().NaturalIDType, device.GetDeviceInfo().NaturalID)

	topic := fmt.Sprintf("$device/%s", announcement.ID)

	announcement.ServiceAnnouncement = model.ServiceAnnouncement{
		Schema: "http://schema.ninjablocks.com/service/device",
	}

	_, err := c.exportService(device, topic, announcement)

	if err != nil {
		return err
	}

	return nil
}

// ExportChannel Exports a device using the given protocol, and announces it
func (c *Connection) ExportChannel(device Device, channel Channel, id string) error {
	return c.ExportChannelWithSupported(device, channel, id, nil, nil)
}

// ExportChannelWithSupported is the same as ExportChannel, but any methods provided must actually be exported by the
// channel, or an error is returned
func (c *Connection) ExportChannelWithSupported(device Device, channel Channel, id string, supportedMethods *[]string, supportedEvents *[]string) error {
	if channel.GetProtocol() == "" {
		return fmt.Errorf("The channel must have a protocol. Channel ID: %s", id)
	}

	announcement := &model.Channel{
		ID:       id,
		Protocol: channel.GetProtocol(),
		Device:   device.GetDeviceInfo(),
	}

	topic := fmt.Sprintf("$device/%s/channel/%s", device.GetDeviceInfo().ID, id)

	announcement.ServiceAnnouncement = model.ServiceAnnouncement{
		Schema:           resolveProtocolURI(channel.GetProtocol()),
		SupportedMethods: supportedMethods,
		SupportedEvents:  supportedEvents,
	}

	_, err := c.exportService(channel, topic, announcement)

	if err != nil {
		return err
	}

	// <TEMPORARY> - Expose channels using the old topic (with the protocol)
	properAnnouncement := announcement.ServiceAnnouncement

	shortProtocol := strings.TrimPrefix(resolveProtocolURI(channel.GetProtocol()), protocolSchemaURL.String())
	oldTopic := fmt.Sprintf("$device/%s/channel/%s/%s", device.GetDeviceInfo().ID, id, shortProtocol)

	deprecated := true
	announcement.ServiceAnnouncement = model.ServiceAnnouncement{
		Schema:           resolveProtocolURI(channel.GetProtocol()),
		SupportedMethods: supportedMethods,
		SupportedEvents:  supportedEvents,
		Deprecated:       &deprecated,
	}

	_, err = c.exportService(channel, oldTopic, announcement)

	announcement.ServiceAnnouncement = properAnnouncement

	if err != nil {
		return err
	}
	// </TEMPORARY>

	return nil
}

type eventingService interface {
	SetEventHandler(func(event string, payload interface{}) error)
}

type serviceAnnouncement interface {
	GetServiceAnnouncement() *model.ServiceAnnouncement
}

// exportService Exports an RPC service, and announces it over TOPIC/event/announce
func (c *Connection) exportService(service interface{}, topic string, announcement serviceAnnouncement) (*rpc.ExportedService, error) {

	exportedService, err := c.rpcServer.RegisterService(service, topic, announcement.GetServiceAnnouncement().Schema)

	if err != nil {
		return nil, fmt.Errorf("Failed to register service on %s : %s", topic, err)
	}

	if announcement.GetServiceAnnouncement().SupportedMethods == nil {
		announcement.GetServiceAnnouncement().SupportedMethods = &exportedService.Methods
	} else {
		// TODO: Check that all strings in announcement.SupportedMethods exist in exportedService.Methods
		if len(*announcement.GetServiceAnnouncement().SupportedMethods) > len(exportedService.Methods) {
			return nil, fmt.Errorf("The number of actual exported methods is less than the number said to be exported. Check the method signatures of the service. topic:%s", topic)
		}
	}

	if announcement.GetServiceAnnouncement().SupportedEvents == nil {
		events := []string{}
		announcement.GetServiceAnnouncement().SupportedEvents = &events
	}

	announcement.GetServiceAnnouncement().Topic = topic

	// send out service announcement
	err = exportedService.SendEvent("announce", announcement)
	if err != nil {
		return nil, fmt.Errorf("Failed sending service announcement: %s", err)
	}

	c.log.Debugf("Exported service on topic: %s (schema: %s) with methods: %s", topic, announcement.GetServiceAnnouncement().Schema, strings.Join(*announcement.GetServiceAnnouncement().SupportedMethods, ", "))

	switch service := service.(type) {
	case eventingService:
		service.SetEventHandler(func(event string, payload interface{}) error {
			return exportedService.SendEvent(event, payload)
		})
	}

	return exportedService, nil
}

// SendNotification Sends a simple json-rpc notification to a topic
func (c *Connection) SendNotification(topic string, params ...interface{}) error {
	return c.rpcServer.SendNotification(topic, params...)
}

// Pull this out into the schema validation package when we have one
var rootSchemaURL, _ = url.Parse("http://schemas.ninjablocks.com")
var protocolSchemaURL, _ = url.Parse("http://schemas.ninjablocks.com/protocol/")

func resolveSchemaURI(uri string) string {
	return resolveSchemaURIWithBase(rootSchemaURL, uri)
}

func resolveProtocolURI(uri string) string {
	return resolveSchemaURIWithBase(protocolSchemaURL, uri)
}

func resolveSchemaURIWithBase(base *url.URL, uri string) string {

	u, err := url.Parse(uri)
	if err != nil {
		log.Fatalf("Expected URL to parse: %q, got error: %v", uri, err)
	}
	return base.ResolveReference(u).String()
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Mattermost, Inc. All Rights Reserved.
// See License.txt for license information.

package api

import (
	"github.com/dotcominternet/platform/model"
)

type MeProvider struct {
}

const (
	CMD_ME = "me"
)

func init() {
	RegisterCommandProvider(&MeProvider{})
}

func (me *MeProvider) GetTrigger() string {
	return CMD_ME
}

func (me *MeProvider) GetCommand(c *Context) *model.Command {
	return &model.Command{
		Trigger:          CMD_ME,
		AutoComplete:     true,
		AutoCompleteDesc: c.T("api.command_me.desc"),
		AutoCompleteHint: c.T("api.command_me.hint"),
		DisplayName:      c.T("api.command_me.name"),
	}
}

func (me *MeProvider) DoCommand(c *Context, channelId string, message string) *model.CommandResponse {
	return &model.CommandResponse{ResponseType: model.COMMAND_RESPONSE_TYPE_IN_CHANNEL, Text: "*" + message + "*"}
}
<commit_msg>Make the /me command prefix the action with the invoking user's name, and then assign an "action" class property to alter the appearance of the post.<commit_after>// Copyright (c) 2016 Mattermost, Inc. All Rights Reserved.
// See License.txt for license information.

package api

import (
	"github.com/dotcominternet/platform/model"
)

type MeProvider struct {
}

const (
	CMD_ME = "me"
)

func init() {
	RegisterCommandProvider(&MeProvider{})
}

func (me *MeProvider) GetTrigger() string {
	return CMD_ME
}

func (me *MeProvider) GetCommand(c *Context) *model.Command {
	return &model.Command{
		Trigger:          CMD_ME,
		AutoComplete:     true,
		AutoCompleteDesc: c.T("api.command_me.desc"),
		AutoCompleteHint: c.T("api.command_me.hint"),
		DisplayName:      c.T("api.command_me.name"),
	}
}

func (me *MeProvider) DoCommand(c *Context, channelId string, message string) *model.CommandResponse {
	userChan := Srv.Store.User().Get(c.Session.UserId)
	var user *model.User
	if ur := <-userChan; ur.Err != nil {
		c.Err = ur.Err
		return nil
	} else {
		user = ur.Data.(*model.User)
	}

	var name = user.Username
	if len(user.Nickname) > 0 {
		name = user.Nickname
	} else if len(user.FirstName) > 0 {
		name = user.FirstName
	}

	return &model.CommandResponse{
		ResponseType: model.COMMAND_RESPONSE_TYPE_IN_CHANNEL,
		Text: "*" + name + " " + message + "*",
		Props: model.StringInterface{
			"class": "action",
		},
	}
}
<|endoftext|>


<commit_before>package lib

import (
	"bufio"
	"fmt"
	"net"
	"os/exec"
	"strings"

	"github.com/subutai-io/base/agent/config"
	"github.com/subutai-io/base/agent/log"
)

func Tunnel(dst, timeout, dstport string, global bool) {
	var tunsrv string
	var args []string
	if dst == "" || timeout == "" {
		log.Error("Please specify destination IP and timeout")
	}

	if len(dstport) == 0 {
		dstport = "22"
	}

	if global {
		cdn, err := net.LookupIP(config.Cdn.Url)
		tunsrv = cdn[0].String()
		log.Check(log.ErrorLevel, "Resolving nearest tunnel node address", err)
		args = []string{"-i", config.Agent.AppPrefix + "etc/ssh.pem", "-N", "-f", "-R", "0:" + dst + ":" + dstport, "-o", "StrictHostKeyChecking=no", "tunnel@" + tunsrv, "sleep", timeout}
	} else {
		wan, err := net.InterfaceByName("wan")
		log.Check(log.ErrorLevel, "Getting WAN interface info", err)
		wanIP, err := wan.Addrs()
		log.Check(log.ErrorLevel, "Getting WAN interface addresses", err)
		if len(wanIP) > 0 {
			ip := strings.Split(wanIP[0].String(), "/")
			if len(ip) > 0 {
				tunsrv = ip[0]
			}
		}
		args = []string{"-N", "-f", "-R", "0:" + dst + ":" + dstport, "-o", "StrictHostKeyChecking=no", "ubuntu@" + tunsrv, "sleep", timeout}
	}

	cmd := exec.Command("ssh", args...)
	stderr, _ := cmd.StderrPipe()
	log.Check(log.FatalLevel, "Creating SSH tunnel to "+dst, cmd.Start())
	r := bufio.NewReader(stderr)
	i := 0
	for line, _, err := r.ReadLine(); err == nil && i < 20; i++ {
		log.Check(log.FatalLevel, "Reading SSH pipe", err)
		if strings.Contains(string(line), "Allocated port") {
			port := strings.Fields(string(line))
			if global {
				fmt.Println(tunsrv + ":" + port[2])
			} else {
				fmt.Println(port[2])
			}
			return
		}
		line, _, err = r.ReadLine()
	}
	log.Error("Cannot get tunnel port")
}
<commit_msg>Tunnel default port changed<commit_after>package lib

import (
	"bufio"
	"fmt"
	"net"
	"os/exec"
	"strings"

	"github.com/subutai-io/base/agent/config"
	"github.com/subutai-io/base/agent/log"
)

func Tunnel(dst, timeout, dstport string, global bool) {
	var tunsrv string
	var args []string
	if dst == "" || timeout == "" {
		log.Error("Please specify destination IP and timeout")
	}

	if len(dstport) == 0 {
		dstport = "22"
	}

	if global {
		cdn, err := net.LookupIP(config.Cdn.Url)
		tunsrv = cdn[0].String()
		log.Check(log.ErrorLevel, "Resolving nearest tunnel node address", err)
		args = []string{"-i", config.Agent.AppPrefix + "etc/ssh.pem", "-N", "-p", "8022", "-f", "-R", "0:" + dst + ":" + dstport, "-o", "StrictHostKeyChecking=no", "tunnel@" + tunsrv, "sleep", timeout}
	} else {
		wan, err := net.InterfaceByName("wan")
		log.Check(log.ErrorLevel, "Getting WAN interface info", err)
		wanIP, err := wan.Addrs()
		log.Check(log.ErrorLevel, "Getting WAN interface addresses", err)
		if len(wanIP) > 0 {
			ip := strings.Split(wanIP[0].String(), "/")
			if len(ip) > 0 {
				tunsrv = ip[0]
			}
		}
		args = []string{"-N", "-f", "-R", "0:" + dst + ":" + dstport, "-o", "StrictHostKeyChecking=no", "ubuntu@" + tunsrv, "sleep", timeout}
	}

	cmd := exec.Command("ssh", args...)
	stderr, _ := cmd.StderrPipe()
	log.Check(log.FatalLevel, "Creating SSH tunnel to "+dst, cmd.Start())
	r := bufio.NewReader(stderr)
	i := 0
	for line, _, err := r.ReadLine(); err == nil && i < 20; i++ {
		log.Check(log.FatalLevel, "Reading SSH pipe", err)
		if strings.Contains(string(line), "Allocated port") {
			port := strings.Fields(string(line))
			if global {
				fmt.Println(tunsrv + ":" + port[2])
			} else {
				fmt.Println(port[2])
			}
			return
		}
		line, _, err = r.ReadLine()
	}
	log.Error("Cannot get tunnel port")
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"github.com/buildkite/agent/api"
	"github.com/buildkite/agent/logger"
	"github.com/buildkite/agent/process"
	"github.com/buildkite/agent/retry"
	"os"
	"path/filepath"
	_ "regexp"
	"time"
)

type JobRunner struct {
	// The job being run
	Job *api.Job

	// The APIClient that will be used when updating the job
	APIClient *api.Client

	// The endpoint that should be used when communicating with the API
	Endpoint string

	// The registred agent API record running this job
	Agent *api.Agent

	// The configuration of the agent from the CLI
	AgentConfiguration *AgentConfiguration

	// The interal process of the job
	process *process.Process

	// The internal header time streamer
	headerTimesStreamer *HeaderTimesStreamer

	// The internal log streamer
	logStreamer *LogStreamer

	// If the job is being cancelled
	cancelled bool
}

// Initializes the job runner
func (r JobRunner) Create() (runner *JobRunner, err error) {
	runner = &r

	// Our own APIClient using the endpoint and the agents access token
	runner.APIClient = APIClient{Endpoint: r.Endpoint, Token: r.Agent.AccessToken}.Create()

	// // Create our header times struct
	runner.headerTimesStreamer = &HeaderTimesStreamer{Callback: r.onUploadHeaderTime}

	// The log streamer that will take the output chunks, and send them to
	// the Buildkite Agent API
	runner.logStreamer = LogStreamer{MaxChunkSizeBytes: r.Job.ChunksMaxSizeBytes, Callback: r.onUploadChunk}.New()

	// The process that will run the bootstrap script
	runner.process = process.Process{
		Script:        r.AgentConfiguration.BootstrapScript,
		Env:           r.createEnvironment(),
		PTY:           r.AgentConfiguration.RunInPty,
		StartCallback: r.onProcessStartCallback,
		LineCallback:  runner.headerTimesStreamer.Scan,
	}.Create()

	return
}

// Runs the job
func (r *JobRunner) Run() error {
	logger.Info("Starting job %s", r.Job.ID)

	// Start the log streamer
	if err := r.logStreamer.Start(); err != nil {
		return err
	}

	// Start the build in the Buildkite Agent API
	if err := r.startJob(time.Now()); err != nil {
		return err
	}

	// Start the process. This will block until it finishes.
	if err := r.process.Start(); err != nil {
		// Send the error as output
		r.logStreamer.Process(fmt.Sprintf("%s", err))
	} else {
		// Add the final output to the streamer
		r.logStreamer.Process(r.process.Output())
	}

	// Store the finished at time
	finishedAt := time.Now()

	// Wait until all the header times have finished uploading
	logger.Debug("[HeaderTimes] Waiting for all times to finish uploading")

	r.headerTimesStreamer.Wait()

	// Stop the log streamer. This will block until all the chunks have
	// been uploaded
	r.logStreamer.Stop()

	// Warn about failed chunks
	if r.logStreamer.ChunksFailedCount > 0 {
		logger.Warn("%d chunks failed to upload for this job", r.logStreamer.ChunksFailedCount)
	}

	// Finish the build in the Buildkite Agent API
	r.finishJob(finishedAt, r.process.ExitStatus, int(r.logStreamer.ChunksFailedCount))

	logger.Info("Finished job %s", r.Job.ID)

	return nil
}

func (r *JobRunner) Kill() error {
	if !r.cancelled {
		logger.Info("Canceling job %s", r.Job.ID)
		r.cancelled = true

		if r.process != nil {
			r.process.Kill()
		} else {
			logger.Error("No process to kill")
		}
	}

	return nil
}

// Creates the environment variables that will be used in the process
func (r *JobRunner) createEnvironment() []string {
	// Create a clone of our jobs environment. We'll then set the
	// environment variables provided by the agent, which will override any
	// sent by Buildkite. The variables below should always take
	// precedence.
	env := make(map[string]string)
	for key, value := range r.Job.Env {
		env[key] = value
	}

	// Add agent environment variables
	env["BUILDKITE_AGENT_ENDPOINT"] = r.Endpoint
	env["BUILDKITE_AGENT_ACCESS_TOKEN"] = r.Agent.AccessToken
	env["BUILDKITE_AGENT_DEBUG"] = fmt.Sprintf("%t", logger.GetLevel() == logger.DEBUG)

	// We know the BUILDKITE_BIN_PATH dir, because it's the path to the
	// currently running file (there is only 1 binary)
	dir, _ := filepath.Abs(filepath.Dir(os.Args[0]))
	env["BUILDKITE_BIN_PATH"] = dir

	// Add misc options
	env["BUILDKITE_BUILD_PATH"] = r.AgentConfiguration.BuildPath
	env["BUILDKITE_HOOKS_PATH"] = r.AgentConfiguration.HooksPath
	env["BUILDKITE_AUTO_SSH_FINGERPRINT_VERIFICATION"] = fmt.Sprintf("%t", r.AgentConfiguration.AutoSSHFingerprintVerification)
	env["BUILDKITE_COMMAND_EVAL"] = fmt.Sprintf("%t", r.AgentConfiguration.CommandEval)

	// Convert the env map into a slice (which is what the script gear
	// needs)
	envSlice := []string{}
	for key, value := range env {
		envSlice = append(envSlice, fmt.Sprintf("%s=%s", key, value))
	}

	return envSlice
}

// Starts the job in the Buildkite Agent API
func (r *JobRunner) startJob(startedAt time.Time) error {
	r.Job.StartedAt = startedAt.UTC().Format(time.RFC3339Nano)
	_, _, err := r.APIClient.Jobs.Start(r.Job)

	return err
}

// Finishes the job in the Buildkite Agent API. This call will keep on retrying
// forever until it finally gets a successfull response from the API.
func (r *JobRunner) finishJob(finishedAt time.Time, exitStatus string, failedChunkCount int) error {
	r.Job.FinishedAt = finishedAt.UTC().Format(time.RFC3339Nano)
	r.Job.ExitStatus = exitStatus
	r.Job.ChunksFailedCount = failedChunkCount

	return retry.Do(func(s *retry.Stats) error {
		_, _, err := r.APIClient.Jobs.Finish(r.Job)
		if err != nil {
			logger.Warn("%s (%s)", err, s)
		}

		return err
	}, &retry.Config{Forever: true, Interval: 1 * time.Second})
}

func (r *JobRunner) onProcessStartCallback() {
	// Start a routine that will grab the output every few seconds and send
	// it back to Buildkite
	go func() {
		for r.process.Running {
			// Send the output of the process to the log streamer
			// for processing
			r.logStreamer.Process(r.process.Output())

			// Check the output in another second
			time.Sleep(1 * time.Second)
		}

		logger.Debug("Routine that processes the log has finished")
	}()

	// Start a routine that will grab the output every few seconds and send it back to Buildkite
	go func() {
		for r.process.Running {
			// Re-get the job and check it's status to see if it's been
			// cancelled
			job, _, err := r.APIClient.Jobs.Get(r.Job.ID)
			if err != nil {
				// We don't really care if it fails, we'll just
				// try again in a second anyway
				logger.Warn("Problem with getting job status %s (%s)", r.Job.ID, err)
			} else if job.State == "canceled" {
				r.Kill()
			}

			// Check for cancellations every few seconds
			time.Sleep(3 * time.Second)
		}

		logger.Debug("Routine that refreshes the job has finished")
	}()
}

func (r *JobRunner) onUploadHeaderTime(cursor int, total int, times map[string]string) {
	retry.Do(func(s *retry.Stats) error {
		logger.Debug("Uploading header times %d..%d (%d)", cursor+1, total, len(times))

		_, err := r.APIClient.HeaderTimes.Save(r.Job.ID, &api.HeaderTimes{Times: times})
		if err != nil {
			logger.Warn("%s (%s)", err, s)
		}

		return err
	}, &retry.Config{Maximum: 10, Interval: 1 * time.Second})
}

// Call when a chunk is ready for upload. It retry the chunk upload with an
// interval before giving up.
func (r *JobRunner) onUploadChunk(chunk *LogStreamerChunk) error {
	return retry.Do(func(s *retry.Stats) error {
		_, err := r.APIClient.Chunks.Upload(r.Job.ID, &api.Chunk{
			Data:     chunk.Data,
			Sequence: chunk.Order,
		})
		if err != nil {
			logger.Warn("%s (%s)", err, s)
		}

		return err
	}, &retry.Config{Maximum: 10, Interval: 1 * time.Second})
}
<commit_msg>Gracefully handle 422 errors when finishing a job.<commit_after>package agent

import (
	"fmt"
	"github.com/buildkite/agent/api"
	"github.com/buildkite/agent/logger"
	"github.com/buildkite/agent/process"
	"github.com/buildkite/agent/retry"
	"os"
	"path/filepath"
	_ "regexp"
	"time"
)

type JobRunner struct {
	// The job being run
	Job *api.Job

	// The APIClient that will be used when updating the job
	APIClient *api.Client

	// The endpoint that should be used when communicating with the API
	Endpoint string

	// The registred agent API record running this job
	Agent *api.Agent

	// The configuration of the agent from the CLI
	AgentConfiguration *AgentConfiguration

	// The interal process of the job
	process *process.Process

	// The internal header time streamer
	headerTimesStreamer *HeaderTimesStreamer

	// The internal log streamer
	logStreamer *LogStreamer

	// If the job is being cancelled
	cancelled bool
}

// Initializes the job runner
func (r JobRunner) Create() (runner *JobRunner, err error) {
	runner = &r

	// Our own APIClient using the endpoint and the agents access token
	runner.APIClient = APIClient{Endpoint: r.Endpoint, Token: r.Agent.AccessToken}.Create()

	// // Create our header times struct
	runner.headerTimesStreamer = &HeaderTimesStreamer{Callback: r.onUploadHeaderTime}

	// The log streamer that will take the output chunks, and send them to
	// the Buildkite Agent API
	runner.logStreamer = LogStreamer{MaxChunkSizeBytes: r.Job.ChunksMaxSizeBytes, Callback: r.onUploadChunk}.New()

	// The process that will run the bootstrap script
	runner.process = process.Process{
		Script:        r.AgentConfiguration.BootstrapScript,
		Env:           r.createEnvironment(),
		PTY:           r.AgentConfiguration.RunInPty,
		StartCallback: r.onProcessStartCallback,
		LineCallback:  runner.headerTimesStreamer.Scan,
	}.Create()

	return
}

// Runs the job
func (r *JobRunner) Run() error {
	logger.Info("Starting job %s", r.Job.ID)

	// Start the log streamer
	if err := r.logStreamer.Start(); err != nil {
		return err
	}

	// Start the build in the Buildkite Agent API
	if err := r.startJob(time.Now()); err != nil {
		return err
	}

	// Start the process. This will block until it finishes.
	if err := r.process.Start(); err != nil {
		// Send the error as output
		r.logStreamer.Process(fmt.Sprintf("%s", err))
	} else {
		// Add the final output to the streamer
		r.logStreamer.Process(r.process.Output())
	}

	// Store the finished at time
	finishedAt := time.Now()

	// Wait until all the header times have finished uploading
	logger.Debug("[HeaderTimes] Waiting for all times to finish uploading")

	r.headerTimesStreamer.Wait()

	// Stop the log streamer. This will block until all the chunks have
	// been uploaded
	r.logStreamer.Stop()

	// Warn about failed chunks
	if r.logStreamer.ChunksFailedCount > 0 {
		logger.Warn("%d chunks failed to upload for this job", r.logStreamer.ChunksFailedCount)
	}

	// Finish the build in the Buildkite Agent API
	r.finishJob(finishedAt, r.process.ExitStatus, int(r.logStreamer.ChunksFailedCount))

	logger.Info("Finished job %s", r.Job.ID)

	return nil
}

func (r *JobRunner) Kill() error {
	if !r.cancelled {
		logger.Info("Canceling job %s", r.Job.ID)
		r.cancelled = true

		if r.process != nil {
			r.process.Kill()
		} else {
			logger.Error("No process to kill")
		}
	}

	return nil
}

// Creates the environment variables that will be used in the process
func (r *JobRunner) createEnvironment() []string {
	// Create a clone of our jobs environment. We'll then set the
	// environment variables provided by the agent, which will override any
	// sent by Buildkite. The variables below should always take
	// precedence.
	env := make(map[string]string)
	for key, value := range r.Job.Env {
		env[key] = value
	}

	// Add agent environment variables
	env["BUILDKITE_AGENT_ENDPOINT"] = r.Endpoint
	env["BUILDKITE_AGENT_ACCESS_TOKEN"] = r.Agent.AccessToken
	env["BUILDKITE_AGENT_DEBUG"] = fmt.Sprintf("%t", logger.GetLevel() == logger.DEBUG)

	// We know the BUILDKITE_BIN_PATH dir, because it's the path to the
	// currently running file (there is only 1 binary)
	dir, _ := filepath.Abs(filepath.Dir(os.Args[0]))
	env["BUILDKITE_BIN_PATH"] = dir

	// Add misc options
	env["BUILDKITE_BUILD_PATH"] = r.AgentConfiguration.BuildPath
	env["BUILDKITE_HOOKS_PATH"] = r.AgentConfiguration.HooksPath
	env["BUILDKITE_AUTO_SSH_FINGERPRINT_VERIFICATION"] = fmt.Sprintf("%t", r.AgentConfiguration.AutoSSHFingerprintVerification)
	env["BUILDKITE_COMMAND_EVAL"] = fmt.Sprintf("%t", r.AgentConfiguration.CommandEval)

	// Convert the env map into a slice (which is what the script gear
	// needs)
	envSlice := []string{}
	for key, value := range env {
		envSlice = append(envSlice, fmt.Sprintf("%s=%s", key, value))
	}

	return envSlice
}

// Starts the job in the Buildkite Agent API
func (r *JobRunner) startJob(startedAt time.Time) error {
	r.Job.StartedAt = startedAt.UTC().Format(time.RFC3339Nano)
	_, _, err := r.APIClient.Jobs.Start(r.Job)

	return err
}

// Finishes the job in the Buildkite Agent API. This call will keep on retrying
// forever until it finally gets a successfull response from the API.
func (r *JobRunner) finishJob(finishedAt time.Time, exitStatus string, failedChunkCount int) error {
	r.Job.FinishedAt = finishedAt.UTC().Format(time.RFC3339Nano)
	r.Job.ExitStatus = exitStatus
	r.Job.ChunksFailedCount = failedChunkCount

	return retry.Do(func(s *retry.Stats) error {
		_, response, err := r.APIClient.Jobs.Finish(r.Job)
		if err != nil {
			// If the API returns with a 422, that means that we
			// succesfully tried to finish the job, but Buildkite
			// rejected the finish for some reason. This can
			// sometimes mean that Buildkite has cancelled the job
			// before we get a chance to send the final API call
			// (maybe this agent took too long to kill the
			// process). In that case, we don't want to keep trying
			// to finish the job forever so we'll just bail out and
			// go find some more work to do.
			if response.StatusCode == 422 {
				logger.Warn("Buildkite rejected the call to finish the job (%s)", err)
				return nil
			} else {
				logger.Warn("%s (%s)", err, s)
			}
		}

		return err
	}, &retry.Config{Forever: true, Interval: 1 * time.Second})
}

func (r *JobRunner) onProcessStartCallback() {
	// Start a routine that will grab the output every few seconds and send
	// it back to Buildkite
	go func() {
		for r.process.Running {
			// Send the output of the process to the log streamer
			// for processing
			r.logStreamer.Process(r.process.Output())

			// Check the output in another second
			time.Sleep(1 * time.Second)
		}

		logger.Debug("Routine that processes the log has finished")
	}()

	// Start a routine that will grab the output every few seconds and send it back to Buildkite
	go func() {
		for r.process.Running {
			// Re-get the job and check it's status to see if it's been
			// cancelled
			job, _, err := r.APIClient.Jobs.Get(r.Job.ID)
			if err != nil {
				// We don't really care if it fails, we'll just
				// try again in a second anyway
				logger.Warn("Problem with getting job status %s (%s)", r.Job.ID, err)
			} else if job.State == "canceled" {
				r.Kill()
			}

			// Check for cancellations every few seconds
			time.Sleep(3 * time.Second)
		}

		logger.Debug("Routine that refreshes the job has finished")
	}()
}

func (r *JobRunner) onUploadHeaderTime(cursor int, total int, times map[string]string) {
	retry.Do(func(s *retry.Stats) error {
		logger.Debug("Uploading header times %d..%d (%d)", cursor+1, total, len(times))

		_, err := r.APIClient.HeaderTimes.Save(r.Job.ID, &api.HeaderTimes{Times: times})
		if err != nil {
			logger.Warn("%s (%s)", err, s)
		}

		return err
	}, &retry.Config{Maximum: 10, Interval: 1 * time.Second})
}

// Call when a chunk is ready for upload. It retry the chunk upload with an
// interval before giving up.
func (r *JobRunner) onUploadChunk(chunk *LogStreamerChunk) error {
	return retry.Do(func(s *retry.Stats) error {
		_, err := r.APIClient.Chunks.Upload(r.Job.ID, &api.Chunk{
			Data:     chunk.Data,
			Sequence: chunk.Order,
		})
		if err != nil {
			logger.Warn("%s (%s)", err, s)
		}

		return err
	}, &retry.Config{Maximum: 10, Interval: 1 * time.Second})
}
<|endoftext|>


<commit_before>// Copyright 2017 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 gc

import (
	"bufio"
	"internal/testenv"
	"io"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"testing"
)

// TestIntendedInlining tests that specific runtime functions are inlined.
// This allows refactoring for code clarity and re-use without fear that
// changes to the compiler will cause silent performance regressions.
func TestIntendedInlining(t *testing.T) {
	if testing.Short() && testenv.Builder() == "" {
		t.Skip("skipping in short mode")
	}
	testenv.MustHaveGoRun(t)
	t.Parallel()

	// want is the list of function names (by package) that should
	// be inlined.
	want := map[string][]string{
		"runtime": {
			// TODO(mvdan): enable these once mid-stack
			// inlining is available
			// "adjustctxt",

			"add",
			"acquirem",
			"add1",
			"addb",
			"adjustpanics",
			"adjustpointer",
			"bucketMask",
			"bucketShift",
			"chanbuf",
			"deferArgs",
			"deferclass",
			"evacuated",
			"fastlog2",
			"fastrand",
			"float64bits",
			"funcPC",
			"getm",
			"isDirectIface",
			"itabHashFunc",
			"maxSliceCap",
			"noescape",
			"readUnaligned32",
			"readUnaligned64",
			"releasem",
			"round",
			"roundupsize",
			"selectsize",
			"stringStructOf",
			"subtract1",
			"subtractb",
			"tophash",
			"totaldefersize",
			"(*bmap).keys",
			"(*bmap).overflow",
			"(*waitq).enqueue",

			// GC-related ones
			"cgoInRange",
			"gclinkptr.ptr",
			"guintptr.ptr",
			"heapBits.bits",
			"heapBits.isPointer",
			"heapBits.morePointers",
			"heapBits.next",
			"heapBitsForAddr",
			"inheap",
			"markBits.isMarked",
			"muintptr.ptr",
			"puintptr.ptr",
			"spanOfUnchecked",
			"(*gcWork).putFast",
			"(*gcWork).tryGetFast",
			"(*guintptr).set",
			"(*markBits).advance",
			"(*mspan).allocBitsForIndex",
			"(*mspan).base",
			"(*mspan).markBitsForBase",
			"(*mspan).markBitsForIndex",
			"(*muintptr).set",
			"(*puintptr).set",
		},
		"runtime/internal/sys": {},
		"bytes": {
			"(*Buffer).Bytes",
			"(*Buffer).Cap",
			"(*Buffer).Len",
			"(*Buffer).Next",
			"(*Buffer).Read",
			"(*Buffer).ReadByte",
			"(*Buffer).Reset",
			"(*Buffer).String",
			"(*Buffer).UnreadByte",
			"(*Buffer).tryGrowByReslice",
		},
		"unicode/utf8": {
			"FullRune",
			"FullRuneInString",
			"RuneLen",
			"ValidRune",
		},
		"reflect": {
			"Value.CanAddr",
			"Value.CanSet",
			"Value.IsValid",
			"add",
			"align",
			"flag.kind",
			"flag.ro",

			// TODO: these use panic, need mid-stack
			// inlining
			// "Value.CanInterface",
			// "Value.pointer",
			// "flag.mustBe",
			// "flag.mustBeAssignable",
			// "flag.mustBeExported",
		},
		"regexp": {
			"(*bitState).push",
		},
	}

	if runtime.GOARCH != "386" {
		// nextFreeFast calls sys.Ctz64, which on 386 is implemented in asm and is not inlinable.
		// We currently don't have midstack inlining so nextFreeFast is also not inlinable on 386.
		// So check for it only on non-386 platforms.
		want["runtime"] = append(want["runtime"], "nextFreeFast")
		// As explained above, Ctz64 and Ctz32 are not Go code on 386.
		// The same applies to Bswap32.
		want["runtime/internal/sys"] = append(want["runtime/internal/sys"], "Ctz64")
		want["runtime/internal/sys"] = append(want["runtime/internal/sys"], "Ctz32")
		want["runtime/internal/sys"] = append(want["runtime/internal/sys"], "Bswap32")
	}
	switch runtime.GOARCH {
	case "amd64", "amd64p32", "arm64", "mips64", "mips64le", "ppc64", "ppc64le", "s390x":
		// rotl_31 is only defined on 64-bit architectures
		want["runtime"] = append(want["runtime"], "rotl_31")
	}

	notInlinedReason := make(map[string]string)
	pkgs := make([]string, 0, len(want))
	for pname, fnames := range want {
		pkgs = append(pkgs, pname)
		for _, fname := range fnames {
			fullName := pname + "." + fname
			if _, ok := notInlinedReason[fullName]; ok {
				t.Errorf("duplicate func: %s", fullName)
			}
			notInlinedReason[fullName] = "unknown reason"
		}
	}

	args := append([]string{"build", "-a", "-gcflags=-m -m"}, pkgs...)
	cmd := testenv.CleanCmdEnv(exec.Command(testenv.GoToolPath(t), args...))
	pr, pw := io.Pipe()
	cmd.Stdout = pw
	cmd.Stderr = pw
	cmdErr := make(chan error, 1)
	go func() {
		cmdErr <- cmd.Run()
		pw.Close()
	}()
	scanner := bufio.NewScanner(pr)
	curPkg := ""
	canInline := regexp.MustCompile(`: can inline ([^ ]*)`)
	cannotInline := regexp.MustCompile(`: cannot inline ([^ ]*): (.*)`)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "# ") {
			curPkg = line[2:]
			continue
		}
		if m := canInline.FindStringSubmatch(line); m != nil {
			fname := m[1]
			delete(notInlinedReason, curPkg+"."+fname)
			continue
		}
		if m := cannotInline.FindStringSubmatch(line); m != nil {
			fname, reason := m[1], m[2]
			fullName := curPkg + "." + fname
			if _, ok := notInlinedReason[fullName]; ok {
				// cmd/compile gave us a reason why
				notInlinedReason[fullName] = reason
			}
			continue
		}
	}
	if err := <-cmdErr; err != nil {
		t.Fatal(err)
	}
	if err := scanner.Err(); err != nil {
		t.Fatal(err)
	}
	for fullName, reason := range notInlinedReason {
		t.Errorf("%s was not inlined: %s", fullName, reason)
	}
}
<commit_msg>cmd/compile: skip runtime.nextFreeFast inlining test on MIPS64x<commit_after>// Copyright 2017 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 gc

import (
	"bufio"
	"internal/testenv"
	"io"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"testing"
)

// TestIntendedInlining tests that specific runtime functions are inlined.
// This allows refactoring for code clarity and re-use without fear that
// changes to the compiler will cause silent performance regressions.
func TestIntendedInlining(t *testing.T) {
	if testing.Short() && testenv.Builder() == "" {
		t.Skip("skipping in short mode")
	}
	testenv.MustHaveGoRun(t)
	t.Parallel()

	// want is the list of function names (by package) that should
	// be inlined.
	want := map[string][]string{
		"runtime": {
			// TODO(mvdan): enable these once mid-stack
			// inlining is available
			// "adjustctxt",

			"add",
			"acquirem",
			"add1",
			"addb",
			"adjustpanics",
			"adjustpointer",
			"bucketMask",
			"bucketShift",
			"chanbuf",
			"deferArgs",
			"deferclass",
			"evacuated",
			"fastlog2",
			"fastrand",
			"float64bits",
			"funcPC",
			"getm",
			"isDirectIface",
			"itabHashFunc",
			"maxSliceCap",
			"noescape",
			"readUnaligned32",
			"readUnaligned64",
			"releasem",
			"round",
			"roundupsize",
			"selectsize",
			"stringStructOf",
			"subtract1",
			"subtractb",
			"tophash",
			"totaldefersize",
			"(*bmap).keys",
			"(*bmap).overflow",
			"(*waitq).enqueue",

			// GC-related ones
			"cgoInRange",
			"gclinkptr.ptr",
			"guintptr.ptr",
			"heapBits.bits",
			"heapBits.isPointer",
			"heapBits.morePointers",
			"heapBits.next",
			"heapBitsForAddr",
			"inheap",
			"markBits.isMarked",
			"muintptr.ptr",
			"puintptr.ptr",
			"spanOfUnchecked",
			"(*gcWork).putFast",
			"(*gcWork).tryGetFast",
			"(*guintptr).set",
			"(*markBits).advance",
			"(*mspan).allocBitsForIndex",
			"(*mspan).base",
			"(*mspan).markBitsForBase",
			"(*mspan).markBitsForIndex",
			"(*muintptr).set",
			"(*puintptr).set",
		},
		"runtime/internal/sys": {},
		"bytes": {
			"(*Buffer).Bytes",
			"(*Buffer).Cap",
			"(*Buffer).Len",
			"(*Buffer).Next",
			"(*Buffer).Read",
			"(*Buffer).ReadByte",
			"(*Buffer).Reset",
			"(*Buffer).String",
			"(*Buffer).UnreadByte",
			"(*Buffer).tryGrowByReslice",
		},
		"unicode/utf8": {
			"FullRune",
			"FullRuneInString",
			"RuneLen",
			"ValidRune",
		},
		"reflect": {
			"Value.CanAddr",
			"Value.CanSet",
			"Value.IsValid",
			"add",
			"align",
			"flag.kind",
			"flag.ro",

			// TODO: these use panic, need mid-stack
			// inlining
			// "Value.CanInterface",
			// "Value.pointer",
			// "flag.mustBe",
			// "flag.mustBeAssignable",
			// "flag.mustBeExported",
		},
		"regexp": {
			"(*bitState).push",
		},
	}

	if runtime.GOARCH != "386" && runtime.GOARCH != "mips64" && runtime.GOARCH != "mips64le" {
		// nextFreeFast calls sys.Ctz64, which on 386 is implemented in asm and is not inlinable.
		// We currently don't have midstack inlining so nextFreeFast is also not inlinable on 386.
		// On MIPS64x, Ctz64 is not intrinsified and causes nextFreeFast too expensive to inline
		// (Issue 22239).
		want["runtime"] = append(want["runtime"], "nextFreeFast")
	}
	if runtime.GOARCH != "386" {
		// As explained above, Ctz64 and Ctz32 are not Go code on 386.
		// The same applies to Bswap32.
		want["runtime/internal/sys"] = append(want["runtime/internal/sys"], "Ctz64")
		want["runtime/internal/sys"] = append(want["runtime/internal/sys"], "Ctz32")
		want["runtime/internal/sys"] = append(want["runtime/internal/sys"], "Bswap32")
	}
	switch runtime.GOARCH {
	case "amd64", "amd64p32", "arm64", "mips64", "mips64le", "ppc64", "ppc64le", "s390x":
		// rotl_31 is only defined on 64-bit architectures
		want["runtime"] = append(want["runtime"], "rotl_31")
	}

	notInlinedReason := make(map[string]string)
	pkgs := make([]string, 0, len(want))
	for pname, fnames := range want {
		pkgs = append(pkgs, pname)
		for _, fname := range fnames {
			fullName := pname + "." + fname
			if _, ok := notInlinedReason[fullName]; ok {
				t.Errorf("duplicate func: %s", fullName)
			}
			notInlinedReason[fullName] = "unknown reason"
		}
	}

	args := append([]string{"build", "-a", "-gcflags=-m -m"}, pkgs...)
	cmd := testenv.CleanCmdEnv(exec.Command(testenv.GoToolPath(t), args...))
	pr, pw := io.Pipe()
	cmd.Stdout = pw
	cmd.Stderr = pw
	cmdErr := make(chan error, 1)
	go func() {
		cmdErr <- cmd.Run()
		pw.Close()
	}()
	scanner := bufio.NewScanner(pr)
	curPkg := ""
	canInline := regexp.MustCompile(`: can inline ([^ ]*)`)
	cannotInline := regexp.MustCompile(`: cannot inline ([^ ]*): (.*)`)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "# ") {
			curPkg = line[2:]
			continue
		}
		if m := canInline.FindStringSubmatch(line); m != nil {
			fname := m[1]
			delete(notInlinedReason, curPkg+"."+fname)
			continue
		}
		if m := cannotInline.FindStringSubmatch(line); m != nil {
			fname, reason := m[1], m[2]
			fullName := curPkg + "." + fname
			if _, ok := notInlinedReason[fullName]; ok {
				// cmd/compile gave us a reason why
				notInlinedReason[fullName] = reason
			}
			continue
		}
	}
	if err := <-cmdErr; err != nil {
		t.Fatal(err)
	}
	if err := scanner.Err(); err != nil {
		t.Fatal(err)
	}
	for fullName, reason := range notInlinedReason {
		t.Errorf("%s was not inlined: %s", fullName, reason)
	}
}
<|endoftext|>


<commit_before>package vegeta

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

func TestAttackRate(t *testing.T) {
	t.Parallel()

	hitCount := uint64(0)
	server := httptest.NewServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			atomic.AddUint64(&hitCount, 1)
		}),
	)

	tgt := Target{Method: "GET", URL: server.URL}
	rate := uint64(5000)
	Attack(Targets{tgt}, rate, 1*time.Second)
	if hits := atomic.LoadUint64(&hitCount); hits != rate {
		t.Fatalf("Wrong number of hits: want %d, got %d\n", rate, hits)
	}
}

func TestAttackBody(t *testing.T) {
	t.Parallel()

	want := []byte("VEGETA!")
	server := httptest.NewServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			got, err := ioutil.ReadAll(r.Body)
			if err != nil {
				t.Fatal(err)
			}
			if !bytes.Equal(want, got) {
				t.Fatalf("Wrong body. Want: %s, Got: %s", want, got)
			}
		}),
	)

	tgt := Target{Method: "GET", URL: server.URL, Body: want}
	rate := uint64(5000)
	Attack(Targets{tgt}, rate, 1*time.Second)
}

func TestDefaultAttackerCertConfig(t *testing.T) {
	t.Parallel()

	server := httptest.NewTLSServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}),
	)
	request, _ := http.NewRequest("GET", server.URL, nil)
	_, err := DefaultAttacker.client.Do(request)
	if err != nil && strings.Contains(err.Error(), "x509: certificate signed by unknown authority") {
		t.Errorf("Invalid certificates should be ignored: Got `%s`", err)
	}
}

func TestRedirects(t *testing.T) {
	t.Parallel()

	var servers [2]*httptest.Server
	var hits uint64

	for i := range servers {
		servers[i] = httptest.NewServer(
			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				atomic.AddUint64(&hits, 1)
				http.Redirect(w, r, servers[(i+1)%2].URL, 302)
			}),
		)
	}

	atk := NewAttacker(2, DefaultTimeout, DefaultLocalAddr)
	tgt := Target{Method: "GET", URL: servers[0].URL}
	var rate uint64 = 100
	results := atk.Attack(Targets{tgt}, rate, 1*time.Second)

	want := fmt.Sprintf("stopped after %d redirects", 2)
	for _, result := range results {
		if !strings.Contains(result.Error, want) {
			t.Fatalf("Expected error to be: %s, Got: %s", want, result.Error)
		}
	}

	if want, got := rate*(2+1), hits; want != got {
		t.Fatalf("Expected hits to be: %d, Got: %d", want, got)
	}
}

func TestTimeout(t *testing.T) {
	t.Parallel()

	server := httptest.NewServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			<-time.After(2 * time.Second)
		}),
	)

	atk := NewAttacker(DefaultRedirects, 500*time.Millisecond, DefaultLocalAddr)
	tgt := Target{Method: "GET", URL: server.URL}
	results := atk.Attack(Targets{tgt}, 100, 1*time.Second)

	want := "net/http: timeout awaiting response headers"
	for _, result := range results {
		if !strings.Contains(result.Error, want) {
			t.Fatalf("Expected error to be: %s, Got: %s", want, result.Error)
		}
	}
}

func TestLocalAddr(t *testing.T) {
	t.Parallel()

	addr, err := net.ResolveIPAddr("ip", "127.0.0.1")
	if err != nil {
		t.Fatal(err)
	}

	server := httptest.NewServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			host, _, err := net.SplitHostPort(r.RemoteAddr)
			if err != nil {
				t.Fatal(err)
			}

			if host != addr.String() {
				t.Fatalf("Wrong source address. Want %s, Got %s", addr, host)
			}
		}),
	)

	atk := NewAttacker(DefaultRedirects, DefaultTimeout, *addr)
	tgt := Target{Method: "GET", URL: server.URL}
	for _, result := range atk.Attack(Targets{tgt}, 100, 1*time.Second) {
		if result.Error != "" {
			t.Fatal(result.Error)
		}
	}
}
<commit_msg>Speed up tests in lib/attack_test.go<commit_after>package vegeta

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

func TestAttackRate(t *testing.T) {
	t.Parallel()

	hitCount := uint64(0)
	server := httptest.NewServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			atomic.AddUint64(&hitCount, 1)
		}),
	)

	tgt := Target{Method: "GET", URL: server.URL}
	rate := uint64(1000)
	Attack(Targets{tgt}, rate, 1*time.Second)
	if hits := atomic.LoadUint64(&hitCount); hits != rate {
		t.Fatalf("Wrong number of hits: want %d, got %d\n", rate, hits)
	}
}

func TestAttackBody(t *testing.T) {
	t.Parallel()

	want := []byte("VEGETA!")
	server := httptest.NewServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			got, err := ioutil.ReadAll(r.Body)
			if err != nil {
				t.Fatal(err)
			}
			if !bytes.Equal(want, got) {
				t.Fatalf("Wrong body. Want: %s, Got: %s", want, got)
			}
		}),
	)

	Attack(Targets{{Method: "GET", URL: server.URL, Body: want}}, 100, 1*time.Second)
}

func TestDefaultAttackerCertConfig(t *testing.T) {
	t.Parallel()

	server := httptest.NewTLSServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}),
	)
	request, _ := http.NewRequest("GET", server.URL, nil)
	_, err := DefaultAttacker.client.Do(request)
	if err != nil && strings.Contains(err.Error(), "x509: certificate signed by unknown authority") {
		t.Errorf("Invalid certificates should be ignored: Got `%s`", err)
	}
}

func TestRedirects(t *testing.T) {
	t.Parallel()

	var servers [2]*httptest.Server
	var hits uint64

	for i := range servers {
		servers[i] = httptest.NewServer(
			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				atomic.AddUint64(&hits, 1)
				http.Redirect(w, r, servers[(i+1)%2].URL, 302)
			}),
		)
	}

	atk := NewAttacker(2, DefaultTimeout, DefaultLocalAddr)
	tgt := Target{Method: "GET", URL: servers[0].URL}
	var rate uint64 = 10
	results := atk.Attack(Targets{tgt}, rate, 1*time.Second)

	want := fmt.Sprintf("stopped after %d redirects", 2)
	for _, result := range results {
		if !strings.Contains(result.Error, want) {
			t.Fatalf("Expected error to be: %s, Got: %s", want, result.Error)
		}
	}

	if want, got := rate*(2+1), hits; want != got {
		t.Fatalf("Expected hits to be: %d, Got: %d", want, got)
	}
}

func TestTimeout(t *testing.T) {
	t.Parallel()

	server := httptest.NewServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			<-time.After(20 * time.Millisecond)
		}),
	)

	atk := NewAttacker(DefaultRedirects, 10*time.Millisecond, DefaultLocalAddr)
	tgt := Target{Method: "GET", URL: server.URL}
	results := atk.Attack(Targets{tgt}, 1, 1*time.Second)

	want := "net/http: timeout awaiting response headers"
	for _, result := range results {
		if !strings.Contains(result.Error, want) {
			t.Fatalf("Expected error to be: %s, Got: %s", want, result.Error)
		}
	}
}

func TestLocalAddr(t *testing.T) {
	t.Parallel()

	addr, err := net.ResolveIPAddr("ip", "127.0.0.1")
	if err != nil {
		t.Fatal(err)
	}

	server := httptest.NewServer(
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			host, _, err := net.SplitHostPort(r.RemoteAddr)
			if err != nil {
				t.Fatal(err)
			}

			if host != addr.String() {
				t.Fatalf("Wrong source address. Want %s, Got %s", addr, host)
			}
		}),
	)

	atk := NewAttacker(DefaultRedirects, DefaultTimeout, *addr)
	tgt := Target{Method: "GET", URL: server.URL}

	for _, result := range atk.Attack(Targets{tgt}, 1, 1*time.Second) {
		if result.Error != "" {
			t.Fatal(result.Error)
		}
	}
}
<|endoftext|>


<commit_before>package gokeepasslib

import (
	"bytes"
	"compress/gzip"
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"reflect"
)

type Decoder struct {
	r io.Reader
}

func (d *Decoder) Decode(db *Database) error {
	if err := d.readSignature(db); err != nil {
		return err
	}

	if err := d.readHeaders(db); err != nil {
		return err
	}

	if err := d.readData(db); err != nil {
		return err
	}

	return nil
}

func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{r: r}
}

func (d *Decoder) readSignature(db *Database) error {
	sig := new(Signature)
	if err := binary.Read(d.r, binary.LittleEndian, sig); err != nil {
		return err
	}

	if sig.BaseSignature != [...]byte{0x03, 0xd9, 0xa2, 0x9a} {
		return errors.New("BaseSignature not valid")
	}
	if sig.VersionSignature != [...]byte{0x67, 0xfb, 0x4b, 0xb5} {
		return errors.New("VersionSignature not valid")
	}

	db.signature = *sig
	return nil
}

func (d *Decoder) readHeaders(db *Database) error {
	headers := new(Headers)

	for {
		var fieldID byte
		if err := binary.Read(d.r, binary.LittleEndian, &fieldID); err != nil {
			return err
		}

		var fieldLength [2]byte
		if err := binary.Read(d.r, binary.LittleEndian, &fieldLength); err != nil {
			return err
		}

		var fieldData = make([]byte, binary.LittleEndian.Uint16(fieldLength[:]))
		if err := binary.Read(d.r, binary.LittleEndian, &fieldData); err != nil {
			return err
		}

		switch fieldID {
		case 1:
			headers.Comment = fieldData
		case 2:
			headers.CipherID = fieldData
		case 3:
			headers.CompressionFlags = binary.LittleEndian.Uint32(fieldData)
		case 4:
			headers.MasterSeed = fieldData
		case 5:
			headers.TransformSeed = fieldData
		case 6:
			data := binary.LittleEndian.Uint32(fieldData[4:8])*
				(1<<16)*(1<<16) +
				binary.LittleEndian.Uint32(fieldData[0:4])
			headers.TransformRounds = data
		case 7:
			headers.EncryptionIV = fieldData
		case 8:
			headers.ProtectedStreamKey = fieldData
		case 9:
			headers.StreamStartBytes = fieldData
		case 10:
			headers.InnerRandomStreamID = fieldData
		}

		if fieldID == 0 {
			break
		}
	}

	db.headers = *headers
	return nil
}

func (d *Decoder) readData(db *Database) error {
	masterKey, nil := d.buildMasterKey(db)

	block, err := aes.NewCipher(masterKey)
	if err != nil {
		return err
	}

	in, err := ioutil.ReadAll(d.r)
	if err != nil {
		return err
	}

	mode := cipher.NewCBCDecrypter(block, db.headers.EncryptionIV)
	decrypted := make([]byte, len(in))
	mode.CryptBlocks(decrypted, in)

	startBytes := db.headers.StreamStartBytes
	if !reflect.DeepEqual(decrypted[0:len(startBytes)], startBytes) {
		return errors.New("Database integrity check failed")
	}
	decrypted = decrypted[len(startBytes):]

	zippedBody, err := d.checkHashBlocks(decrypted)
	if err != nil {
		return err
	}

	b := bytes.NewBuffer(zippedBody)
	r, err := gzip.NewReader(b)
	if err != nil {
		return err
	}
	defer r.Close()
	result := []byte{}
	r.Read(result)

	fmt.Println(result)

	return nil
}

func (d *Decoder) buildMasterKey(db *Database) ([]byte, error) {
	masterKey := make([]byte, 32)
	copy(masterKey, db.credentials.Key)

	tmp := sha256.Sum256(masterKey)
	masterKey = tmp[:]

	block, err := aes.NewCipher(db.headers.TransformSeed)
	if err != nil {
		return nil, err
	}

	// http://crypto.stackexchange.com/questions/21048/can-i-simulate-iterated-aes-ecb-with-other-block-cipher-modes
	for i := uint32(0); i < 6000; i++ {
		result := make([]byte, 16)
		crypter := cipher.NewCBCEncrypter(block, result)
		crypter.CryptBlocks(masterKey[:16], masterKey[:16])
		crypter = cipher.NewCBCEncrypter(block, result)
		crypter.CryptBlocks(masterKey[16:], masterKey[16:])
	}

	tmp = sha256.Sum256(masterKey)
	masterKey = tmp[:]

	masterKey = append(db.headers.MasterSeed, masterKey...)
	masterHash := sha256.Sum256(masterKey)
	masterKey = masterHash[:]

	return masterKey, nil
}

func (d *Decoder) checkHashBlocks(hashedBody []byte) ([]byte, error) {
	result := make([]byte, 0)

	for len(hashedBody) > 0 {
		fmt.Println(len(hashedBody))
		index := binary.LittleEndian.Uint32(hashedBody[:4])
		hashedBody = hashedBody[4:]
		blockHash := hashedBody[:32]
		hashedBody = hashedBody[32:]
		blockLength := binary.LittleEndian.Uint32(hashedBody[:4])
		hashedBody = hashedBody[4:]

		if blockLength > 0 {
			blockData := hashedBody[:blockLength]
			hashedBody = hashedBody[blockLength:]
			calculatedHash := sha256.Sum256(blockData)

			if !reflect.DeepEqual(calculatedHash[:], blockHash[:]) {
				return nil, fmt.Errorf("Hash mismatch. Database seems to be corrupt at index %d", index)
			} else {
				result = append(result, blockData...)
			}
		}
	}

	return result, nil
}
<commit_msg>made hash check and unzip work<commit_after>package gokeepasslib

import (
	"bytes"
	"compress/gzip"
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"reflect"
)

type Decoder struct {
	r io.Reader
}

func (d *Decoder) Decode(db *Database) error {
	if err := d.readSignature(db); err != nil {
		return err
	}

	if err := d.readHeaders(db); err != nil {
		return err
	}

	if err := d.readData(db); err != nil {
		return err
	}

	return nil
}

func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{r: r}
}

func (d *Decoder) readSignature(db *Database) error {
	sig := new(Signature)
	if err := binary.Read(d.r, binary.LittleEndian, sig); err != nil {
		return err
	}

	if sig.BaseSignature != [...]byte{0x03, 0xd9, 0xa2, 0x9a} {
		return errors.New("BaseSignature not valid")
	}
	if sig.VersionSignature != [...]byte{0x67, 0xfb, 0x4b, 0xb5} {
		return errors.New("VersionSignature not valid")
	}

	db.signature = *sig
	return nil
}

func (d *Decoder) readHeaders(db *Database) error {
	headers := new(Headers)

	for {
		var fieldID byte
		if err := binary.Read(d.r, binary.LittleEndian, &fieldID); err != nil {
			return err
		}

		var fieldLength [2]byte
		if err := binary.Read(d.r, binary.LittleEndian, &fieldLength); err != nil {
			return err
		}

		var fieldData = make([]byte, binary.LittleEndian.Uint16(fieldLength[:]))
		if err := binary.Read(d.r, binary.LittleEndian, &fieldData); err != nil {
			return err
		}

		switch fieldID {
		case 1:
			headers.Comment = fieldData
		case 2:
			headers.CipherID = fieldData
		case 3:
			headers.CompressionFlags = binary.LittleEndian.Uint32(fieldData)
		case 4:
			headers.MasterSeed = fieldData
		case 5:
			headers.TransformSeed = fieldData
		case 6:
			data := binary.LittleEndian.Uint32(fieldData[4:8])*
				(1<<16)*(1<<16) +
				binary.LittleEndian.Uint32(fieldData[0:4])
			headers.TransformRounds = data
		case 7:
			headers.EncryptionIV = fieldData
		case 8:
			headers.ProtectedStreamKey = fieldData
		case 9:
			headers.StreamStartBytes = fieldData
		case 10:
			headers.InnerRandomStreamID = fieldData
		}

		if fieldID == 0 {
			break
		}
	}

	db.headers = *headers
	return nil
}

func (d *Decoder) readData(db *Database) error {
	masterKey, nil := d.buildMasterKey(db)

	block, err := aes.NewCipher(masterKey)
	if err != nil {
		return err
	}

	in, err := ioutil.ReadAll(d.r)
	if err != nil {
		return err
	}

	mode := cipher.NewCBCDecrypter(block, db.headers.EncryptionIV)
	decrypted := make([]byte, len(in))
	mode.CryptBlocks(decrypted, in)

	startBytes := db.headers.StreamStartBytes
	if !reflect.DeepEqual(decrypted[0:len(startBytes)], startBytes) {
		return errors.New("Database integrity check failed")
	}
	decrypted = decrypted[len(startBytes):]

	zippedBody, err := d.checkHashBlocks(decrypted)
	if err != nil {
		return err
	}

	b := bytes.NewBuffer(zippedBody)
	r, err := gzip.NewReader(b)
	if err != nil {
		return err
	}
	defer r.Close()

	data, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}

	fmt.Printf("%s\n", data)

	return nil
}

func (d *Decoder) buildMasterKey(db *Database) ([]byte, error) {
	masterKey := make([]byte, 32)
	copy(masterKey, db.credentials.Key)

	tmp := sha256.Sum256(masterKey)
	masterKey = tmp[:]

	block, err := aes.NewCipher(db.headers.TransformSeed)
	if err != nil {
		return nil, err
	}

	// http://crypto.stackexchange.com/questions/21048/can-i-simulate-iterated-aes-ecb-with-other-block-cipher-modes
	for i := uint32(0); i < 6000; i++ {
		result := make([]byte, 16)
		crypter := cipher.NewCBCEncrypter(block, result)
		crypter.CryptBlocks(masterKey[:16], masterKey[:16])
		crypter = cipher.NewCBCEncrypter(block, result)
		crypter.CryptBlocks(masterKey[16:], masterKey[16:])
	}

	tmp = sha256.Sum256(masterKey)
	masterKey = tmp[:]

	masterKey = append(db.headers.MasterSeed, masterKey...)
	masterHash := sha256.Sum256(masterKey)
	masterKey = masterHash[:]

	return masterKey, nil
}

func (d *Decoder) checkHashBlocks(hashedBody []byte) ([]byte, error) {
	result := make([]byte, 0)

	for len(hashedBody) > 0 {
		index := binary.LittleEndian.Uint32(hashedBody[:4])
		hashedBody = hashedBody[4:]
		blockHash := hashedBody[:32]
		hashedBody = hashedBody[32:]
		blockLength := binary.LittleEndian.Uint32(hashedBody[:4])
		hashedBody = hashedBody[4:]

		if blockLength > 0 {
			blockData := hashedBody[:blockLength]
			hashedBody = hashedBody[blockLength:]
			calculatedHash := sha256.Sum256(blockData)

			if !reflect.DeepEqual(calculatedHash[:], blockHash[:]) {
				return nil, fmt.Errorf("Hash mismatch. Database seems to be corrupt at index %d", index)
			} else {
				result = append(result, blockData...)
			}
		} else {
			break
		}
	}

	return result, nil
}
<|endoftext|>


<commit_before>package weixinmp

import (
	"fmt"
	"io/ioutil"
	"os"
	"time"
)

type AccessToken struct {
	AppId     string
	AppSecret string
	TmpName   string
}

// get fresh access_token string
func (this *AccessToken) Fresh() (string, error) {
	if this.TmpName == "" {
		this.TmpName = "accesstoken.tmp"
	}
	fi, err := os.Stat(this.TmpName)
	if err != nil && !os.IsExist(err) {
		return this.fetchAndStore()
	}
	expires := fi.ModTime().Unix() + 7200
	if expires <= time.Now().Unix() {
		return this.fetchAndStore()
	}
	tmp, err := os.Open(this.TmpName)
	if err != nil {
		return "", err
	}
	defer tmp.Close()
	data, err := ioutil.ReadAll(tmp)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func (this *AccessToken) fetchAndStore() (string, error) {
	token, err := this.fetch()
	if err != nil {
		return "", err
	}
	if err := this.store(token); err != nil {
		return "", err
	}
	return token, nil
}

func (this *AccessToken) store(token string) error {
	tmp, err := os.OpenFile(this.TmpName, os.O_WRONLY|os.O_CREATE, os.ModeTemporary)
	if err != nil {
		return err
	}
	defer tmp.Close()
	if _, err := tmp.Write([]byte(token)); err != nil {
		return err
	}
	return nil
}

func (this *AccessToken) fetch() (string, error) {
	rtn, err := get(fmt.Sprintf(
		"%stoken?grant_type=client_credential&appid=%s&secret=%s",
		UrlPrefix,
		this.AppId,
		this.AppSecret,
	))
	if err != nil {
		return "", err
	}
	return rtn.AccessToken, nil
}
<commit_msg>auto create directories for accesstoken temporary file<commit_after>package weixinmp

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"time"
)

type AccessToken struct {
	AppId     string
	AppSecret string
	TmpName   string
}

// get fresh access_token string
func (this *AccessToken) Fresh() (string, error) {
	if this.TmpName == "" {
		this.TmpName = "accesstoken.tmp"
	}
	fi, err := os.Stat(this.TmpName)
	if err != nil && !os.IsExist(err) {
		return this.fetchAndStore()
	}
	expires := fi.ModTime().Unix() + 7200
	if expires <= time.Now().Unix() {
		return this.fetchAndStore()
	}
	tmp, err := os.Open(this.TmpName)
	if err != nil {
		return "", err
	}
	defer tmp.Close()
	data, err := ioutil.ReadAll(tmp)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func (this *AccessToken) fetchAndStore() (string, error) {
	token, err := this.fetch()
	if err != nil {
		return "", err
	}
	if err := this.store(token); err != nil {
		return "", err
	}
	return token, nil
}

func (this *AccessToken) store(token string) error {
	dir := path.Dir(this.TmpName)
	fi, err := os.Stat(dir)
	if os.IsNotExist(err) {
		if err := os.MkdirAll(dir, os.ModeTemporary); err != nil {
			return err
		}
	}
	if !fi.IsDir() {
		return errors.New("path is not a directory")
	}
	tmp, err := os.OpenFile(this.TmpName, os.O_WRONLY|os.O_CREATE, os.ModeTemporary)
	if err != nil {
		return err
	}
	defer tmp.Close()
	if _, err := tmp.Write([]byte(token)); err != nil {
		return err
	}
	return nil
}

func (this *AccessToken) fetch() (string, error) {
	rtn, err := get(fmt.Sprintf(
		"%stoken?grant_type=client_credential&appid=%s&secret=%s",
		UrlPrefix,
		this.AppId,
		this.AppSecret,
	))
	if err != nil {
		return "", err
	}
	return rtn.AccessToken, nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"sync"
)

var (
	//ErrBigWindow happens when we try to get more datapoints to average what
	//we can put in our accumulator. Good place to panic
	ErrBigWindow = errors.New("averaging window is bigger than accumulator")

	//ErrSmallSample happens when we just began and got not enough datapoints
	//to fill our window. Normal working situation.
	ErrSmallSample = errors.New("not enough data saved")

	//ErrWrongSize happens when signed int was expected to be positive but
	//ended up zero or negative. Sad, sad situation. Panic.
	ErrWrongSize = errors.New("size can not be less than one")
)

//Acc is a write-only ring buffer of finite and static capacity, with methods that
//provide sum and average of last n values pushed inside
type Acc struct {
	head int //as Acc is  read-only, there is no need to know position of tail
	vals []int64
	mu   sync.RWMutex
	full bool
}

//NewAcc returns accumulator of given size
func NewAcc(s int) *Acc {
	if s < 1 {
		panic(ErrWrongSize)
	}

	return &Acc{
		vals: make([]int64, s, s),
	}
}

//Push adds new value into Acc, overwriting old ones and wrapping
//around the ring as needed
func (a *Acc) Push(v int64) {
	a.mu.Lock()
	defer a.mu.Unlock()

	a.vals[a.head] = v

	if a.head == cap(a.vals)-1 {
		a.full = true
		a.head = 0
	} else {
		a.head = a.head + 1
	}
}

//Sum returns sum total of deltas in latest window of given size
func (a *Acc) Sum(w int) (sum int64, err error) {
	//if window is bigger than our accumulator, we fail horribly
	if w > cap(a.vals) {
		return 0, ErrBigWindow
	}

	//Critical section - DeltaAcc should not chage while we are in it
	a.mu.RLock()
	defer a.mu.RUnlock()

	//if we don't have enough data to fill window, we fail, less horribly
	if !a.full && a.head < w {
		return 0, ErrSmallSample
	}
	//if we need to get bits from different ends of our slice, let it be so
	if a.head < w {
		for _, v := range a.vals[:a.head] {
			sum += v
		}
		//pointer math, circular buffer, yay!
		for _, v := range a.vals[cap(a.vals)+a.head-w:] {
			sum += v
		}
	} else { //sane, classic situation - window is inside the slice
		for _, v := range a.vals[a.head-w : a.head] {
			sum += v
		}
	}
	return sum, nil
}

//Average returns sum average of deltas in latest window of given size
func (a *Acc) Average(w int) (avg float32, err error) {
	sum, err := a.Sum(w)
	avg = float32(sum) / float32(w)
	return
}

//DeltaAcc accumulates changes between pushed values.
//For DeltaAcc amount of remembered points is either size of underlying Acc or
//pushed amount minus one.
type DeltaAcc struct {
	last  uint64
	initd bool
	Acc
}

//NewDeltaAcc returns delta-accumulator of given size
func NewDeltaAcc(s int) *DeltaAcc {

	return &DeltaAcc{
		Acc: *NewAcc(s),
	}
}

//Push takes a value and adds difference between it and previous value into
//accumulator, with circular overwrite on full capacity
//First delta happens only after two values were pushed in
func (a *DeltaAcc) Push(v uint64) {
	defer func() { a.last = v }()

	if !a.initd {
		a.initd = true
		return
	}

	if a.last >= v {
		a.Acc.Push(-int64(a.last - v))
	} else { //stupid unsigned math
		a.Acc.Push(int64(v - a.last))
	}

}
<commit_msg>aligncheck fix<commit_after>package main

import (
	"errors"
	"sync"
)

var (
	//ErrBigWindow happens when we try to get more datapoints to average what
	//we can put in our accumulator. Good place to panic
	ErrBigWindow = errors.New("averaging window is bigger than accumulator")

	//ErrSmallSample happens when we just began and got not enough datapoints
	//to fill our window. Normal working situation.
	ErrSmallSample = errors.New("not enough data saved")

	//ErrWrongSize happens when signed int was expected to be positive but
	//ended up zero or negative. Sad, sad situation. Panic.
	ErrWrongSize = errors.New("size can not be less than one")
)

//Acc is a write-only ring buffer of finite and static capacity, with methods that
//provide sum and average of last n values pushed inside
type Acc struct {
	head int //as Acc is  read-only, there is no need to know position of tail
	vals []int64
	mu   sync.RWMutex
	full bool
}

//NewAcc returns accumulator of given size
func NewAcc(s int) *Acc {
	if s < 1 {
		panic(ErrWrongSize)
	}

	return &Acc{
		vals: make([]int64, s, s),
	}
}

//Push adds new value into Acc, overwriting old ones and wrapping
//around the ring as needed
func (a *Acc) Push(v int64) {
	a.mu.Lock()
	defer a.mu.Unlock()

	a.vals[a.head] = v

	if a.head == cap(a.vals)-1 {
		a.full = true
		a.head = 0
	} else {
		a.head = a.head + 1
	}
}

//Sum returns sum total of deltas in latest window of given size
func (a *Acc) Sum(w int) (sum int64, err error) {
	//if window is bigger than our accumulator, we fail horribly
	if w > cap(a.vals) {
		return 0, ErrBigWindow
	}

	//Critical section - DeltaAcc should not chage while we are in it
	a.mu.RLock()
	defer a.mu.RUnlock()

	//if we don't have enough data to fill window, we fail, less horribly
	if !a.full && a.head < w {
		return 0, ErrSmallSample
	}
	//if we need to get bits from different ends of our slice, let it be so
	if a.head < w {
		for _, v := range a.vals[:a.head] {
			sum += v
		}
		//pointer math, circular buffer, yay!
		for _, v := range a.vals[cap(a.vals)+a.head-w:] {
			sum += v
		}
	} else { //sane, classic situation - window is inside the slice
		for _, v := range a.vals[a.head-w : a.head] {
			sum += v
		}
	}
	return sum, nil
}

//Average returns sum average of deltas in latest window of given size
func (a *Acc) Average(w int) (avg float32, err error) {
	sum, err := a.Sum(w)
	avg = float32(sum) / float32(w)
	return
}

//DeltaAcc accumulates changes between pushed values.
//For DeltaAcc amount of remembered points is either size of underlying Acc or
//pushed amount minus one.
type DeltaAcc struct {
	last uint64
	Acc
	initd bool
}

//NewDeltaAcc returns delta-accumulator of given size
func NewDeltaAcc(s int) *DeltaAcc {

	return &DeltaAcc{
		Acc: *NewAcc(s),
	}
}

//Push takes a value and adds difference between it and previous value into
//accumulator, with circular overwrite on full capacity
//First delta happens only after two values were pushed in
func (a *DeltaAcc) Push(v uint64) {
	defer func() { a.last = v }()

	if !a.initd {
		a.initd = true
		return
	}

	if a.last >= v {
		a.Acc.Push(-int64(a.last - v))
	} else { //stupid unsigned math
		a.Acc.Push(int64(v - a.last))
	}

}
<|endoftext|>