<commit_before>// Container provides a Reader and Writer which serialize and deserialize gogen-avro structs to the Avro Object Container File (OCF) format. These are low-level primitives you probably don't need to use directly.
package container

import (
	"bytes"
	"compress/flate"
	"io"

	"github.com/actgardner/gogen-avro/container/avro"
)

// A Codec specifies how the blocks within a container file should be compressed.
type Codec string

const (
	// No compression
	Null Codec = "null"

	// Deflate compression
	Deflate Codec = "deflate"

	// Snappy compression
	Snappy Codec = "snappy"
)

type CloseableResettableWriter interface {
	Close() error
	Reset(io.Writer)
}

// Writer wraps an io.Writer and writes the file and block-level framing required for an OCF file.
// You can create a Writer for a given struct by calling the generated method `New<RecordType>Writer`.
type Writer struct {
	writer           io.Writer
	syncMarker       [16]byte
	codec            Codec
	recordsPerBlock  int64
	blockBuffer      *bytes.Buffer
	compressedWriter io.Writer
	nextBlockRecords int64
}

//  Create a new Writer wrapping the provided io.Writer with the given Codec and number of records per block.
//  The Writer will lazily write the container file header when WriteRecord is called the first time.
//  You must call Flush on the Writer before closing the underlying io.Writer, to ensure the final block is written.
//  A schema string must be passed to ensure that a correct header is written even if no records are written. This
//  is required to produce valid empty Avro container files.
func NewWriter(writer io.Writer, codec Codec, recordsPerBlock int64, schema string) (*Writer, error) {
	blockBytes := make([]byte, 0)
	blockBuffer := bytes.NewBuffer(blockBytes)

	avroWriter := &Writer{
		writer:          writer,
		syncMarker:      [16]byte{'g', 'o', 'g', 'e', 'n', 'a', 'v', 'r', 'o', 'm', 'a', 'g', 'i', 'c', '1', '0'},
		codec:           codec,
		recordsPerBlock: recordsPerBlock,
		blockBuffer:     blockBuffer,
	}
	var err error
	if codec == Deflate {
		avroWriter.compressedWriter, err = flate.NewWriter(avroWriter.blockBuffer, flate.DefaultCompression)
		if err != nil {
			return nil, err
		}
	} else if codec == Snappy {
		avroWriter.compressedWriter = newSnappyWriter(avroWriter.blockBuffer)
	} else {
		avroWriter.compressedWriter = avroWriter.blockBuffer
	}

	err = avroWriter.writeHeader(schema)
	if err != nil {
		return nil, err
	}

	return avroWriter, nil
}

func (avroWriter *Writer) writeHeader(schema string) error {
	header := &avro.AvroContainerHeader{
		Magic: [4]byte{'O', 'b', 'j', 1},
		Meta: map[string][]byte{
			"avro.schema": []byte(schema),
			"avro.codec":  []byte(avroWriter.codec),
		},
		Sync: avroWriter.syncMarker,
	}
	return header.Serialize(avroWriter.writer)
}

//  Write an AvroRecord to the container file. All gogen-avro generated structs
//  fulfill the AvroRecord interface. Note that all records in a given container file
//  must be of the same Avro type.
func (avroWriter *Writer) WriteRecord(record AvroRecord) error {
	var err error
	// Serialize the new record into the compressed writer
	err = record.Serialize(avroWriter.compressedWriter)
	if err != nil {
		return err
	}
	avroWriter.nextBlockRecords += 1

	// If the block if full, flush and reset the compressed writer,
	// write the header and the block contents
	if avroWriter.nextBlockRecords >= avroWriter.recordsPerBlock {
		return avroWriter.Flush()
	}

	return nil
}

//  Write the current block to the file if it has been filled.  It is
//  best-practise to always call this before the underlying io.Writer is closed.
func (avroWriter *Writer) Flush() error {
	if avroWriter.nextBlockRecords == 0 {
		return nil
	}

	// Write out all of the buffered records as a new block
	// Must be called before closing to ensure the last block is written
	if fwWriter, ok := avroWriter.compressedWriter.(CloseableResettableWriter); ok {
		fwWriter.Close()
		fwWriter.Reset(avroWriter.blockBuffer)
	}

	if avroWriter.nextBlockRecords > 0 {
		block := &avro.AvroContainerBlock{
			NumRecords:  avroWriter.nextBlockRecords,
			RecordBytes: avroWriter.blockBuffer.Bytes(),
			Sync:        avroWriter.syncMarker,
		}
		err := block.Serialize(avroWriter.writer)
		if err != nil {
			return err
		}
	}

	avroWriter.blockBuffer.Reset()
	avroWriter.nextBlockRecords = 0

	return nil
}
<commit_msg>add BlockBufferSize method (#106)<commit_after>// Container provides a Reader and Writer which serialize and deserialize gogen-avro structs to the Avro Object Container File (OCF) format. These are low-level primitives you probably don't need to use directly.
package container

import (
	"bytes"
	"compress/flate"
	"io"

	"github.com/actgardner/gogen-avro/container/avro"
)

// A Codec specifies how the blocks within a container file should be compressed.
type Codec string

const (
	// No compression
	Null Codec = "null"

	// Deflate compression
	Deflate Codec = "deflate"

	// Snappy compression
	Snappy Codec = "snappy"
)

type CloseableResettableWriter interface {
	Close() error
	Reset(io.Writer)
}

// Writer wraps an io.Writer and writes the file and block-level framing required for an OCF file.
// You can create a Writer for a given struct by calling the generated method `New<RecordType>Writer`.
type Writer struct {
	writer           io.Writer
	syncMarker       [16]byte
	codec            Codec
	recordsPerBlock  int64
	blockBuffer      *bytes.Buffer
	compressedWriter io.Writer
	nextBlockRecords int64
}

//  Create a new Writer wrapping the provided io.Writer with the given Codec and number of records per block.
//  The Writer will lazily write the container file header when WriteRecord is called the first time.
//  You must call Flush on the Writer before closing the underlying io.Writer, to ensure the final block is written.
//  A schema string must be passed to ensure that a correct header is written even if no records are written. This
//  is required to produce valid empty Avro container files.
func NewWriter(writer io.Writer, codec Codec, recordsPerBlock int64, schema string) (*Writer, error) {
	blockBytes := make([]byte, 0)
	blockBuffer := bytes.NewBuffer(blockBytes)

	avroWriter := &Writer{
		writer:          writer,
		syncMarker:      [16]byte{'g', 'o', 'g', 'e', 'n', 'a', 'v', 'r', 'o', 'm', 'a', 'g', 'i', 'c', '1', '0'},
		codec:           codec,
		recordsPerBlock: recordsPerBlock,
		blockBuffer:     blockBuffer,
	}
	var err error
	if codec == Deflate {
		avroWriter.compressedWriter, err = flate.NewWriter(avroWriter.blockBuffer, flate.DefaultCompression)
		if err != nil {
			return nil, err
		}
	} else if codec == Snappy {
		avroWriter.compressedWriter = newSnappyWriter(avroWriter.blockBuffer)
	} else {
		avroWriter.compressedWriter = avroWriter.blockBuffer
	}

	err = avroWriter.writeHeader(schema)
	if err != nil {
		return nil, err
	}

	return avroWriter, nil
}

func (avroWriter *Writer) writeHeader(schema string) error {
	header := &avro.AvroContainerHeader{
		Magic: [4]byte{'O', 'b', 'j', 1},
		Meta: map[string][]byte{
			"avro.schema": []byte(schema),
			"avro.codec":  []byte(avroWriter.codec),
		},
		Sync: avroWriter.syncMarker,
	}
	return header.Serialize(avroWriter.writer)
}

//  Write an AvroRecord to the container file. All gogen-avro generated structs
//  fulfill the AvroRecord interface. Note that all records in a given container file
//  must be of the same Avro type.
func (avroWriter *Writer) WriteRecord(record AvroRecord) error {
	var err error
	// Serialize the new record into the compressed writer
	err = record.Serialize(avroWriter.compressedWriter)
	if err != nil {
		return err
	}
	avroWriter.nextBlockRecords += 1

	// If the block if full, flush and reset the compressed writer,
	// write the header and the block contents
	if avroWriter.nextBlockRecords >= avroWriter.recordsPerBlock {
		return avroWriter.Flush()
	}

	return nil
}

//  Write the current block to the file if it has been filled.  It is
//  best-practise to always call this before the underlying io.Writer is closed.
func (avroWriter *Writer) Flush() error {
	if avroWriter.nextBlockRecords == 0 {
		return nil
	}

	// Write out all of the buffered records as a new block
	// Must be called before closing to ensure the last block is written
	if fwWriter, ok := avroWriter.compressedWriter.(CloseableResettableWriter); ok {
		fwWriter.Close()
		fwWriter.Reset(avroWriter.blockBuffer)
	}

	if avroWriter.nextBlockRecords > 0 {
		block := &avro.AvroContainerBlock{
			NumRecords:  avroWriter.nextBlockRecords,
			RecordBytes: avroWriter.blockBuffer.Bytes(),
			Sync:        avroWriter.syncMarker,
		}
		err := block.Serialize(avroWriter.writer)
		if err != nil {
			return err
		}
	}

	avroWriter.blockBuffer.Reset()
	avroWriter.nextBlockRecords = 0

	return nil
}

// Get the current block buffer size.
// caller might trigger an early Flush if current block size gets huge.
func (avroWriter *Writer) BlockBufferSize() int {
	return avroWriter.blockBuffer.Len()
}
<|endoftext|>


<commit_before>package bubbles

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math"
	"net"
	"net/http"
	"sync"
	"time"
)

const (
	// DefaultMaxDocumentsPerBatch is the number of documents a batch needs to
	// have before it is send. This is per connection.
	DefaultMaxDocumentsPerBatch = 1000

	// DefaultFlushTimeout is the maximum time we batch something before we try
	// to send it to a server.
	DefaultFlushTimeout = 10 * time.Second

	// DefaultServerTimeout is the time we give ElasticSearch to respond. This
	// is also the maximum time Stop() will take.
	DefaultServerTimeout = 10 * time.Second

	// DefaultConnCount is the number of connections per hosts.
	DefaultConnCount = 2

	// backoffTimeoutRatio determines when we start backing off.
	backoffTimeoutRatio = 2

	serverErrorWait    = 500 * time.Millisecond
	serverErrorWaitMax = 16 * time.Second

	defaultElasticSearchPort = "9200"
)

var (
	errInvalidResponse = errors.New("invalid response")
)

// Bubbles is the main struct to control a queue of Actions going to the
// ElasticSearch servers.
type Bubbles struct {
	q                chan Action
	retryQ           chan Action
	quit             chan struct{}
	wg               sync.WaitGroup
	maxDocumentCount int
	connCount        int
	flushTimeout     time.Duration
	serverTimeout    time.Duration
	c                Counter
	e                Errer
}

// Opt is any option to New().
type Opt func(*Bubbles)

// OptConnCount is an option to New() to specify the number of connections per
// host. The default is DefaultConnCount.
func OptConnCount(n int) Opt {
	return func(b *Bubbles) {
		b.connCount = n
	}
}

// OptFlush is an option to New() to specify the flush timeout of a batch. The
// default is DefaultFlushTimeout.
func OptFlush(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.flushTimeout = d
	}
}

// OptServerTimeout is an option to New() to specify the timeout of a single
// batch POST to ElasticSearch. This value is also the maximum time Stop() will
// take. All actions in a bulk which is timed out will be retried. The default
// is DefaultServerTimeout.
func OptServerTimeout(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.serverTimeout = d
	}
}

// OptMaxDocs is an option to New() to specify maximum number of documents in a
// single batch. The default is DefaultMaxDocumentsPerBatch.
func OptMaxDocs(n int) Opt {
	return func(b *Bubbles) {
		b.maxDocumentCount = n
	}
}

// OptCounter is an option to New() to specify something that counts documents.
func OptCounter(c Counter) Opt {
	return func(b *Bubbles) {
		b.c = c
	}
}

// OptErrer is an option to New() to specify an error handler. The default
// handler uses the log module.
func OptErrer(e Errer) Opt {
	return func(b *Bubbles) {
		b.e = e
	}
}

// New makes a new ElasticSearch bulk inserter. It needs a list with 'ip' or
// 'ip:port' addresses, options are added via the Opt* functions.
func New(addrs []string, opts ...Opt) *Bubbles {
	b := Bubbles{
		q:                make(chan Action),
		quit:             make(chan struct{}),
		maxDocumentCount: DefaultMaxDocumentsPerBatch,
		connCount:        DefaultConnCount,
		flushTimeout:     DefaultFlushTimeout,
		serverTimeout:    DefaultServerTimeout,
		c:                DefaultCounter{},
		e:                DefaultErrer{},
	}
	for _, o := range opts {
		o(&b)
	}
	b.retryQ = make(chan Action, len(addrs)*b.connCount*b.maxDocumentCount)

	cl := &http.Client{
		Timeout: b.serverTimeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return errors.New("no redirect")
		},
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout:   b.serverTimeout,
				KeepAlive: 30 * time.Second,
			}).Dial,
			MaxIdleConnsPerHost: b.connCount,
			DisableCompression:  false,
		},
	}
	// Start a go routine per connection per host
	for _, a := range addrs {
		addr := withPort(a, defaultElasticSearchPort)
		for i := 0; i < b.connCount; i++ {
			b.wg.Add(1)
			go func(a string) {
				client(&b, cl, a)
				b.wg.Done()
			}(addr)
		}
	}
	return &b
}

// Enqueue returns the queue to add Actions in a routine. It will block if all bulk
// processors are busy.
func (b *Bubbles) Enqueue() chan<- Action {
	return b.q
}

// Stop shuts down all ElasticSearch clients. It'll return all Action entries
// which were not yet processed, or were up for a retry. It can take
// OptServerTimeout to complete.
func (b *Bubbles) Stop() []Action {
	close(b.quit)
	// There is no explicit timeout, we rely on b.serverTimeout to shut down
	// everything.
	b.wg.Wait()

	// Collect and return elements which are in flight.
	close(b.retryQ)
	close(b.q)
	pending := make([]Action, 0, len(b.q)+len(b.retryQ))
	for a := range b.q {
		pending = append(pending, a)
	}
	for a := range b.retryQ {
		pending = append(pending, a)
	}
	return pending
}

type backoff struct {
	level        uint8
	max          uint8
	maxBatchSize int
}

func newBackoff(maxBatchSize int) *backoff {
	b := &backoff{
		maxBatchSize: maxBatchSize,
	}
	// Max out when both components do.
	for ; (b.wait() < serverErrorWaitMax || b.size() > 1) && b.level < math.MaxUint8; b.level++ {
	}
	b.max = b.level
	b.level = 0
	return b
}

// wait calculates the delay based on the current backoff level.
func (b *backoff) wait() time.Duration {
	if b.level == 0 {
		return 0 * time.Second
	}
	w := (1 << (b.level - 1)) * serverErrorWait
	if w >= serverErrorWaitMax {
		return serverErrorWaitMax
	}
	return w
}

// batchSize calculates the batchsize based on the current backoff level.
func (b *backoff) size() int {
	s := b.maxBatchSize / (1 << b.level)
	if s <= 1 {
		return 1
	}
	return s
}

// inc increases the backoff level.
func (b *backoff) inc() {
	if b.level < b.max {
		b.level++
	}
}

// dec decreases the backoff level.
func (b *backoff) dec() {
	if b.level > 0 {
		b.level--
	}
}

// client talks to ElasticSearch. This runs in a go routine in a loop and deals
// with a single ElasticSearch address.
func client(b *Bubbles, cl *http.Client, addr string) {
	url := fmt.Sprintf("http://%s/_bulk", addr)

	backoff := newBackoff(b.maxDocumentCount)
	backoffTime := b.serverTimeout / backoffTimeoutRatio
	for {
		select {
		case <-b.quit:
			return
		case <-time.After(backoff.wait()):
		}
		trouble, batchTime := runBatch(b, cl, url, backoff.size())
		if trouble || batchTime > backoffTime {
			backoff.inc()
			b.c.Timeout()
		} else {
			backoff.dec()
		}
		b.c.BatchTime(batchTime)
	}
}

// runBatch gathers and deals with a batch of actions. It returns
// whether there was trouble, and how long the actual request took.
func runBatch(b *Bubbles, cl *http.Client, url string, batchSize int) (bool, time.Duration) {
	actions := make([]Action, 0, b.maxDocumentCount)
	// First use all retry actions.
retry:
	for len(actions) < batchSize {
		select {
		case a := <-b.retryQ:
			actions = append(actions, a)
		default:
			// no more retry actions queued
			break retry
		}
	}

	var t <-chan time.Time
gather:
	for len(actions) < batchSize {
		if t == nil && len(actions) > 0 {
			// Set timeout on the first element we read
			t = time.After(b.flushTimeout)
		}
		select {
		case <-b.quit:
			for _, a := range actions {
				b.retryQ <- a
			}
			return false, 0
		case <-t:
			// this case is not enabled until we've got an action
			break gather
		case a := <-b.retryQ:
			actions = append(actions, a)
		case a := <-b.q:
			actions = append(actions, a)
		}
	}

	t0 := time.Now()
	res, err := postActions(b.c, cl, url, actions)
	dt := time.Since(t0)
	if err != nil {
		// A server error. Retry these actions later.
		b.e.Error(err)
		for _, a := range actions {
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
		}
		return true, dt
	}

	// Server has accepted the request an sich, but there can be errors in the
	// individual actions.
	if !res.Errors {
		// Simple case, no errors present.
		return false, dt
	}

	// Invalid response from ElasticSearch.
	if len(actions) != len(res.Items) {
		b.e.Error(errInvalidResponse)
		for _, a := range actions {
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
		}
		return true, dt
	}
	// Figure out which actions have errors.
	for i, e := range res.Items {
		a := actions[i]
		el, ok := e[string(a.Type)]
		if !ok {
			// Unexpected reply from ElasticSearch.
			b.e.Error(errInvalidResponse)
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
			continue
		}

		c := el.Status
		switch {
		case c >= 200 && c < 300:
			// Document accepted by ElasticSearch.
		case c == 429 || (c >= 500 && c < 600):
			// Server error. Retry it.
			// We get a 429 when the bulk queue is full, which we just retry as
			// well.
			b.e.Warning(ActionError{
				Action:     a,
				StatusCode: c,
				Msg:        fmt.Sprintf("transient error %d: %s", c, el.Error),
				Server:     url,
			})
			b.c.Retry(RetryTransient, a.Type, len(a.Document))
			b.retryQ <- a
		case c >= 400 && c < 500:
			// Some error. Nothing we can do with it.
			b.e.Error(ActionError{
				Action:     a,
				StatusCode: c,
				Msg:        fmt.Sprintf("error %d: %s", c, el.Error),
				Server:     url,
			})
		default:
			// No idea.
			b.e.Error(fmt.Errorf("unwelcome response %d: %s", c, el.Error))
		}
	}
	return true, dt
}

type bulkRes struct {
	Took   int  `json:"took"`
	Errors bool `json:"errors"`
	Items  []map[string]struct {
		Index   string `json:"_index"`
		Type    string `json:"_type"`
		ID      string `json:"_id"`
		Version int    `json:"_version"`
		Status  int    `json:"status"`
		Error   string `json:"error"`
	} `json:"items"`
}

func postActions(c Counter, cl *http.Client, url string, actions []Action) (*bulkRes, error) {
	buf := bytes.Buffer{}
	for _, a := range actions {
		c.Send(a.Type, len(a.Document))
		buf.Write(a.Buf())
	}
	c.SendTotal(buf.Len())

	// This doesn't Chunk.
	resp, err := cl.Post(url, "application/json", &buf)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("status %d: %s", resp.StatusCode, string(body))
	}

	var bulk bulkRes
	if err := json.Unmarshal(body, &bulk); err != nil {
		return nil, err
	}
	return &bulk, nil
}

// ActionError wraps an Action we won't retry. It implements the error interface.
type ActionError struct {
	Action     Action
	StatusCode int
	Msg        string
	Server     string
}

func (e ActionError) Error() string {
	return fmt.Sprintf("%s: %s %s", e.Server, e.Action.Type, e.Msg)
}

// withPort adds a default port to an address string.
func withPort(a, port string) string {
	if _, _, err := net.SplitHostPort(a); err != nil {
		// no port found.
		return net.JoinHostPort(a, port)
	}
	return a
}
<commit_msg>Cancel in-flight requests on quit.<commit_after>package bubbles

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math"
	"net"
	"net/http"
	"sync"
	"time"
)

const (
	// DefaultMaxDocumentsPerBatch is the number of documents a batch needs to
	// have before it is send. This is per connection.
	DefaultMaxDocumentsPerBatch = 1000

	// DefaultFlushTimeout is the maximum time we batch something before we try
	// to send it to a server.
	DefaultFlushTimeout = 10 * time.Second

	// DefaultServerTimeout is the time we give ElasticSearch to respond. This
	// is also the maximum time Stop() will take.
	DefaultServerTimeout = 10 * time.Second

	// DefaultConnCount is the number of connections per hosts.
	DefaultConnCount = 2

	// backoffTimeoutRatio determines when we start backing off.
	backoffTimeoutRatio = 2

	serverErrorWait    = 500 * time.Millisecond
	serverErrorWaitMax = 16 * time.Second

	defaultElasticSearchPort = "9200"
)

var (
	errInvalidResponse = errors.New("invalid response")
)

// Bubbles is the main struct to control a queue of Actions going to the
// ElasticSearch servers.
type Bubbles struct {
	q                chan Action
	retryQ           chan Action
	quit             chan struct{}
	wg               sync.WaitGroup
	maxDocumentCount int
	connCount        int
	flushTimeout     time.Duration
	serverTimeout    time.Duration
	c                Counter
	e                Errer
}

// Opt is any option to New().
type Opt func(*Bubbles)

// OptConnCount is an option to New() to specify the number of connections per
// host. The default is DefaultConnCount.
func OptConnCount(n int) Opt {
	return func(b *Bubbles) {
		b.connCount = n
	}
}

// OptFlush is an option to New() to specify the flush timeout of a batch. The
// default is DefaultFlushTimeout.
func OptFlush(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.flushTimeout = d
	}
}

// OptServerTimeout is an option to New() to specify the timeout of a single
// batch POST to ElasticSearch. This value is also the maximum time Stop() will
// take. All actions in a bulk which is timed out will be retried. The default
// is DefaultServerTimeout.
func OptServerTimeout(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.serverTimeout = d
	}
}

// OptMaxDocs is an option to New() to specify maximum number of documents in a
// single batch. The default is DefaultMaxDocumentsPerBatch.
func OptMaxDocs(n int) Opt {
	return func(b *Bubbles) {
		b.maxDocumentCount = n
	}
}

// OptCounter is an option to New() to specify something that counts documents.
func OptCounter(c Counter) Opt {
	return func(b *Bubbles) {
		b.c = c
	}
}

// OptErrer is an option to New() to specify an error handler. The default
// handler uses the log module.
func OptErrer(e Errer) Opt {
	return func(b *Bubbles) {
		b.e = e
	}
}

// New makes a new ElasticSearch bulk inserter. It needs a list with 'ip' or
// 'ip:port' addresses, options are added via the Opt* functions.
func New(addrs []string, opts ...Opt) *Bubbles {
	b := Bubbles{
		q:                make(chan Action),
		quit:             make(chan struct{}),
		maxDocumentCount: DefaultMaxDocumentsPerBatch,
		connCount:        DefaultConnCount,
		flushTimeout:     DefaultFlushTimeout,
		serverTimeout:    DefaultServerTimeout,
		c:                DefaultCounter{},
		e:                DefaultErrer{},
	}
	for _, o := range opts {
		o(&b)
	}
	b.retryQ = make(chan Action, len(addrs)*b.connCount*b.maxDocumentCount)

	cl := &http.Client{
		Timeout: b.serverTimeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return errors.New("no redirect")
		},
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout:   b.serverTimeout,
				KeepAlive: 30 * time.Second,
			}).Dial,
			MaxIdleConnsPerHost: b.connCount,
			DisableCompression:  false,
		},
	}
	// Start a go routine per connection per host
	for _, a := range addrs {
		addr := withPort(a, defaultElasticSearchPort)
		for i := 0; i < b.connCount; i++ {
			b.wg.Add(1)
			go func(a string) {
				client(&b, cl, a)
				b.wg.Done()
			}(addr)
		}
	}
	return &b
}

// Enqueue returns the queue to add Actions in a routine. It will block if all bulk
// processors are busy.
func (b *Bubbles) Enqueue() chan<- Action {
	return b.q
}

// Stop shuts down all ElasticSearch clients. It'll return all Action entries
// which were not yet processed, or were up for a retry. It can take
// OptServerTimeout to complete.
func (b *Bubbles) Stop() []Action {
	close(b.quit)
	// There is no explicit timeout, we rely on b.serverTimeout to shut down
	// everything.
	b.wg.Wait()

	// Collect and return elements which are in flight.
	close(b.retryQ)
	close(b.q)
	pending := make([]Action, 0, len(b.q)+len(b.retryQ))
	for a := range b.q {
		pending = append(pending, a)
	}
	for a := range b.retryQ {
		pending = append(pending, a)
	}
	return pending
}

type backoff struct {
	level        uint8
	max          uint8
	maxBatchSize int
}

func newBackoff(maxBatchSize int) *backoff {
	b := &backoff{
		maxBatchSize: maxBatchSize,
	}
	// Max out when both components do.
	for ; (b.wait() < serverErrorWaitMax || b.size() > 1) && b.level < math.MaxUint8; b.level++ {
	}
	b.max = b.level
	b.level = 0
	return b
}

// wait calculates the delay based on the current backoff level.
func (b *backoff) wait() time.Duration {
	if b.level == 0 {
		return 0 * time.Second
	}
	w := (1 << (b.level - 1)) * serverErrorWait
	if w >= serverErrorWaitMax {
		return serverErrorWaitMax
	}
	return w
}

// batchSize calculates the batchsize based on the current backoff level.
func (b *backoff) size() int {
	s := b.maxBatchSize / (1 << b.level)
	if s <= 1 {
		return 1
	}
	return s
}

// inc increases the backoff level.
func (b *backoff) inc() {
	if b.level < b.max {
		b.level++
	}
}

// dec decreases the backoff level.
func (b *backoff) dec() {
	if b.level > 0 {
		b.level--
	}
}

// client talks to ElasticSearch. This runs in a go routine in a loop and deals
// with a single ElasticSearch address.
func client(b *Bubbles, cl *http.Client, addr string) {
	url := fmt.Sprintf("http://%s/_bulk", addr)

	backoff := newBackoff(b.maxDocumentCount)
	backoffTime := b.serverTimeout / backoffTimeoutRatio
	for {
		select {
		case <-b.quit:
			return
		case <-time.After(backoff.wait()):
		}
		trouble, batchTime := runBatch(b, cl, url, backoff.size())
		if trouble || batchTime > backoffTime {
			backoff.inc()
			b.c.Timeout()
		} else {
			backoff.dec()
		}
		b.c.BatchTime(batchTime)
	}
}

// runBatch gathers and deals with a batch of actions. It returns
// whether there was trouble, and how long the actual request took.
func runBatch(b *Bubbles, cl *http.Client, url string, batchSize int) (bool, time.Duration) {
	actions := make([]Action, 0, b.maxDocumentCount)
	// First use all retry actions.
retry:
	for len(actions) < batchSize {
		select {
		case a := <-b.retryQ:
			actions = append(actions, a)
		default:
			// no more retry actions queued
			break retry
		}
	}

	var t <-chan time.Time
gather:
	for len(actions) < batchSize {
		if t == nil && len(actions) > 0 {
			// Set timeout on the first element we read
			t = time.After(b.flushTimeout)
		}
		select {
		case <-b.quit:
			for _, a := range actions {
				b.retryQ <- a
			}
			return false, 0
		case <-t:
			// this case is not enabled until we've got an action
			break gather
		case a := <-b.retryQ:
			actions = append(actions, a)
		case a := <-b.q:
			actions = append(actions, a)
		}
	}

	t0 := time.Now()
	res, err := postActions(b.c, cl, url, actions, b.quit)
	dt := time.Since(t0)
	if err != nil {
		// A server error. Retry these actions later.
		b.e.Error(err)
		for _, a := range actions {
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
		}
		return true, dt
	}

	// Server has accepted the request an sich, but there can be errors in the
	// individual actions.
	if !res.Errors {
		// Simple case, no errors present.
		return false, dt
	}

	// Invalid response from ElasticSearch.
	if len(actions) != len(res.Items) {
		b.e.Error(errInvalidResponse)
		for _, a := range actions {
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
		}
		return true, dt
	}
	// Figure out which actions have errors.
	for i, e := range res.Items {
		a := actions[i]
		el, ok := e[string(a.Type)]
		if !ok {
			// Unexpected reply from ElasticSearch.
			b.e.Error(errInvalidResponse)
			b.c.Retry(RetryUnlikely, a.Type, len(a.Document))
			b.retryQ <- a
			continue
		}

		c := el.Status
		switch {
		case c >= 200 && c < 300:
			// Document accepted by ElasticSearch.
		case c == 429 || (c >= 500 && c < 600):
			// Server error. Retry it.
			// We get a 429 when the bulk queue is full, which we just retry as
			// well.
			b.e.Warning(ActionError{
				Action:     a,
				StatusCode: c,
				Msg:        fmt.Sprintf("transient error %d: %s", c, el.Error),
				Server:     url,
			})
			b.c.Retry(RetryTransient, a.Type, len(a.Document))
			b.retryQ <- a
		case c >= 400 && c < 500:
			// Some error. Nothing we can do with it.
			b.e.Error(ActionError{
				Action:     a,
				StatusCode: c,
				Msg:        fmt.Sprintf("error %d: %s", c, el.Error),
				Server:     url,
			})
		default:
			// No idea.
			b.e.Error(fmt.Errorf("unwelcome response %d: %s", c, el.Error))
		}
	}
	return true, dt
}

type bulkRes struct {
	Took   int  `json:"took"`
	Errors bool `json:"errors"`
	Items  []map[string]struct {
		Index   string `json:"_index"`
		Type    string `json:"_type"`
		ID      string `json:"_id"`
		Version int    `json:"_version"`
		Status  int    `json:"status"`
		Error   string `json:"error"`
	} `json:"items"`
}

func interruptibleDo(cl *http.Client, req *http.Request, interrupt <-chan struct{}) (*http.Response, error) {
	done := make(chan struct{})
	go func() {
		select {
		case <-interrupt:
			cl.Transport.(*http.Transport).CancelRequest(req)
		case <-done:
		}
	}()
	defer close(done)
	return cl.Do(req)
}

func postActions(c Counter, cl *http.Client, url string, actions []Action, quit <-chan struct{}) (*bulkRes, error) {
	buf := bytes.Buffer{}
	for _, a := range actions {
		c.Send(a.Type, len(a.Document))
		buf.Write(a.Buf())
	}
	c.SendTotal(buf.Len())

	// This doesn't Chunk.
	req, err := http.NewRequest("POST", url, &buf)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := interruptibleDo(cl, req, quit)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("status %d: %s", resp.StatusCode, string(body))
	}

	var bulk bulkRes
	if err := json.Unmarshal(body, &bulk); err != nil {
		return nil, err
	}
	return &bulk, nil
}

// ActionError wraps an Action we won't retry. It implements the error interface.
type ActionError struct {
	Action     Action
	StatusCode int
	Msg        string
	Server     string
}

func (e ActionError) Error() string {
	return fmt.Sprintf("%s: %s %s", e.Server, e.Action.Type, e.Msg)
}

// withPort adds a default port to an address string.
func withPort(a, port string) string {
	if _, _, err := net.SplitHostPort(a); err != nil {
		// no port found.
		return net.JoinHostPort(a, port)
	}
	return a
}
<|endoftext|>


<commit_before>package common

import (
	"sync"

	"github.com/hashicorp/packer/helper/multistep"
)

type DriverMock struct {
	sync.Mutex

	CloneCalled bool
	CloneDst    string
	CloneSrc    string
	CloneErr    error

	CompactDiskCalled bool
	CompactDiskPath   string
	CompactDiskErr    error

	CreateDiskCalled bool
	CreateDiskOutput string
	CreateDiskSize   string
	CreateDiskTypeId string
	CreateDiskErr    error

	IsRunningCalled bool
	IsRunningPath   string
	IsRunningResult bool
	IsRunningErr    error

	CommHostCalled bool
	CommHostState  multistep.StateBag
	CommHostResult string
	CommHostErr    error

	HostAddressCalled bool
	HostAddressState  multistep.StateBag
	HostAddressResult string
	HostAddressErr    error

	HostIPCalled bool
	HostIPState  multistep.StateBag
	HostIPResult string
	HostIPErr    error

	GuestAddressCalled bool
	GuestAddressState  multistep.StateBag
	GuestAddressResult string
	GuestAddressErr    error

	GuestIPCalled bool
	GuestIPState  multistep.StateBag
	GuestIPResult string
	GuestIPErr    error

	StartCalled   bool
	StartPath     string
	StartHeadless bool
	StartErr      error

	StopCalled bool
	StopPath   string
	StopErr    error

	SuppressMessagesCalled bool
	SuppressMessagesPath   string
	SuppressMessagesErr    error

	ToolsIsoPathCalled bool
	ToolsIsoPathFlavor string
	ToolsIsoPathResult string

	ToolsInstallCalled bool
	ToolsInstallErr    error

	DhcpLeasesPathCalled bool
	DhcpLeasesPathDevice string
	DhcpLeasesPathResult string

	DhcpConfPathCalled bool
	DhcpConfPathResult string

	VmnetnatConfPathCalled bool
	VmnetnatConfPathResult string

	NetmapConfPathCalled bool
	NetmapConfPathResult string

	VerifyCalled bool
	VerifyErr    error
}

func (d *DriverMock) Clone(dst string, src string) error {
	d.CloneCalled = true
	d.CloneDst = dst
	d.CloneSrc = src
	return d.CloneErr
}

func (d *DriverMock) CompactDisk(path string) error {
	d.CompactDiskCalled = true
	d.CompactDiskPath = path
	return d.CompactDiskErr
}

func (d *DriverMock) CreateDisk(output string, size string, typeId string) error {
	d.CreateDiskCalled = true
	d.CreateDiskOutput = output
	d.CreateDiskSize = size
	d.CreateDiskTypeId = typeId
	return d.CreateDiskErr
}

func (d *DriverMock) IsRunning(path string) (bool, error) {
	d.Lock()
	defer d.Unlock()

	d.IsRunningCalled = true
	d.IsRunningPath = path
	return d.IsRunningResult, d.IsRunningErr
}

func (d *DriverMock) CommHost(state multistep.StateBag) (string, error) {
	d.CommHostCalled = true
	d.CommHostState = state
	return d.CommHostResult, d.CommHostErr
}

func (d *DriverMock) HostAddress(state multistep.StateBag) (string, error) {
	d.HostAddressCalled = true
	d.HostAddressState = state
	return d.HostAddressResult, d.HostAddressErr
}

func (d *DriverMock) HostIP(state multistep.StateBag) (string, error) {
	d.HostIPCalled = true
	d.HostIPState = state
	return d.HostIPResult, d.HostIPErr
}

func (d *DriverMock) GuestAddress(state multistep.StateBag) (string, error) {
	d.GuestAddressCalled = true
	d.GuestAddressState = state
	return d.GuestAddressResult, d.GuestAddressErr
}

func (d *DriverMock) GuestIP(state multistep.StateBag) (string, error) {
	d.GuestIPCalled = true
	d.GuestIPState = state
	return d.GuestIPResult, d.GuestIPErr
}

func (d *DriverMock) Start(path string, headless bool) error {
	d.StartCalled = true
	d.StartPath = path
	d.StartHeadless = headless
	return d.StartErr
}

func (d *DriverMock) Stop(path string) error {
	d.StopCalled = true
	d.StopPath = path
	return d.StopErr
}

func (d *DriverMock) SuppressMessages(path string) error {
	d.SuppressMessagesCalled = true
	d.SuppressMessagesPath = path
	return d.SuppressMessagesErr
}

func (d *DriverMock) ToolsIsoPath(flavor string) string {
	d.ToolsIsoPathCalled = true
	d.ToolsIsoPathFlavor = flavor
	return d.ToolsIsoPathResult
}

func (d *DriverMock) ToolsInstall() error {
	d.ToolsInstallCalled = true
	return d.ToolsInstallErr
}

func (d *DriverMock) DhcpLeasesPath(device string) string {
	d.DhcpLeasesPathCalled = true
	d.DhcpLeasesPathDevice = device
	return d.DhcpLeasesPathResult
}

func (d *DriverMock) DhcpConfPath(device string) string {
	d.DhcpConfPathCalled = true
	return d.DhcpConfPathResult
}

func (d *DriverMock) VmnetnatConfPath(device string) string {
	d.VmnetnatConfPathCalled = true
	return d.VmnetnatConfPathResult
}

func (d *DriverMock) NetmapConfPath() string {
	d.NetmapConfPathCalled = true
	return d.NetmapConfPathResult
}

func (d *DriverMock) Verify() error {
	d.VerifyCalled = true
	return d.VerifyErr
}

func (d *DriverMock) GetVmwareDriver() VmwareDriver {
	var state VmwareDriver
	state.DhcpLeasesPath = func(string) string {
		return "/path/to/dhcp.leases"
	}
	state.DhcpConfPath = func(string) string {
		return "/path/to/dhcp.conf"
	}
	state.VmnetnatConfPath = func(string) string {
		return "/path/to/vmnetnat.conf"
	}
	state.NetmapConfPath = func() string {
		return "/path/to/netmap.conf"
	}
	return state
}
<commit_msg>Fixed oversight in VMware builder's mock-driver that neglected to initialize 'HostAddressResult'.<commit_after>package common

import (
	"net"
	"sync"

	"github.com/hashicorp/packer/helper/multistep"
)

type DriverMock struct {
	sync.Mutex

	CloneCalled bool
	CloneDst    string
	CloneSrc    string
	CloneErr    error

	CompactDiskCalled bool
	CompactDiskPath   string
	CompactDiskErr    error

	CreateDiskCalled bool
	CreateDiskOutput string
	CreateDiskSize   string
	CreateDiskTypeId string
	CreateDiskErr    error

	IsRunningCalled bool
	IsRunningPath   string
	IsRunningResult bool
	IsRunningErr    error

	CommHostCalled bool
	CommHostState  multistep.StateBag
	CommHostResult string
	CommHostErr    error

	HostAddressCalled bool
	HostAddressState  multistep.StateBag
	HostAddressResult string
	HostAddressErr    error

	HostIPCalled bool
	HostIPState  multistep.StateBag
	HostIPResult string
	HostIPErr    error

	GuestAddressCalled bool
	GuestAddressState  multistep.StateBag
	GuestAddressResult string
	GuestAddressErr    error

	GuestIPCalled bool
	GuestIPState  multistep.StateBag
	GuestIPResult string
	GuestIPErr    error

	StartCalled   bool
	StartPath     string
	StartHeadless bool
	StartErr      error

	StopCalled bool
	StopPath   string
	StopErr    error

	SuppressMessagesCalled bool
	SuppressMessagesPath   string
	SuppressMessagesErr    error

	ToolsIsoPathCalled bool
	ToolsIsoPathFlavor string
	ToolsIsoPathResult string

	ToolsInstallCalled bool
	ToolsInstallErr    error

	DhcpLeasesPathCalled bool
	DhcpLeasesPathDevice string
	DhcpLeasesPathResult string

	DhcpConfPathCalled bool
	DhcpConfPathResult string

	VmnetnatConfPathCalled bool
	VmnetnatConfPathResult string

	NetmapConfPathCalled bool
	NetmapConfPathResult string

	VerifyCalled bool
	VerifyErr    error
}

func (d *DriverMock) Clone(dst string, src string) error {
	d.CloneCalled = true
	d.CloneDst = dst
	d.CloneSrc = src
	return d.CloneErr
}

func (d *DriverMock) CompactDisk(path string) error {
	d.CompactDiskCalled = true
	d.CompactDiskPath = path
	return d.CompactDiskErr
}

func (d *DriverMock) CreateDisk(output string, size string, typeId string) error {
	d.CreateDiskCalled = true
	d.CreateDiskOutput = output
	d.CreateDiskSize = size
	d.CreateDiskTypeId = typeId
	return d.CreateDiskErr
}

func (d *DriverMock) IsRunning(path string) (bool, error) {
	d.Lock()
	defer d.Unlock()

	d.IsRunningCalled = true
	d.IsRunningPath = path
	return d.IsRunningResult, d.IsRunningErr
}

func (d *DriverMock) CommHost(state multistep.StateBag) (string, error) {
	d.CommHostCalled = true
	d.CommHostState = state
	return d.CommHostResult, d.CommHostErr
}

func MockInterface() net.Interface {
	interfaces, err := net.Interfaces()

	// Build a dummy interface due to being unable to enumerate interfaces
	if err != nil || len(interfaces) == 0 {
		return net.Interface{
			Index:        0,
			MTU:          -1,
			Name:         "dummy",
			HardwareAddr: net.HardwareAddr{0, 0, 0, 0, 0, 0},
			Flags:        net.FlagLoopback,
		}
	}

	// Find the first loopback interface
	for _, intf := range interfaces {
		if intf.Flags&net.FlagLoopback == net.FlagLoopback {
			return intf
		}
	}

	// Fall-back to just the first one
	return interfaces[0]
}

func (d *DriverMock) HostAddress(state multistep.StateBag) (string, error) {
	intf := MockInterface()
	d.HostAddressResult = intf.HardwareAddr.String()
	d.HostAddressCalled = true
	d.HostAddressState = state
	return d.HostAddressResult, d.HostAddressErr
}

func (d *DriverMock) HostIP(state multistep.StateBag) (string, error) {
	d.HostIPResult = "127.0.0.1"
	d.HostIPCalled = true
	d.HostIPState = state
	return d.HostIPResult, d.HostIPErr
}

func (d *DriverMock) GuestAddress(state multistep.StateBag) (string, error) {
	d.GuestAddressCalled = true
	d.GuestAddressState = state
	return d.GuestAddressResult, d.GuestAddressErr
}

func (d *DriverMock) GuestIP(state multistep.StateBag) (string, error) {
	d.GuestIPCalled = true
	d.GuestIPState = state
	return d.GuestIPResult, d.GuestIPErr
}

func (d *DriverMock) Start(path string, headless bool) error {
	d.StartCalled = true
	d.StartPath = path
	d.StartHeadless = headless
	return d.StartErr
}

func (d *DriverMock) Stop(path string) error {
	d.StopCalled = true
	d.StopPath = path
	return d.StopErr
}

func (d *DriverMock) SuppressMessages(path string) error {
	d.SuppressMessagesCalled = true
	d.SuppressMessagesPath = path
	return d.SuppressMessagesErr
}

func (d *DriverMock) ToolsIsoPath(flavor string) string {
	d.ToolsIsoPathCalled = true
	d.ToolsIsoPathFlavor = flavor
	return d.ToolsIsoPathResult
}

func (d *DriverMock) ToolsInstall() error {
	d.ToolsInstallCalled = true
	return d.ToolsInstallErr
}

func (d *DriverMock) DhcpLeasesPath(device string) string {
	d.DhcpLeasesPathCalled = true
	d.DhcpLeasesPathDevice = device
	return d.DhcpLeasesPathResult
}

func (d *DriverMock) DhcpConfPath(device string) string {
	d.DhcpConfPathCalled = true
	return d.DhcpConfPathResult
}

func (d *DriverMock) VmnetnatConfPath(device string) string {
	d.VmnetnatConfPathCalled = true
	return d.VmnetnatConfPathResult
}

func (d *DriverMock) NetmapConfPath() string {
	d.NetmapConfPathCalled = true
	return d.NetmapConfPathResult
}

func (d *DriverMock) Verify() error {
	d.VerifyCalled = true
	return d.VerifyErr
}

func (d *DriverMock) GetVmwareDriver() VmwareDriver {
	var state VmwareDriver
	state.DhcpLeasesPath = func(string) string {
		return "/path/to/dhcp.leases"
	}
	state.DhcpConfPath = func(string) string {
		return "/path/to/dhcp.conf"
	}
	state.VmnetnatConfPath = func(string) string {
		return "/path/to/vmnetnat.conf"
	}
	state.NetmapConfPath = func() string {
		return "/path/to/netmap.conf"
	}
	return state
}
<|endoftext|>


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

package tls

import (
	"crypto/x509"
	"reflect"
	"syscall"
	"unsafe"
)

func loadStore(roots *x509.CertPool, name string) {
	store, errno := syscall.CertOpenSystemStore(syscall.InvalidHandle, syscall.StringToUTF16Ptr(name))
	if errno != 0 {
		return
	}

	var prev *syscall.CertContext
	for {
		cur := syscall.CertEnumCertificatesInStore(store, prev)
		if cur == nil {
			break
		}

		var buf []byte
		hdrp := (*reflect.SliceHeader)(unsafe.Pointer(&buf))
		hdrp.Data = cur.EncodedCert
		hdrp.Len = int(cur.Length)
		hdrp.Cap = int(cur.Length)

		cert, err := x509.ParseCertificate(buf)
		if err != nil {
			continue
		}

		roots.AddCert(cert)
		prev = cur
	}

	syscall.CertCloseStore(store, 0)
}

func initDefaultRoots() {
	roots := x509.NewCertPool()

	// Roots
	loadStore(roots, "ROOT")

	// Intermediates
	loadStore(roots, "CA")

	varDefaultRoots = roots
}
<commit_msg>crypto/tls: disable root cert fetching to fix windows build<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tls

import (
	"crypto/x509"
	"reflect"
	"syscall"
	"unsafe"
)

func loadStore(roots *x509.CertPool, name string) {
	store, errno := syscall.CertOpenSystemStore(syscall.InvalidHandle, syscall.StringToUTF16Ptr(name))
	if errno != 0 {
		return
	}

	var prev *syscall.CertContext
	for {
		cur := syscall.CertEnumCertificatesInStore(store, prev)
		if cur == nil {
			break
		}

		var buf []byte
		hdrp := (*reflect.SliceHeader)(unsafe.Pointer(&buf))
		hdrp.Data = cur.EncodedCert
		hdrp.Len = int(cur.Length)
		hdrp.Cap = int(cur.Length)

		cert, err := x509.ParseCertificate(buf)
		if err != nil {
			continue
		}

		roots.AddCert(cert)
		prev = cur
	}

	syscall.CertCloseStore(store, 0)
}

func initDefaultRoots() {
	// TODO(brainman): To be fixed
	return

	roots := x509.NewCertPool()

	// Roots
	loadStore(roots, "ROOT")

	// Intermediates
	loadStore(roots, "CA")

	varDefaultRoots = roots
}
<|endoftext|>


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

package state_test

import (
	"github.com/jacobsa/comeback/state"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestScoreMap(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type ScoreMapTest struct {
	m state.ScoreMap
}

func init() { RegisterTestSuite(&ScoreMapTest{}) }

func (t *ScoreMapTest) SetUp(i *TestInfo) {
	t.m = state.NewScoreMap()
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *ScoreMapTest) DoesFoo() {
	ExpectEq("TODO", "")
}
<commit_msg>Added test names.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package state_test

import (
	"github.com/jacobsa/comeback/state"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestScoreMap(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type ScoreMapTest struct {
	m state.ScoreMap
}

func init() { RegisterTestSuite(&ScoreMapTest{}) }

func (t *ScoreMapTest) SetUp(i *TestInfo) {
	t.m = state.NewScoreMap()
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *ScoreMapTest) EmptyMap() {
	ExpectEq("TODO", "")
}

func (t *ScoreMapTest) SomeElements() {
	ExpectEq("TODO", "")
}

func (t *ScoreMapTest) AddTwice() {
	ExpectEq("TODO", "")
}

func (t *ScoreMapTest) GobRoundTrip() {
	ExpectEq("TODO", "")
}

func (t *ScoreMapTest) DecodingOverwritesContents() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>package version

const (
	Version = "0.1.3"
)
<commit_msg>bump(version): v0.1.3+git<commit_after>package version

const (
	Version = "0.1.3+git"
)
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"net"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"text/template"
	"time"
)

// The Simple Push server version.
const VERSION = "1.5.0"

var (
	ErrMissingOrigin = errors.New("Missing WebSocket origin")
	ErrInvalidOrigin = errors.New("WebSocket origin not allowed")
)

type ApplicationConfig struct {
	Hostname           string `toml:"current_host" env:"current_host"`
	TokenKey           string `toml:"token_key" env:"token_key"`
	PushEndpoint       string `toml:"push_endpoint_template" env:"push_endpoint_template"`
	UseAwsHost         bool   `toml:"use_aws_host" env:"use_aws_host"`
	ResolveHost        bool   `toml:"resolve_host" env:"resolve_host"`
	ClientMinPing      string `toml:"client_min_ping_interval" env:"client_min_ping_interval"`
	ClientHelloTimeout string `toml:"client_hello_timeout" env:"client_hello_timeout"`
	PushLongPongs      bool   `toml:"push_long_pongs" env:"push_long_pongs"`
	ClientPongInterval string `toml:"client_pong_interval" env:"client_pong_interval"`
}

func NewApplication() (a *Application) {
	a = &Application{
		workers:   make(map[string]Worker),
		closeChan: make(chan bool),
	}
	return a
}

type Application struct {
	info               InstanceInfo
	hostname           string
	host               string
	port               int
	clientMinPing      time.Duration
	clientHelloTimeout time.Duration
	clientPongInterval time.Duration
	pushLongPongs      bool
	tokenKey           []byte
	endpointTemplate   *template.Template
	log                *SimpleLogger
	metrics            Statistician
	workers            map[string]Worker
	workerMux          sync.RWMutex
	workerCount        int32
	store              Store
	router             Router
	locator            Locator
	balancer           Balancer
	sh                 Handler // WebSocket handler.
	eh                 Handler // HTTP update handler.
	ph                 Handler // Performance profiling handlers.
	propping           PropPinger
	closeChan          chan bool
	closeOnce          Once
}

func (a *Application) ConfigStruct() interface{} {
	defaultHost, _ := os.Hostname()
	return &ApplicationConfig{
		Hostname:           defaultHost,
		PushEndpoint:       "{{.CurrentHost}}/update/{{.Token}}",
		UseAwsHost:         false,
		ResolveHost:        false,
		ClientMinPing:      "20s",
		ClientHelloTimeout: "30s",
		ClientPongInterval: "5m",
	}
}

// Fully initialize the application, this initializes all the other components
// as well.
// Note: We implement the Init method to comply with the interface, so the app
// passed here will be nil.
func (a *Application) Init(_ *Application, config interface{}) (err error) {
	conf := config.(*ApplicationConfig)

	if conf.UseAwsHost {
		a.info = new(EC2Info)
	} else if conf.ResolveHost {
		addr, err := net.ResolveIPAddr("ip", conf.Hostname)
		if err != nil {
			return fmt.Errorf("Error resolving hostname: %s", err)
		}
		a.info = LocalInfo{addr.String()}
	} else {
		a.info = LocalInfo{conf.Hostname}
	}
	if a.hostname, err = a.info.PublicHostname(); err != nil {
		return fmt.Errorf("Error determining hostname: %s", err)
	}

	if err = a.SetTokenKey(conf.TokenKey); err != nil {
		return fmt.Errorf("Malformed token key: %s", err)
	}
	if a.endpointTemplate, err = template.New("Push").Parse(conf.PushEndpoint); err != nil {
		return fmt.Errorf("Error parsing push endpoint template: %s", err)
	}

	if a.clientMinPing, err = time.ParseDuration(conf.ClientMinPing); err != nil {
		return fmt.Errorf("Unable to parse 'client_min_ping_interval': %s",
			err.Error())
	}
	if a.clientPongInterval, err = time.ParseDuration(conf.ClientPongInterval); err != nil {
		return fmt.Errorf("Unable to parse 'client_pong_interval': %s",
			err.Error())
	}
	if a.clientHelloTimeout, err = time.ParseDuration(conf.ClientHelloTimeout); err != nil {
		return fmt.Errorf("Unable to parse 'client_hello_timeout': %s",
			err.Error())
	}
	a.pushLongPongs = conf.PushLongPongs
	return
}

// Set a logger
func (a *Application) SetLogger(logger Logger) (err error) {
	a.log, err = NewLogger(logger)
	return
}

func (a *Application) SetPropPinger(ping PropPinger) (err error) {
	a.propping = ping
	return
}

func (a *Application) SetMetrics(metrics Statistician) error {
	a.metrics = metrics
	return nil
}

func (a *Application) SetStore(store Store) error {
	a.store = store
	return nil
}

func (a *Application) SetRouter(router Router) error {
	a.router = router
	return nil
}

func (a *Application) SetLocator(locator Locator) error {
	a.locator = locator
	return nil
}

func (a *Application) SetBalancer(b Balancer) error {
	a.balancer = b
	return nil
}

func (a *Application) SetSocketHandler(h Handler) error {
	a.sh = h
	return nil
}

func (a *Application) SetEndpointHandler(h Handler) error {
	a.eh = h
	return nil
}

func (a *Application) SetProfileHandlers(h Handler) error {
	a.ph = h
	return nil
}

// Start the application
func (a *Application) Run() (errChan chan error) {
	errChan = make(chan error, 4)

	go a.sh.Start(errChan)
	go a.eh.Start(errChan)
	go a.router.Start(errChan)
	go a.ph.Start(errChan)

	go a.sendClientCount()
	return errChan
}

func (a *Application) Hostname() string {
	return a.hostname
}

func (a *Application) InstanceInfo() InstanceInfo {
	return a.info
}

func (a *Application) Logger() *SimpleLogger {
	return a.log
}

//TODO: move these to handler so we can deal with multiple prop.ping formats
func (a *Application) PropPinger() PropPinger {
	return a.propping
}

func (a *Application) Store() Store {
	return a.store
}

func (a *Application) Metrics() Statistician {
	return a.metrics
}

func (a *Application) Router() Router {
	return a.router
}

func (a *Application) Locator() Locator {
	return a.locator
}

func (a *Application) Balancer() Balancer {
	return a.balancer
}

func (a *Application) SocketHandler() Handler {
	return a.sh
}

func (a *Application) EndpointHandler() Handler {
	return a.eh
}

func (a *Application) ProfileHandlers() Handler {
	return a.ph
}

func (a *Application) TokenKey() []byte {
	return a.tokenKey
}

func (a *Application) SetTokenKey(key string) (err error) {
	if len(key) == 0 {
		a.tokenKey = nil
	} else {
		a.tokenKey, err = base64.URLEncoding.DecodeString(key)
	}
	return
}

func (a *Application) WorkerCount() (count int) {
	return int(atomic.LoadInt32(&a.workerCount))
}

func (a *Application) WorkerExists(uaid string) (collision bool) {
	_, collision = a.GetWorker(uaid)
	return
}

func (a *Application) GetWorker(uaid string) (worker Worker, ok bool) {
	a.workerMux.RLock()
	worker, ok = a.workers[uaid]
	a.workerMux.RUnlock()
	return
}

func (a *Application) AddWorker(uaid string, worker Worker) (replaced bool) {
	if a.closeOnce.IsDone() {
		worker.Close()
		return
	}
	a.workerMux.Lock()
	// Avoid incrementing the worker count for duplicate handshakes. Callers
	// can use this to short-circuit other operations (e.g., re-registering
	// with the router).
	_, replaced = a.workers[uaid]
	a.workers[uaid] = worker
	a.workerMux.Unlock()
	if !replaced {
		atomic.AddInt32(&a.workerCount, 1)
	}
	return
}

func (a *Application) RemoveWorker(uaid string, worker Worker) (removed bool) {
	if a.closeOnce.IsDone() {
		return
	}
	a.workerMux.Lock()
	if prevWorker, ok := a.workers[uaid]; ok && prevWorker == worker {
		delete(a.workers, uaid)
		removed = true
	}
	a.workerMux.Unlock()
	if removed {
		atomic.AddInt32(&a.workerCount, -1)
	}
	return removed
}

func (a *Application) closeWorkers() {
	a.workerMux.Lock()
	defer a.workerMux.Unlock()
	for uaid, worker := range a.workers {
		delete(a.workers, uaid)
		worker.Close()
	}
}

// CreateEndpoint allocates an update endpoint with the given primary key.
func (a *Application) CreateEndpoint(key string) (string, error) {
	token, err := a.encodePK(key)
	if err != nil {
		return "", err
	}
	return a.genEndpoint(token)
}

// encodePK encodes a primary key if a token key is specified.
func (a *Application) encodePK(key string) (token string, err error) {
	tokenKey := a.TokenKey()
	if len(tokenKey) == 0 {
		return key, nil
	}
	btoken := []byte(key)
	return Encode(tokenKey, btoken)
}

// genEndpoint generates an update endpoint.
func (a *Application) genEndpoint(token string) (string, error) {
	var currentHost string
	if eh := a.EndpointHandler(); eh != nil {
		currentHost = eh.URL()
	}
	// cheezy variable replacement.
	endpoint := new(bytes.Buffer)
	if err := a.endpointTemplate.Execute(endpoint, struct {
		Token       string
		CurrentHost string
	}{
		token,
		currentHost,
	}); err != nil {
		return "", err
	}
	return endpoint.String(), nil
}

func (a *Application) Close() error {
	return a.closeOnce.Do(a.close)
}

func (a *Application) close() error {
	var errors MultipleError
	if eh := a.EndpointHandler(); eh != nil {
		// Stop the update listener; close all connections.
		if err := eh.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if b := a.Balancer(); b != nil {
		// Deregister from the balancer.
		if err := b.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if sh := a.SocketHandler(); sh != nil {
		// Close the WebSocket listener.
		if err := sh.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	// Disconnect existing clients.
	a.closeWorkers()
	// Stop publishing client counts.
	close(a.closeChan)
	if l := a.Locator(); l != nil {
		// Deregister from the discovery service.
		if err := a.locator.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if r := a.Router(); r != nil {
		// Close the routing listener.
		if err := r.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if ph := a.ProfileHandlers(); ph != nil {
		// Stop the profiling listener.
		if err := ph.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if len(errors) > 0 {
		return errors
	}
	return nil
}

func (a *Application) sendClientCount() {
	metrics := a.Metrics()
	ticker := time.NewTicker(1 * time.Second)
	for ok := true; ok; {
		select {
		case ok = <-a.closeChan:
		case <-ticker.C:
			metrics.Gauge("goroutines", int64(runtime.NumGoroutine()))
			metrics.Gauge("update.client.connections", int64(a.WorkerCount()))
		}
	}
	ticker.Stop()
}
<commit_msg>Disable server-sent pongs by default. Closes #216.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"net"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"text/template"
	"time"
)

// The Simple Push server version.
const VERSION = "1.5.0"

var (
	ErrMissingOrigin = errors.New("Missing WebSocket origin")
	ErrInvalidOrigin = errors.New("WebSocket origin not allowed")
)

type ApplicationConfig struct {
	Hostname           string `toml:"current_host" env:"current_host"`
	TokenKey           string `toml:"token_key" env:"token_key"`
	PushEndpoint       string `toml:"push_endpoint_template" env:"push_endpoint_template"`
	UseAwsHost         bool   `toml:"use_aws_host" env:"use_aws_host"`
	ResolveHost        bool   `toml:"resolve_host" env:"resolve_host"`
	ClientMinPing      string `toml:"client_min_ping_interval" env:"client_min_ping_interval"`
	ClientHelloTimeout string `toml:"client_hello_timeout" env:"client_hello_timeout"`
	PushLongPongs      bool   `toml:"push_long_pongs" env:"push_long_pongs"`
	ClientPongInterval string `toml:"client_pong_interval" env:"client_pong_interval"`
}

func NewApplication() (a *Application) {
	a = &Application{
		workers:   make(map[string]Worker),
		closeChan: make(chan bool),
	}
	return a
}

type Application struct {
	info               InstanceInfo
	hostname           string
	host               string
	port               int
	clientMinPing      time.Duration
	clientHelloTimeout time.Duration
	clientPongInterval time.Duration
	pushLongPongs      bool
	tokenKey           []byte
	endpointTemplate   *template.Template
	log                *SimpleLogger
	metrics            Statistician
	workers            map[string]Worker
	workerMux          sync.RWMutex
	workerCount        int32
	store              Store
	router             Router
	locator            Locator
	balancer           Balancer
	sh                 Handler // WebSocket handler.
	eh                 Handler // HTTP update handler.
	ph                 Handler // Performance profiling handlers.
	propping           PropPinger
	closeChan          chan bool
	closeOnce          Once
}

func (a *Application) ConfigStruct() interface{} {
	defaultHost, _ := os.Hostname()
	return &ApplicationConfig{
		Hostname:           defaultHost,
		PushEndpoint:       "{{.CurrentHost}}/update/{{.Token}}",
		UseAwsHost:         false,
		ResolveHost:        false,
		ClientMinPing:      "20s",
		ClientHelloTimeout: "30s",
	}
}

// Fully initialize the application, this initializes all the other components
// as well.
// Note: We implement the Init method to comply with the interface, so the app
// passed here will be nil.
func (a *Application) Init(_ *Application, config interface{}) (err error) {
	conf := config.(*ApplicationConfig)

	if conf.UseAwsHost {
		a.info = new(EC2Info)
	} else if conf.ResolveHost {
		addr, err := net.ResolveIPAddr("ip", conf.Hostname)
		if err != nil {
			return fmt.Errorf("Error resolving hostname: %s", err)
		}
		a.info = LocalInfo{addr.String()}
	} else {
		a.info = LocalInfo{conf.Hostname}
	}
	if a.hostname, err = a.info.PublicHostname(); err != nil {
		return fmt.Errorf("Error determining hostname: %s", err)
	}

	if err = a.SetTokenKey(conf.TokenKey); err != nil {
		return fmt.Errorf("Malformed token key: %s", err)
	}
	if a.endpointTemplate, err = template.New("Push").Parse(conf.PushEndpoint); err != nil {
		return fmt.Errorf("Error parsing push endpoint template: %s", err)
	}

	if a.clientMinPing, err = time.ParseDuration(conf.ClientMinPing); err != nil {
		return fmt.Errorf("Unable to parse 'client_min_ping_interval': %s",
			err.Error())
	}
	if len(conf.ClientPongInterval) > 0 {
		if a.clientPongInterval, err = time.ParseDuration(conf.ClientPongInterval); err != nil {
			return fmt.Errorf("Unable to parse 'client_pong_interval': %s",
				err.Error())
		}
	}
	if a.clientHelloTimeout, err = time.ParseDuration(conf.ClientHelloTimeout); err != nil {
		return fmt.Errorf("Unable to parse 'client_hello_timeout': %s",
			err.Error())
	}
	a.pushLongPongs = conf.PushLongPongs
	return
}

// Set a logger
func (a *Application) SetLogger(logger Logger) (err error) {
	a.log, err = NewLogger(logger)
	return
}

func (a *Application) SetPropPinger(ping PropPinger) (err error) {
	a.propping = ping
	return
}

func (a *Application) SetMetrics(metrics Statistician) error {
	a.metrics = metrics
	return nil
}

func (a *Application) SetStore(store Store) error {
	a.store = store
	return nil
}

func (a *Application) SetRouter(router Router) error {
	a.router = router
	return nil
}

func (a *Application) SetLocator(locator Locator) error {
	a.locator = locator
	return nil
}

func (a *Application) SetBalancer(b Balancer) error {
	a.balancer = b
	return nil
}

func (a *Application) SetSocketHandler(h Handler) error {
	a.sh = h
	return nil
}

func (a *Application) SetEndpointHandler(h Handler) error {
	a.eh = h
	return nil
}

func (a *Application) SetProfileHandlers(h Handler) error {
	a.ph = h
	return nil
}

// Start the application
func (a *Application) Run() (errChan chan error) {
	errChan = make(chan error, 4)

	go a.sh.Start(errChan)
	go a.eh.Start(errChan)
	go a.router.Start(errChan)
	go a.ph.Start(errChan)

	go a.sendClientCount()
	return errChan
}

func (a *Application) Hostname() string {
	return a.hostname
}

func (a *Application) InstanceInfo() InstanceInfo {
	return a.info
}

func (a *Application) Logger() *SimpleLogger {
	return a.log
}

//TODO: move these to handler so we can deal with multiple prop.ping formats
func (a *Application) PropPinger() PropPinger {
	return a.propping
}

func (a *Application) Store() Store {
	return a.store
}

func (a *Application) Metrics() Statistician {
	return a.metrics
}

func (a *Application) Router() Router {
	return a.router
}

func (a *Application) Locator() Locator {
	return a.locator
}

func (a *Application) Balancer() Balancer {
	return a.balancer
}

func (a *Application) SocketHandler() Handler {
	return a.sh
}

func (a *Application) EndpointHandler() Handler {
	return a.eh
}

func (a *Application) ProfileHandlers() Handler {
	return a.ph
}

func (a *Application) TokenKey() []byte {
	return a.tokenKey
}

func (a *Application) SetTokenKey(key string) (err error) {
	if len(key) == 0 {
		a.tokenKey = nil
	} else {
		a.tokenKey, err = base64.URLEncoding.DecodeString(key)
	}
	return
}

func (a *Application) WorkerCount() (count int) {
	return int(atomic.LoadInt32(&a.workerCount))
}

func (a *Application) WorkerExists(uaid string) (collision bool) {
	_, collision = a.GetWorker(uaid)
	return
}

func (a *Application) GetWorker(uaid string) (worker Worker, ok bool) {
	a.workerMux.RLock()
	worker, ok = a.workers[uaid]
	a.workerMux.RUnlock()
	return
}

func (a *Application) AddWorker(uaid string, worker Worker) (replaced bool) {
	if a.closeOnce.IsDone() {
		worker.Close()
		return
	}
	a.workerMux.Lock()
	// Avoid incrementing the worker count for duplicate handshakes. Callers
	// can use this to short-circuit other operations (e.g., re-registering
	// with the router).
	_, replaced = a.workers[uaid]
	a.workers[uaid] = worker
	a.workerMux.Unlock()
	if !replaced {
		atomic.AddInt32(&a.workerCount, 1)
	}
	return
}

func (a *Application) RemoveWorker(uaid string, worker Worker) (removed bool) {
	if a.closeOnce.IsDone() {
		return
	}
	a.workerMux.Lock()
	if prevWorker, ok := a.workers[uaid]; ok && prevWorker == worker {
		delete(a.workers, uaid)
		removed = true
	}
	a.workerMux.Unlock()
	if removed {
		atomic.AddInt32(&a.workerCount, -1)
	}
	return removed
}

func (a *Application) closeWorkers() {
	a.workerMux.Lock()
	defer a.workerMux.Unlock()
	for uaid, worker := range a.workers {
		delete(a.workers, uaid)
		worker.Close()
	}
}

// CreateEndpoint allocates an update endpoint with the given primary key.
func (a *Application) CreateEndpoint(key string) (string, error) {
	token, err := a.encodePK(key)
	if err != nil {
		return "", err
	}
	return a.genEndpoint(token)
}

// encodePK encodes a primary key if a token key is specified.
func (a *Application) encodePK(key string) (token string, err error) {
	tokenKey := a.TokenKey()
	if len(tokenKey) == 0 {
		return key, nil
	}
	btoken := []byte(key)
	return Encode(tokenKey, btoken)
}

// genEndpoint generates an update endpoint.
func (a *Application) genEndpoint(token string) (string, error) {
	var currentHost string
	if eh := a.EndpointHandler(); eh != nil {
		currentHost = eh.URL()
	}
	// cheezy variable replacement.
	endpoint := new(bytes.Buffer)
	if err := a.endpointTemplate.Execute(endpoint, struct {
		Token       string
		CurrentHost string
	}{
		token,
		currentHost,
	}); err != nil {
		return "", err
	}
	return endpoint.String(), nil
}

func (a *Application) Close() error {
	return a.closeOnce.Do(a.close)
}

func (a *Application) close() error {
	var errors MultipleError
	if eh := a.EndpointHandler(); eh != nil {
		// Stop the update listener; close all connections.
		if err := eh.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if b := a.Balancer(); b != nil {
		// Deregister from the balancer.
		if err := b.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if sh := a.SocketHandler(); sh != nil {
		// Close the WebSocket listener.
		if err := sh.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	// Disconnect existing clients.
	a.closeWorkers()
	// Stop publishing client counts.
	close(a.closeChan)
	if l := a.Locator(); l != nil {
		// Deregister from the discovery service.
		if err := a.locator.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if r := a.Router(); r != nil {
		// Close the routing listener.
		if err := r.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if ph := a.ProfileHandlers(); ph != nil {
		// Stop the profiling listener.
		if err := ph.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if len(errors) > 0 {
		return errors
	}
	return nil
}

func (a *Application) sendClientCount() {
	metrics := a.Metrics()
	ticker := time.NewTicker(1 * time.Second)
	for ok := true; ok; {
		select {
		case ok = <-a.closeChan:
		case <-ticker.C:
			metrics.Gauge("goroutines", int64(runtime.NumGoroutine()))
			metrics.Gauge("update.client.connections", int64(a.WorkerCount()))
		}
	}
	ticker.Stop()
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	"time"
)

// User maps to users table
type User struct {
	id             string
	Email          string
	HashedPassword string
	IsAdmin        bool
	DateCreated    time.Time
}

// ID returns read-only Primary Key ID of User
func (user *User) ID() string {
	return user.id
}

// IsTransient determines if User record has been saved to the database,
// true means User struct has NOT been saved, false means it has.
func (user *User) IsTransient() bool {
	return len(user.id) == 0
}

// UserSave saves the User struct to the database.
var UserSave = func(user *User) error {
	if user.IsTransient() {
		cmd := `INSERT INTO users(email, hashed_password, is_admin, date_created)
				VALUES($1, $2, $3, $4)
				RETURNING id`

		statement, err := db.Prepare(cmd)
		if err != nil {
			return err
		}
		defer statement.Close()

		err = statement.
			QueryRow(user.Email, user.HashedPassword, user.IsAdmin, user.DateCreated).
			Scan(&user.id)
		if err != nil {
			return err
		}
	} else {
		cmd := `UPDATE users
				SET email = $2, hashed_password = $3 is_admin = $4, date_created = $5
				WHERE id = $1`

		_, err := db.Exec(cmd, user.id, user.HashedPassword, user.Email, user.IsAdmin, user.DateCreated)
		if err != nil {
			return err
		}
	}

	return nil
}

// Save saves the User struct to the database.
func (user *User) Save() error {
	return UserSave(user)
}

// UserDelete deletes the User from the database
var UserDelete = func(user *User) error {
	cmd := `DELETE FROM users
			WHERE id = $1`

	_, err := db.Exec(cmd, user.id)
	return err
}

// Delete deletes the User from the database
func (user *User) Delete() error {
	return UserDelete(user)
}

// UserFromID returns a User record with given ID
var UserFromID = func(id string) (User, error) {
	var user User

	if !isUUID.MatchString(id) {
		return user, ErrEntityNotFound
	}

	cmd := `SELECT email, hashed_password, is_admin, date_created
			FROM users
			WHERE id = $1`

	err := db.QueryRow(cmd, id).
		Scan(&user.Email, &user.HashedPassword, &user.IsAdmin, &user.DateCreated)
	if err == sql.ErrNoRows {
		return user, ErrEntityNotFound
	} else if err != nil {
		return user, err
	}

	user.id = id
	return user, nil
}

// UserFromEmail returns the User record matching an email address
var UserFromEmail = func(email string) (User, error) {
	var user User

	cmd := `SELECT id, email, hashed_password, is_admin, date_created
			FROM users
			WHERE email = $1`

	err := db.QueryRow(cmd, email).
		Scan(&user.id, &user.Email, &user.HashedPassword, &user.IsAdmin, &user.DateCreated)
	if err == sql.ErrNoRows {
		return user, ErrEntityNotFound
	} else if err != nil {
		return user, err
	}

	return user, nil
}

const (
	// UsersOrderByDateCreated is for ordering users by DateCreated
	UsersOrderByDateCreated = "date_created"
	// UsersOrderByEmail is for ordering users by Email address
	UsersOrderByEmail = "email"
)

// UsersAll returns all User records from the database
var UsersAll = func(query QueryAll) ([]User, error) {
	var users []User

	cmd := `SELECT id, email, hashed_password, is_admin, date_created
			FROM users
			ORDER BY $1
			LIMIT $2`

	orderBy := query.OrderBy
	if query.OrderAsc {
		orderBy += " ASC"
	} else {
		orderBy += " DESC"
	}

	rows, err := db.Query(cmd, orderBy, query.Limit)
	if err != nil {
		return nil, err
	}

	defer rows.Close()

	for rows.Next() {
		user := User{}
		err = rows.Scan(&user.id, &user.Email, &user.HashedPassword, &user.IsAdmin, &user.DateCreated)
		if err != nil {
			return nil, err
		}

		users = append(users, user)
	}

	return users, nil
}
<commit_msg>fixed: failing unit test.<commit_after>package db

import (
	"database/sql"
	"time"
)

// User maps to users table
type User struct {
	id             string
	Email          string
	HashedPassword string
	IsAdmin        bool
	DateCreated    time.Time
}

// ID returns read-only Primary Key ID of User
func (user *User) ID() string {
	return user.id
}

// IsTransient determines if User record has been saved to the database,
// true means User struct has NOT been saved, false means it has.
func (user *User) IsTransient() bool {
	return len(user.id) == 0
}

// UserSave saves the User struct to the database.
var UserSave = func(user *User) error {
	if user.IsTransient() {
		cmd := `INSERT INTO users(email, hashed_password, is_admin, date_created)
				VALUES($1, $2, $3, $4)
				RETURNING id`

		statement, err := db.Prepare(cmd)
		if err != nil {
			return err
		}
		defer statement.Close()

		err = statement.
			QueryRow(user.Email, user.HashedPassword, user.IsAdmin, user.DateCreated).
			Scan(&user.id)
		if err != nil {
			return err
		}
	} else {
		cmd := `UPDATE users
				SET email = $2, hashed_password = $3 is_admin = $4, date_created = $5
				WHERE id = $1`

		_, err := db.Exec(cmd, user.id, user.HashedPassword, user.Email, user.IsAdmin, user.DateCreated)
		if err != nil {
			return err
		}
	}

	return nil
}

// Save saves the User struct to the database.
func (user *User) Save() error {
	return UserSave(user)
}

// UserDelete deletes the User from the database
var UserDelete = func(user *User) error {
	cmd := `DELETE FROM users
			WHERE id = $1`

	_, err := db.Exec(cmd, user.id)
	return err
}

// Delete deletes the User from the database
func (user *User) Delete() error {
	return UserDelete(user)
}

// UserFromID returns a User record with given ID
var UserFromID = func(id string) (User, error) {
	var user User

	if !isUUID.MatchString(id) {
		return user, ErrEntityNotFound
	}

	cmd := `SELECT email, hashed_password, is_admin, date_created
			FROM users
			WHERE id = $1`

	err := db.QueryRow(cmd, id).
		Scan(&user.Email, &user.HashedPassword, &user.IsAdmin, &user.DateCreated)
	if err == sql.ErrNoRows {
		return user, ErrEntityNotFound
	} else if err != nil {
		return user, err
	}

	user.id = id
	return user, nil
}

// UserFromEmail returns the User record matching an email address
var UserFromEmail = func(email string) (User, error) {
	var user User

	cmd := `SELECT id, email, hashed_password, is_admin, date_created
			FROM users
			WHERE email = $1`

	err := db.QueryRow(cmd, email).
		Scan(&user.id, &user.Email, &user.HashedPassword, &user.IsAdmin, &user.DateCreated)
	if err == sql.ErrNoRows {
		return user, ErrEntityNotFound
	} else if err != nil {
		return user, err
	}

	return user, nil
}

const (
	// UsersOrderByDateCreated is for ordering users by DateCreated
	UsersOrderByDateCreated = "date_created"
	// UsersOrderByEmail is for ordering users by Email address
	UsersOrderByEmail = "email"
)

// UsersAll returns all User records from the database
var UsersAll = func(query QueryAll) ([]User, error) {
	var users []User

	orderBy := query.OrderBy
	if query.OrderAsc {
		orderBy += " ASC"
	} else {
		orderBy += " DESC"
	}

	var rows *sql.Rows
	var err error

	if after, ok := query.After.(string); ok && query.OrderBy == UsersOrderByEmail {
		cmd := `SELECT id, email, hashed_password, is_admin, date_created
			    FROM users
				WHERE id > (SELECT id FROM users WHERE email = $1)
				ORDER BY $2
			    LIMIT $3`

		rows, err = db.Query(cmd, after, orderBy, query.Limit)
	} else if after, ok := query.After.(time.Time); ok && query.OrderBy == UsersOrderByDateCreated {
		cmd := `SELECT id, email, hashed_password, is_admin, date_created
			    FROM users
				WHERE date_created > $1
				ORDER BY $2
			    LIMIT $3`

		rows, err = db.Query(cmd, after, orderBy, query.Limit)
	} else {
		cmd := `SELECT id, email, hashed_password, is_admin, date_created
			    FROM users
			    ORDER BY $1
			    LIMIT $2`

		rows, err = db.Query(cmd, orderBy, query.Limit)
	}

	if err != nil {
		return nil, err
	}

	defer rows.Close()

	for rows.Next() {
		user := User{}
		err := rows.Scan(&user.id, &user.Email, &user.HashedPassword, &user.IsAdmin, &user.DateCreated)
		if err != nil {
			return nil, err
		}

		users = append(users, user)
	}

	return users, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google Inc. All rights reserved.

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

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

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

package websockets

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"path"
	"strings"
	"sync"
	"sync/atomic"
	"text/template"

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

const (
	contentTypeHeader = "Content-Type"
	shimTemplate      = `
<!--START_WEBSOCKET_SHIM-->
<!--
    This file was served from behind a reverse proxy that does not support websockets.

    The following code snippet has been inserted by the proxy to replace websockets
    with socket.io which is based on HTTP and will work with this proxy.

    If this snippet insertion is causing issues, then contact the server administrator.
-->
<script>
(function() {
    if (typeof window.nativeWebSocket !== 'undefined') {
      // We have already replaced websockets
      return;
    }
    console.log('Replacing native websockets with a shim');
    window.nativeWebSocket = window.WebSocket;
    const  location = window.location;
    const shimUri = location.protocol + '//' + location.host + '/{{.ShimPath}}/';

    function shouldShimWebsockets(url) {
      var parsedURL = new URL(url);
      if (typeof parsedURL.host == 'undefined') {
        parsedURL.host = location.host;
      }
      return (parsedURL.host == location.host);
    }

    function WebSocketShim(url, protocols) {
      if (!shouldShimWebsockets(url)) {
        console.log("Not shimming websockets for " + parsedURL.host + " as it does not match the location of " + location.host);
        return new window.nativeWebSocket(url, protocols);
      }

      // We need to reference "this" within nested functions, so we alias it to "self"
      var self = this;

      this.readyState = WebSocketShim.CONNECTING;
      function openedHandler(msg) {
        self.readyState = WebSocketShim.OPEN;
        if (self.onopen) {
          self.onopen({ target: self });
        }
      }
      function receiveHandler(msg) {
        if (self.onmessage) {
          self.onmessage({ target: self, data: msg });
        }
      }
      function errorHandler() {
        if (self.onerror) {
          self.onerror({ target: self });
        }
      }
      self.xhr = function(action, msg, onsuccess, onexit) {
        var req = new XMLHttpRequest();
        req.onreadystatechange = function() {
          if (req.readyState === 4) {
            if (req.status === 200) {
              if (onsuccess) {
                onsuccess(req.responseText);
              }
            } else if (req.status !== 408) {
              errorHandler();
            }
            if (onexit) {
              onexit();
            }
          }
        };
        req.open("POST", shimUri + action, true);
        if (typeof msg !== 'string') {
          msg = JSON.stringify(msg);
        }
        req.send(msg);
      }

      self.closedHandler = function() {
        self.readyState = WebSocketShim.CLOSED;
        if (self.onclose) {
          self.onclose({ target: self });
        }
      }

      function poll() {
        if (self.readyState != WebSocketShim.OPEN) {
          return;
        }
        self.xhr('poll', {'id': self._sessionID}, receiveHandler, poll);
      }

      self.xhr('open', url, function(resp) {
        respJSON = JSON.parse(resp);
        self._sessionID = respJSON.id;
        openedHandler(respJSON.msg);
        poll();
      });
    }
    WebSocketShim.prototype = {
      binaryType: "blob",
      onopen: null,
      onclose: null,
      onmessage: null,
      onerror: null,

      send: function(data) {
        if (this.readyState != WebSocketShim.OPEN) {
          throw new Error('WebSocket is not yet opened');
        }
        this.xhr('data', {'id': this._sessionID, 'msg': data});
      },
      close: function() {
        if (this.readyState != WebSocketShim.OPEN) {
          return;
        }
        this.readyState = WebSocketShim.CLOSING;
        this.xhr('close', {'id': this._sessionID}, false, this.closedHandler);
      },
    };
    WebSocketShim.CONNECTING = 0;
    WebSocketShim.OPEN = 1;
    WebSocketShim.CLOSING = 2;
    WebSocketShim.CLOSED = 3;

    window.WebSocket = WebSocketShim;
})();
</script>
<!--END_WEBSOCKET_SHIM-->
`
)

var shimTmpl = template.Must(template.New("client-shim").Parse(shimTemplate))

type shimmedBody struct {
	reader io.Reader
	closer io.Closer
}

func (sb *shimmedBody) Read(p []byte) (n int, err error) {
	return sb.reader.Read(p)
}

func (sb *shimmedBody) Close() error {
	return sb.closer.Close()
}

func shimBody(resp *http.Response, shimCode string) error {
	if resp == nil || resp.Body == nil {
		// We have nothing to do on an empty response
		return nil
	}
	contentType := strings.ToLower(resp.Header.Get(contentTypeHeader))
	if !strings.Contains(contentType, "html") {
		// We only want to modify HTML responses
		return nil
	}
	wrapped := resp.Body

	// Read in the first kilobyte to see if the <head> tag exists in it
	buf := make([]byte, 1024)
	count, err := wrapped.Read(buf)
	if err != nil && err != io.EOF {
		return err
	}
	prefix := strings.Replace(string(buf[0:count]), "<head>", "<head>"+shimCode, 1)
	resp.Body = &shimmedBody{
		reader: io.MultiReader(strings.NewReader(prefix), wrapped),
		closer: wrapped,
	}
	return nil
}

type sessionMessage struct {
	ID      string `json:"id,omitempty"`
	Message string `json:"msg,omitempty"`
}

func createShimChannel(ctx context.Context, host, shimPath string) http.Handler {
	var connections sync.Map
	var sessionCount uint64
	mux := http.NewServeMux()
	mux.HandleFunc(path.Join(shimPath, "open"), func(w http.ResponseWriter, r *http.Request) {
		sessionID := fmt.Sprintf("%d", atomic.AddUint64(&sessionCount, 1))
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, fmt.Sprintf("internal error reading a shim request: %v", err), http.StatusInternalServerError)
			return
		}
		targetURL, err := url.Parse(string(body))
		if err != nil {
			http.Error(w, fmt.Sprintf("malformed shim open request: %v", err), http.StatusBadRequest)
			return
		}
		targetURL.Scheme = "ws"
		targetURL.Host = host
		conn, err := NewConnection(ctx, targetURL.String(),
			func(err error) {
				log.Printf("Websocket failure: %v", err)
			})
		if err != nil {
			log.Printf("Failed to dial the websocket server %q: %v\n", targetURL.String(), err)
			http.Error(w, fmt.Sprintf("internal error opening a shim connection: %v", err), http.StatusInternalServerError)
			return
		}
		connections.Store(sessionID, conn)
		log.Printf("Websocket connection to the server %q established for session: %v\n", targetURL.String(), sessionID)
		resp := &sessionMessage{
			ID:      sessionID,
			Message: targetURL.String(),
		}
		respBytes, err := json.Marshal(resp)
		if err != nil {
			log.Printf("Failed to serialize the response to a websocket open request: %v", err)
			http.Error(w, fmt.Sprintf("internal error opening a shim connection: %v", err), http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusOK)
		w.Write(respBytes)
	})
	mux.HandleFunc(path.Join(shimPath, "close"), func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, fmt.Sprintf("internal error reading a shim request: %v", err), http.StatusInternalServerError)
			return
		}
		var msg sessionMessage
		if err := json.Unmarshal(body, &msg); err != nil {
			http.Error(w, fmt.Sprintf("error parsing a shim request: %v", err), http.StatusBadRequest)
			return
		}
		c, ok := connections.Load(msg.ID)
		if !ok {
			http.Error(w, fmt.Sprintf("unknown shim session ID: %q", msg.ID), http.StatusBadRequest)
			return
		}
		conn, ok := c.(*Connection)
		if !ok {
			http.Error(w, "internal error reading a shim session", http.StatusInternalServerError)
			return
		}
		connections.Delete(msg.ID)
		conn.Close()
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})
	mux.HandleFunc(path.Join(shimPath, "data"), func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, fmt.Sprintf("internal error reading a shim request: %v", err), http.StatusInternalServerError)
			return
		}
		var msg sessionMessage
		if err := json.Unmarshal(body, &msg); err != nil {
			http.Error(w, fmt.Sprintf("error parsing a shim request: %v", err), http.StatusBadRequest)
			return
		}
		c, ok := connections.Load(msg.ID)
		if !ok {
			http.Error(w, fmt.Sprintf("unknown shim session ID: %q", msg.ID), http.StatusBadRequest)
			return
		}
		conn, ok := c.(*Connection)
		if !ok {
			http.Error(w, "internal error reading a shim session", http.StatusInternalServerError)
			return
		}
		if err := conn.SendClientMessage(msg.Message); err != nil {
			http.Error(w, fmt.Sprintf("attempt to send data on a closed session: %q", msg.ID), http.StatusBadRequest)
			return
		}
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})
	mux.HandleFunc(path.Join(shimPath, "poll"), func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, fmt.Sprintf("internal error reading a shim request: %v", err), http.StatusInternalServerError)
			return
		}
		var msg sessionMessage
		if err := json.Unmarshal(body, &msg); err != nil {
			http.Error(w, fmt.Sprintf("error parsing a shim request: %v", err), http.StatusBadRequest)
			return
		}
		c, ok := connections.Load(msg.ID)
		if !ok {
			http.Error(w, fmt.Sprintf("unknown shim session ID: %q", msg.ID), http.StatusBadRequest)
			return
		}
		conn, ok := c.(*Connection)
		if !ok {
			http.Error(w, "internal error reading a shim session", http.StatusInternalServerError)
			return
		}
		serverMsg, err := conn.ReadServerMessage()
		if err != nil {
			http.Error(w, fmt.Sprintf("attempt to read data from a closed session: %q", msg.ID), http.StatusBadRequest)
			return
		} else if serverMsg == nil {
			w.WriteHeader(http.StatusRequestTimeout)
			return
		}
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(*serverMsg))
	})
	return mux
}

// Proxy creates a reverse proxy that inserts websocket-shim code into all HTML responses.
func Proxy(ctx context.Context, wrapped *httputil.ReverseProxy, host, shimPath string, injectShimCode bool) (http.Handler, error) {
	var templateBuf bytes.Buffer
	if err := shimTmpl.Execute(&templateBuf, &struct{ ShimPath string }{ShimPath: shimPath}); err != nil {
		return nil, err
	}
	shimCode := templateBuf.String()
	if injectShimCode {
		wrapped.ModifyResponse = func(resp *http.Response) error {
			return shimBody(resp, shimCode)
		}
	}
	mux := http.NewServeMux()
	if shimPath != "" {
		shimPath = path.Clean("/"+shimPath) + "/"
		shimServer := createShimChannel(ctx, host, shimPath)
		mux.Handle(shimPath, shimServer)
	}
	mux.Handle("/", wrapped)
	return mux, nil
}
<commit_msg>agent: maintain ordering of the shimmed websocket messages<commit_after>/*
Copyright 2018 Google Inc. All rights reserved.

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

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

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

package websockets

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"path"
	"strings"
	"sync"
	"sync/atomic"
	"text/template"

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

const (
	contentTypeHeader = "Content-Type"
	shimTemplate      = `
<!--START_WEBSOCKET_SHIM-->
<!--
    This file was served from behind a reverse proxy that does not support websockets.

    The following code snippet has been inserted by the proxy to replace websockets
    with socket.io which is based on HTTP and will work with this proxy.

    If this snippet insertion is causing issues, then contact the server administrator.
-->
<script>
(function() {
    if (typeof window.nativeWebSocket !== 'undefined') {
      // We have already replaced websockets
      return;
    }
    console.log('Replacing native websockets with a shim');
    window.nativeWebSocket = window.WebSocket;
    const  location = window.location;
    const shimUri = location.protocol + '//' + location.host + '/{{.ShimPath}}/';

    function shouldShimWebsockets(url) {
      var parsedURL = new URL(url);
      if (typeof parsedURL.host == 'undefined') {
        parsedURL.host = location.host;
      }
      return (parsedURL.host == location.host);
    }

    function WebSocketShim(url, protocols) {
      if (!shouldShimWebsockets(url)) {
        console.log("Not shimming websockets for " + parsedURL.host + " as it does not match the location of " + location.host);
        return new window.nativeWebSocket(url, protocols);
      }

      // We need to reference "this" within nested functions, so we alias it to "self"
      var self = this;

      this.readyState = WebSocketShim.CONNECTING;
      function openedHandler(msg) {
        self.readyState = WebSocketShim.OPEN;
        if (self.onopen) {
          self.onopen({ target: self });
        }
      }
      function receiveHandler(msg) {
        if (self.onmessage) {
          self.onmessage({ target: self, data: msg });
        }
      }
      function errorHandler() {
        if (self.onerror) {
          self.onerror({ target: self });
        }
      }
      self.xhr = function(action, msg, onsuccess, onexit) {
        var req = new XMLHttpRequest();
        req.onreadystatechange = function() {
          if (req.readyState === 4) {
            if (req.status === 200) {
              if (onsuccess) {
                onsuccess(req.responseText);
              }
            } else if (req.status !== 408) {
              errorHandler();
            }
            if (onexit) {
              onexit();
            }
          }
        };
        req.open("POST", shimUri + action, true);
        if (typeof msg !== 'string') {
          msg = JSON.stringify(msg);
        }
        req.send(msg);
      }

      self.closedHandler = function() {
        self.readyState = WebSocketShim.CLOSED;
        if (self.onclose) {
          self.onclose({ target: self });
        }
      }

      self.pendingMessages = [];
      self.pushing = false;
      self.push = function() {
         if (self.pushing) {
           return;
         }
         if (self.pendingMessages.length == 0) {
           return;
         }
         self.pushing = true;
         var msg = self.pendingMessages.shift();
         self.xhr('data', msg, null, function() {
           self.pushing = false;
           self.push();
         })
      }

      function poll() {
        if (self.readyState != WebSocketShim.OPEN) {
          return;
        }
        self.xhr('poll', {'id': self._sessionID}, receiveHandler, poll);
      }

      self.xhr('open', url, function(resp) {
        respJSON = JSON.parse(resp);
        self._sessionID = respJSON.id;
        openedHandler(respJSON.msg);
        poll();
      });
    }
    WebSocketShim.prototype = {
      binaryType: "blob",
      onopen: null,
      onclose: null,
      onmessage: null,
      onerror: null,

      send: function(data) {
        if (this.readyState != WebSocketShim.OPEN) {
          throw new Error('WebSocket is not yet opened');
        }
        this.pendingMessages.push({'id': this._sessionID, 'msg': data});
        self.push();
      },
      close: function() {
        if (this.readyState != WebSocketShim.OPEN) {
          return;
        }
        this.readyState = WebSocketShim.CLOSING;
        this.xhr('close', {'id': this._sessionID}, false, this.closedHandler);
      },
    };
    WebSocketShim.CONNECTING = 0;
    WebSocketShim.OPEN = 1;
    WebSocketShim.CLOSING = 2;
    WebSocketShim.CLOSED = 3;

    window.WebSocket = WebSocketShim;
})();
</script>
<!--END_WEBSOCKET_SHIM-->
`
)

var shimTmpl = template.Must(template.New("client-shim").Parse(shimTemplate))

type shimmedBody struct {
	reader io.Reader
	closer io.Closer
}

func (sb *shimmedBody) Read(p []byte) (n int, err error) {
	return sb.reader.Read(p)
}

func (sb *shimmedBody) Close() error {
	return sb.closer.Close()
}

func shimBody(resp *http.Response, shimCode string) error {
	if resp == nil || resp.Body == nil {
		// We have nothing to do on an empty response
		return nil
	}
	contentType := strings.ToLower(resp.Header.Get(contentTypeHeader))
	if !strings.Contains(contentType, "html") {
		// We only want to modify HTML responses
		return nil
	}
	wrapped := resp.Body

	// Read in the first kilobyte to see if the <head> tag exists in it
	buf := make([]byte, 1024)
	count, err := wrapped.Read(buf)
	if err != nil && err != io.EOF {
		return err
	}
	prefix := strings.Replace(string(buf[0:count]), "<head>", "<head>"+shimCode, 1)
	resp.Body = &shimmedBody{
		reader: io.MultiReader(strings.NewReader(prefix), wrapped),
		closer: wrapped,
	}
	return nil
}

type sessionMessage struct {
	ID      string `json:"id,omitempty"`
	Message string `json:"msg,omitempty"`
}

func createShimChannel(ctx context.Context, host, shimPath string) http.Handler {
	var connections sync.Map
	var sessionCount uint64
	mux := http.NewServeMux()
	mux.HandleFunc(path.Join(shimPath, "open"), func(w http.ResponseWriter, r *http.Request) {
		sessionID := fmt.Sprintf("%d", atomic.AddUint64(&sessionCount, 1))
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, fmt.Sprintf("internal error reading a shim request: %v", err), http.StatusInternalServerError)
			return
		}
		targetURL, err := url.Parse(string(body))
		if err != nil {
			http.Error(w, fmt.Sprintf("malformed shim open request: %v", err), http.StatusBadRequest)
			return
		}
		targetURL.Scheme = "ws"
		targetURL.Host = host
		conn, err := NewConnection(ctx, targetURL.String(),
			func(err error) {
				log.Printf("Websocket failure: %v", err)
			})
		if err != nil {
			log.Printf("Failed to dial the websocket server %q: %v\n", targetURL.String(), err)
			http.Error(w, fmt.Sprintf("internal error opening a shim connection: %v", err), http.StatusInternalServerError)
			return
		}
		connections.Store(sessionID, conn)
		log.Printf("Websocket connection to the server %q established for session: %v\n", targetURL.String(), sessionID)
		resp := &sessionMessage{
			ID:      sessionID,
			Message: targetURL.String(),
		}
		respBytes, err := json.Marshal(resp)
		if err != nil {
			log.Printf("Failed to serialize the response to a websocket open request: %v", err)
			http.Error(w, fmt.Sprintf("internal error opening a shim connection: %v", err), http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusOK)
		w.Write(respBytes)
	})
	mux.HandleFunc(path.Join(shimPath, "close"), func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, fmt.Sprintf("internal error reading a shim request: %v", err), http.StatusInternalServerError)
			return
		}
		var msg sessionMessage
		if err := json.Unmarshal(body, &msg); err != nil {
			http.Error(w, fmt.Sprintf("error parsing a shim request: %v", err), http.StatusBadRequest)
			return
		}
		c, ok := connections.Load(msg.ID)
		if !ok {
			http.Error(w, fmt.Sprintf("unknown shim session ID: %q", msg.ID), http.StatusBadRequest)
			return
		}
		conn, ok := c.(*Connection)
		if !ok {
			http.Error(w, "internal error reading a shim session", http.StatusInternalServerError)
			return
		}
		connections.Delete(msg.ID)
		conn.Close()
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})
	mux.HandleFunc(path.Join(shimPath, "data"), func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, fmt.Sprintf("internal error reading a shim request: %v", err), http.StatusInternalServerError)
			return
		}
		var msg sessionMessage
		if err := json.Unmarshal(body, &msg); err != nil {
			http.Error(w, fmt.Sprintf("error parsing a shim request: %v", err), http.StatusBadRequest)
			return
		}
		c, ok := connections.Load(msg.ID)
		if !ok {
			http.Error(w, fmt.Sprintf("unknown shim session ID: %q", msg.ID), http.StatusBadRequest)
			return
		}
		conn, ok := c.(*Connection)
		if !ok {
			http.Error(w, "internal error reading a shim session", http.StatusInternalServerError)
			return
		}
		if err := conn.SendClientMessage(msg.Message); err != nil {
			http.Error(w, fmt.Sprintf("attempt to send data on a closed session: %q", msg.ID), http.StatusBadRequest)
			return
		}
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})
	mux.HandleFunc(path.Join(shimPath, "poll"), func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, fmt.Sprintf("internal error reading a shim request: %v", err), http.StatusInternalServerError)
			return
		}
		var msg sessionMessage
		if err := json.Unmarshal(body, &msg); err != nil {
			http.Error(w, fmt.Sprintf("error parsing a shim request: %v", err), http.StatusBadRequest)
			return
		}
		c, ok := connections.Load(msg.ID)
		if !ok {
			http.Error(w, fmt.Sprintf("unknown shim session ID: %q", msg.ID), http.StatusBadRequest)
			return
		}
		conn, ok := c.(*Connection)
		if !ok {
			http.Error(w, "internal error reading a shim session", http.StatusInternalServerError)
			return
		}
		serverMsg, err := conn.ReadServerMessage()
		if err != nil {
			http.Error(w, fmt.Sprintf("attempt to read data from a closed session: %q", msg.ID), http.StatusBadRequest)
			return
		} else if serverMsg == nil {
			w.WriteHeader(http.StatusRequestTimeout)
			return
		}
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(*serverMsg))
	})
	return mux
}

// Proxy creates a reverse proxy that inserts websocket-shim code into all HTML responses.
func Proxy(ctx context.Context, wrapped *httputil.ReverseProxy, host, shimPath string, injectShimCode bool) (http.Handler, error) {
	var templateBuf bytes.Buffer
	if err := shimTmpl.Execute(&templateBuf, &struct{ ShimPath string }{ShimPath: shimPath}); err != nil {
		return nil, err
	}
	shimCode := templateBuf.String()
	if injectShimCode {
		wrapped.ModifyResponse = func(resp *http.Response) error {
			return shimBody(resp, shimCode)
		}
	}
	mux := http.NewServeMux()
	if shimPath != "" {
		shimPath = path.Clean("/"+shimPath) + "/"
		shimServer := createShimChannel(ctx, host, shimPath)
		mux.Handle(shimPath, shimServer)
	}
	mux.Handle("/", wrapped)
	return mux, nil
}
<|endoftext|>