<commit_before>package kafka

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"math"
	"net"
	"sync"
	"time"

	"github.com/optiopay/kafka/proto"
)

// ErrClosed is returned as result of any request made using closed connection.
var ErrClosed = errors.New("closed")

// Low level abstraction over connection to Kafka.
type connection struct {
	rw     net.Conn
	stop   chan struct{}
	nextID chan int32
	logger Logger

	mu          sync.Mutex
	respc       map[int32]chan []byte
	stopErr     error
	readTimeout time.Duration
	apiVersions map[int16]proto.SupportedVersion
}

func newTLSConnection(address string, ca, cert, key []byte, timeout, readTimeout time.Duration) (*connection, error) {
	var fetchVersions = true
	for {
		roots := x509.NewCertPool()
		ok := roots.AppendCertsFromPEM(ca)
		if !ok {
			return nil, fmt.Errorf("Cannot parse root certificate")
		}

		certificate, err := tls.X509KeyPair(cert, key)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse key/cert for TLS: %s", err)
		}

		conf := &tls.Config{
			Certificates: []tls.Certificate{certificate},
			RootCAs:      roots,
		}

		dialer := net.Dialer{
			Timeout:   timeout,
			KeepAlive: 30 * time.Second,
		}
		conn, err := tls.DialWithDialer(&dialer, "tcp", address, conf)
		if err != nil {
			return nil, err
		}
		c := &connection{
			stop:        make(chan struct{}),
			nextID:      make(chan int32),
			rw:          conn,
			respc:       make(map[int32]chan []byte),
			logger:      &nullLogger{},
			readTimeout: readTimeout,
			apiVersions: make(map[int16]proto.SupportedVersion),
		}
		go c.nextIDLoop()
		go c.readRespLoop()
		if fetchVersions {
			if c.cacheApiVersions() != nil {
				fetchVersions = false
				c.Close()
			}
		}

		return c, nil
	}

}

// newConnection returns new, initialized connection or error
func newTCPConnection(address string, timeout, readTimeout time.Duration) (*connection, error) {
	var fetchVersions = true
	for {
		dialer := net.Dialer{
			Timeout:   timeout,
			KeepAlive: 30 * time.Second,
		}
		conn, err := dialer.Dial("tcp", address)
		if err != nil {
			return nil, err
		}
		c := &connection{
			stop:        make(chan struct{}),
			nextID:      make(chan int32),
			rw:          conn,
			respc:       make(map[int32]chan []byte),
			logger:      &nullLogger{},
			readTimeout: readTimeout,
			apiVersions: make(map[int16]proto.SupportedVersion),
		}
		go c.nextIDLoop()
		go c.readRespLoop()

		if fetchVersions {
			if c.cacheApiVersions() != nil {
				fetchVersions = false
				c.Close()
			}
		}
		return c, nil
	}

}

func (c *connection) cacheApiVersions() error {
	apiVersions, err := c.APIVersions(&proto.APIVersionsReq{})
	if err != nil {
		c.logger.Debug("cannot fetch apiversions",
			"error", err)
		return err
	}
	for _, api := range apiVersions.APIVersions {
		c.apiVersions[api.APIKey] = api
	}
	return nil
}

//getBestVersion returns version for passed apiKey which best fit server and client requirements
func (c *connection) getBestVersion(apiKey int16) int16 {
	if requested, ok := c.apiVersions[apiKey]; ok {
		supported := proto.SupportedByDriver[apiKey]
		if min(supported.MaxVersion, requested.MaxVersion) >= max(supported.MinVersion, requested.MinVersion) {
			return min(supported.MaxVersion, requested.MaxVersion)
		}
	}
	return 0
}

func min(a int16, b int16) int16 {
	if a < b {
		return a
	}
	return b
}

func max(a int16, b int16) int16 {
	if a > b {
		return a
	}
	return b
}

// nextIDLoop generates correlation IDs, making sure they are always in order
// and within the scope of request-response mapping array.
func (c *connection) nextIDLoop() {
	var id int32 = 1
	for {
		select {
		case <-c.stop:
			close(c.nextID)
			return
		case c.nextID <- id:
			id++
			if id == math.MaxInt32 {
				id = 1
			}
		}
	}
}

// readRespLoop constantly reading response messages from the socket and after
// partial parsing, sends byte representation of the whole message to request
// sending process.
func (c *connection) readRespLoop() {
	defer func() {
		c.mu.Lock()
		for _, cc := range c.respc {
			close(cc)
		}
		c.respc = make(map[int32]chan []byte)
		c.mu.Unlock()
	}()

	rd := bufio.NewReader(c.rw)
	for {
		if c.readTimeout > 0 {
			err := c.rw.SetReadDeadline(time.Now().Add(c.readTimeout))
			if err != nil {
				c.logger.Error("msg", "SetReadDeadline failed",
					"error", err)
			}
		}
		correlationID, b, err := proto.ReadResp(rd)
		if err != nil {
			c.mu.Lock()
			if c.stopErr == nil {
				c.stopErr = err
				close(c.stop)
			}
			c.mu.Unlock()
			return
		}

		c.mu.Lock()
		rc, ok := c.respc[correlationID]
		delete(c.respc, correlationID)
		c.mu.Unlock()
		if !ok {
			c.logger.Warn(
				"msg", "response to unknown request",
				"correlationID", correlationID)
			continue
		}

		select {
		case <-c.stop:
			c.mu.Lock()
			if c.stopErr == nil {
				c.stopErr = ErrClosed
			}
			c.mu.Unlock()
		case rc <- b:
		}
		close(rc)
	}
}

// respWaiter register listener to response message with given correlationID
// and return channel that single response message will be pushed to once it
// will arrive.
// After pushing response message, channel is closed.
//
// Upon connection close, all unconsumed channels are closed.
func (c *connection) respWaiter(correlationID int32) (respc chan []byte, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.stopErr != nil {
		return nil, c.stopErr
	}
	if _, ok := c.respc[correlationID]; ok {
		c.logger.Error("msg", "correlation conflict", "correlationID", correlationID)
		return nil, fmt.Errorf("correlation conflict: %d", correlationID)
	}
	respc = make(chan []byte)
	c.respc[correlationID] = respc
	return respc, nil
}

// releaseWaiter removes response channel from waiters pool and close it.
// Calling this method for unknown correlationID has no effect.
func (c *connection) releaseWaiter(correlationID int32) {
	c.mu.Lock()
	rc, ok := c.respc[correlationID]
	if ok {
		delete(c.respc, correlationID)
		close(rc)
	}
	c.mu.Unlock()
}

// Close close underlying transport connection and cancel all pending response
// waiters.
func (c *connection) Close() error {
	c.mu.Lock()
	if c.stopErr == nil {
		c.stopErr = ErrClosed
		close(c.stop)
	}
	c.mu.Unlock()
	return c.rw.Close()
}

func (c *connection) sendRequest(req proto.Request) ([]byte, error) {
	req.SetVersion(c.getBestVersion(req.Kind()))
	var ok bool
	var correlationID int32
	if correlationID, ok = <-c.nextID; !ok {
		return nil, c.stopErr
	}
	req.SetCorrelationID(correlationID)

	respc, err := c.respWaiter(req.GetCorrelationID())
	if err != nil {
		c.logger.Error("msg", "failed waiting for response", "error", err)
		return nil, fmt.Errorf("wait for response: %s", err)
	}

	if _, err := req.WriteTo(c.rw); err != nil {
		c.logger.Error("msg", "cannot write", "error", err)
		c.releaseWaiter(req.GetCorrelationID())
		return nil, err
	}
	b, ok := <-respc
	if !ok {
		return nil, c.stopErr
	}
	return b, nil
}

func (c *connection) sendRequestWithoutAcks(req proto.Request) error {
	var ok bool
	var correlationID int32
	if correlationID, ok = <-c.nextID; !ok {
		return c.stopErr
	}
	req.SetCorrelationID(correlationID)

	req.SetVersion(c.getBestVersion(req.Kind()))

	_, err := req.WriteTo(c.rw)
	return err
}

// APIVersions sends a request to fetch the supported versions for each API.
// Versioning is only supported in Kafka versions above 0.10.0.0
func (c *connection) APIVersions(req *proto.APIVersionsReq) (*proto.APIVersionsResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedAPIVersionsResp(bytes.NewReader(b), req.Version)
}

// Metadata sends given metadata request to kafka node and returns related
// metadata response.
// Calling this method on closed connection will always return ErrClosed.
func (c *connection) Metadata(req *proto.MetadataReq) (*proto.MetadataResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedMetadataResp(bytes.NewReader(b), req.Version)
}

// Produce sends given produce request to kafka node and returns related
// response. Sending request with no ACKs flag will result with returning nil
// right after sending request, without waiting for response.
// Calling this method on closed connection will always return ErrClosed.
func (c *connection) Produce(req *proto.ProduceReq) (*proto.ProduceResp, error) {

	if req.RequiredAcks == proto.RequiredAcksNone {
		return nil, c.sendRequestWithoutAcks(req)
	}

	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}

	return proto.ReadVersionedProduceResp(bytes.NewReader(b), req.Version)
}

// Fetch sends given fetch request to kafka node and returns related response.
// Calling this method on closed connection will always return ErrClosed.
func (c *connection) Fetch(req *proto.FetchReq) (*proto.FetchResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}

	resp, err := proto.ReadVersionedFetchResp(bytes.NewReader(b), req.Version)
	if err != nil {
		return nil, err
	}

	// Compressed messages are returned in full batches for efficiency
	// (the broker doesn't need to decompress).
	// This means that it's possible to get some leading messages
	// with a smaller offset than requested. Trim those.
	for ti := range resp.Topics {
		topic := &resp.Topics[ti]
		reqTopic := &req.Topics[ti]
		for pi := range topic.Partitions {
			partition := &topic.Partitions[pi]
			reqPartition := &reqTopic.Partitions[pi]
			i := 0
			for _, msg := range partition.Messages {
				if msg.Offset >= reqPartition.FetchOffset {
					break
				}
				i++
			}
			partition.Messages = partition.Messages[i:]
		}
	}
	return resp, nil
}

// Offset sends given offset request to kafka node and returns related response.
// Calling this method on closed connection will always return ErrClosed.
func (c *connection) Offset(req *proto.OffsetReq) (*proto.OffsetResp, error) {
	// TODO(husio) documentation is not mentioning this directly, but I assume
	// -1 is for non node clients
	req.ReplicaID = -1

	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedOffsetResp(bytes.NewReader(b), req.Version)
}

func (c *connection) ConsumerMetadata(req *proto.ConsumerMetadataReq) (*proto.ConsumerMetadataResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedConsumerMetadataResp(bytes.NewReader(b), req.Version)
}

func (c *connection) OffsetCommit(req *proto.OffsetCommitReq) (*proto.OffsetCommitResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedOffsetCommitResp(bytes.NewReader(b), req.Version)
}

func (c *connection) OffsetFetch(req *proto.OffsetFetchReq) (*proto.OffsetFetchResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedOffsetFetchResp(bytes.NewReader(b), req.Version)
}
<commit_msg>fix getting versions for a new connection<commit_after>package kafka

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"math"
	"net"
	"sync"
	"time"

	"github.com/optiopay/kafka/proto"
)

// ErrClosed is returned as result of any request made using closed connection.
var ErrClosed = errors.New("closed")

// Low level abstraction over connection to Kafka.
type connection struct {
	rw     net.Conn
	stop   chan struct{}
	nextID chan int32
	logger Logger

	mu          sync.Mutex
	respc       map[int32]chan []byte
	stopErr     error
	readTimeout time.Duration
	apiVersions map[int16]proto.SupportedVersion
}

func newTLSConnection(address string, ca, cert, key []byte, timeout, readTimeout time.Duration) (*connection, error) {
	var fetchVersions = true
	for {
		roots := x509.NewCertPool()
		ok := roots.AppendCertsFromPEM(ca)
		if !ok {
			return nil, fmt.Errorf("Cannot parse root certificate")
		}

		certificate, err := tls.X509KeyPair(cert, key)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse key/cert for TLS: %s", err)
		}

		conf := &tls.Config{
			Certificates: []tls.Certificate{certificate},
			RootCAs:      roots,
		}

		dialer := net.Dialer{
			Timeout:   timeout,
			KeepAlive: 30 * time.Second,
		}
		conn, err := tls.DialWithDialer(&dialer, "tcp", address, conf)
		if err != nil {
			return nil, err
		}
		c := &connection{
			stop:        make(chan struct{}),
			nextID:      make(chan int32),
			rw:          conn,
			respc:       make(map[int32]chan []byte),
			logger:      &nullLogger{},
			readTimeout: readTimeout,
			apiVersions: make(map[int16]proto.SupportedVersion),
		}
		go c.nextIDLoop()
		go c.readRespLoop()
		if fetchVersions {
			if c.cacheApiVersions() != nil {
				fetchVersions = false
				c.Close()
			}
		}

		return c, nil
	}

}

// newConnection returns new, initialized connection or error
func newTCPConnection(address string, timeout, readTimeout time.Duration) (*connection, error) {
	var fetchVersions = true
	for {
		dialer := net.Dialer{
			Timeout:   timeout,
			KeepAlive: 30 * time.Second,
		}
		conn, err := dialer.Dial("tcp", address)
		if err != nil {
			return nil, err
		}
		c := &connection{
			stop:        make(chan struct{}),
			nextID:      make(chan int32),
			rw:          conn,
			respc:       make(map[int32]chan []byte),
			logger:      &nullLogger{},
			readTimeout: readTimeout,
			apiVersions: make(map[int16]proto.SupportedVersion),
		}
		go c.nextIDLoop()
		go c.readRespLoop()

		if fetchVersions {
			if c.cacheApiVersions() != nil {
				fetchVersions = false
				c.Close()
				continue
			}
		}
		return c, nil
	}

}

func (c *connection) cacheApiVersions() error {
	apiVersions, err := c.APIVersions(&proto.APIVersionsReq{})
	if err != nil {
		c.logger.Debug("cannot fetch apiversions",
			"error", err)
		return err
	}
	for _, api := range apiVersions.APIVersions {
		c.apiVersions[api.APIKey] = api
	}
	return nil
}

//getBestVersion returns version for passed apiKey which best fit server and client requirements
func (c *connection) getBestVersion(apiKey int16) int16 {
	if requested, ok := c.apiVersions[apiKey]; ok {
		supported := proto.SupportedByDriver[apiKey]
		if min(supported.MaxVersion, requested.MaxVersion) >= max(supported.MinVersion, requested.MinVersion) {
			return min(supported.MaxVersion, requested.MaxVersion)
		}
	}
	return 0
}

func min(a int16, b int16) int16 {
	if a < b {
		return a
	}
	return b
}

func max(a int16, b int16) int16 {
	if a > b {
		return a
	}
	return b
}

// nextIDLoop generates correlation IDs, making sure they are always in order
// and within the scope of request-response mapping array.
func (c *connection) nextIDLoop() {
	var id int32 = 1
	for {
		select {
		case <-c.stop:
			close(c.nextID)
			return
		case c.nextID <- id:
			id++
			if id == math.MaxInt32 {
				id = 1
			}
		}
	}
}

// readRespLoop constantly reading response messages from the socket and after
// partial parsing, sends byte representation of the whole message to request
// sending process.
func (c *connection) readRespLoop() {
	defer func() {
		c.mu.Lock()
		for _, cc := range c.respc {
			close(cc)
		}
		c.respc = make(map[int32]chan []byte)
		c.mu.Unlock()
	}()

	rd := bufio.NewReader(c.rw)
	for {
		if c.readTimeout > 0 {
			err := c.rw.SetReadDeadline(time.Now().Add(c.readTimeout))
			if err != nil {
				c.logger.Error("msg", "SetReadDeadline failed",
					"error", err)
			}
		}
		correlationID, b, err := proto.ReadResp(rd)
		if err != nil {
			c.mu.Lock()
			if c.stopErr == nil {
				c.stopErr = err
				close(c.stop)
			}
			c.mu.Unlock()
			return
		}

		c.mu.Lock()
		rc, ok := c.respc[correlationID]
		delete(c.respc, correlationID)
		c.mu.Unlock()
		if !ok {
			c.logger.Warn(
				"msg", "response to unknown request",
				"correlationID", correlationID)
			continue
		}

		select {
		case <-c.stop:
			c.mu.Lock()
			if c.stopErr == nil {
				c.stopErr = ErrClosed
			}
			c.mu.Unlock()
		case rc <- b:
		}
		close(rc)
	}
}

// respWaiter register listener to response message with given correlationID
// and return channel that single response message will be pushed to once it
// will arrive.
// After pushing response message, channel is closed.
//
// Upon connection close, all unconsumed channels are closed.
func (c *connection) respWaiter(correlationID int32) (respc chan []byte, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.stopErr != nil {
		return nil, c.stopErr
	}
	if _, ok := c.respc[correlationID]; ok {
		c.logger.Error("msg", "correlation conflict", "correlationID", correlationID)
		return nil, fmt.Errorf("correlation conflict: %d", correlationID)
	}
	respc = make(chan []byte)
	c.respc[correlationID] = respc
	return respc, nil
}

// releaseWaiter removes response channel from waiters pool and close it.
// Calling this method for unknown correlationID has no effect.
func (c *connection) releaseWaiter(correlationID int32) {
	c.mu.Lock()
	rc, ok := c.respc[correlationID]
	if ok {
		delete(c.respc, correlationID)
		close(rc)
	}
	c.mu.Unlock()
}

// Close close underlying transport connection and cancel all pending response
// waiters.
func (c *connection) Close() error {
	c.mu.Lock()
	if c.stopErr == nil {
		c.stopErr = ErrClosed
		close(c.stop)
	}
	c.mu.Unlock()
	return c.rw.Close()
}

func (c *connection) sendRequest(req proto.Request) ([]byte, error) {
	req.SetVersion(c.getBestVersion(req.Kind()))
	var ok bool
	var correlationID int32
	if correlationID, ok = <-c.nextID; !ok {
		return nil, c.stopErr
	}
	req.SetCorrelationID(correlationID)

	respc, err := c.respWaiter(req.GetCorrelationID())
	if err != nil {
		c.logger.Error("msg", "failed waiting for response", "error", err)
		return nil, fmt.Errorf("wait for response: %s", err)
	}

	if _, err := req.WriteTo(c.rw); err != nil {
		c.logger.Error("msg", "cannot write", "error", err)
		c.releaseWaiter(req.GetCorrelationID())
		return nil, err
	}
	b, ok := <-respc
	if !ok {
		return nil, c.stopErr
	}
	return b, nil
}

func (c *connection) sendRequestWithoutAcks(req proto.Request) error {
	var ok bool
	var correlationID int32
	if correlationID, ok = <-c.nextID; !ok {
		return c.stopErr
	}
	req.SetCorrelationID(correlationID)

	req.SetVersion(c.getBestVersion(req.Kind()))

	_, err := req.WriteTo(c.rw)
	return err
}

// APIVersions sends a request to fetch the supported versions for each API.
// Versioning is only supported in Kafka versions above 0.10.0.0
func (c *connection) APIVersions(req *proto.APIVersionsReq) (*proto.APIVersionsResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedAPIVersionsResp(bytes.NewReader(b), req.Version)
}

// Metadata sends given metadata request to kafka node and returns related
// metadata response.
// Calling this method on closed connection will always return ErrClosed.
func (c *connection) Metadata(req *proto.MetadataReq) (*proto.MetadataResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedMetadataResp(bytes.NewReader(b), req.Version)
}

// Produce sends given produce request to kafka node and returns related
// response. Sending request with no ACKs flag will result with returning nil
// right after sending request, without waiting for response.
// Calling this method on closed connection will always return ErrClosed.
func (c *connection) Produce(req *proto.ProduceReq) (*proto.ProduceResp, error) {

	if req.RequiredAcks == proto.RequiredAcksNone {
		return nil, c.sendRequestWithoutAcks(req)
	}

	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}

	return proto.ReadVersionedProduceResp(bytes.NewReader(b), req.Version)
}

// Fetch sends given fetch request to kafka node and returns related response.
// Calling this method on closed connection will always return ErrClosed.
func (c *connection) Fetch(req *proto.FetchReq) (*proto.FetchResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}

	resp, err := proto.ReadVersionedFetchResp(bytes.NewReader(b), req.Version)
	if err != nil {
		return nil, err
	}

	// Compressed messages are returned in full batches for efficiency
	// (the broker doesn't need to decompress).
	// This means that it's possible to get some leading messages
	// with a smaller offset than requested. Trim those.
	for ti := range resp.Topics {
		topic := &resp.Topics[ti]
		reqTopic := &req.Topics[ti]
		for pi := range topic.Partitions {
			partition := &topic.Partitions[pi]
			reqPartition := &reqTopic.Partitions[pi]
			i := 0
			for _, msg := range partition.Messages {
				if msg.Offset >= reqPartition.FetchOffset {
					break
				}
				i++
			}
			partition.Messages = partition.Messages[i:]
		}
	}
	return resp, nil
}

// Offset sends given offset request to kafka node and returns related response.
// Calling this method on closed connection will always return ErrClosed.
func (c *connection) Offset(req *proto.OffsetReq) (*proto.OffsetResp, error) {
	// TODO(husio) documentation is not mentioning this directly, but I assume
	// -1 is for non node clients
	req.ReplicaID = -1

	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedOffsetResp(bytes.NewReader(b), req.Version)
}

func (c *connection) ConsumerMetadata(req *proto.ConsumerMetadataReq) (*proto.ConsumerMetadataResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedConsumerMetadataResp(bytes.NewReader(b), req.Version)
}

func (c *connection) OffsetCommit(req *proto.OffsetCommitReq) (*proto.OffsetCommitResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedOffsetCommitResp(bytes.NewReader(b), req.Version)
}

func (c *connection) OffsetFetch(req *proto.OffsetFetchReq) (*proto.OffsetFetchResp, error) {
	b, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}
	return proto.ReadVersionedOffsetFetchResp(bytes.NewReader(b), req.Version)
}
<|endoftext|>


<commit_before>package main

import (
	//	"bufio"
	"errors"
	//	proto "github.com/shirou/mqtt"
	proto "github.com/huin/mqtt"
	"io"
	"log"
	"net"
	"strings"
)

// ConnectionErrors is an array of errors corresponding to the
// Connect return codes specified in the specification.
var ConnectionErrors = [6]error{
	nil, // Connection Accepted (not an error)
	errors.New("Connection Refused: unacceptable protocol version"),
	errors.New("Connection Refused: identifier rejected"),
	errors.New("Connection Refused: server unavailable"),
	errors.New("Connection Refused: bad user name or password"),
	errors.New("Connection Refused: not authorized"),
}

const (
	ClientAvailable   int = 0
	ClientUnAvailable     // no PINGACK, no DISCONNECT
)

type Connection struct {
	broker   *Broker
	conn     net.Conn
	clientid string
	storage  Storage
	jobs     chan job
	Done     chan struct{}
	Status   int
}

type job struct {
	m proto.Message
	r receipt
}

type receipt chan struct{}

// Wait for the receipt to indicate that the job is done.
func (r receipt) wait() {
	// TODO: timeout
	<-r
}

func (c *Connection) handleConnection() {
	defer func() {
		c.conn.Close()
		c.broker.stats.clientDisconnect(c)
		close(c.jobs)
	}()

	for {
		m, err := proto.DecodeOneMessage(c.conn, nil)
		if err != nil {
			if err == io.EOF {
				return
			}
			if strings.HasSuffix(err.Error(), "use of closed network connection") {
				return
			}
			log.Print("reader: ", err)
			return
		}
		log.Printf("incoming: %T", m)
		switch m := m.(type) {
		case *proto.Connect:
			c.handleConnect(m)
		case *proto.Publish:
			c.handlePublish(m)
		case *proto.PingReq:
			c.submit(&proto.PingResp{})
		case *proto.Disconnect:
			// finish this goroutine
			return
		case *proto.Subscribe:
			c.handleSubscribe(m)
		case *proto.Unsubscribe:
			c.handleUnsubscribe(m)
		default:
			log.Printf("reader: unknown msg type %T, continue anyway", m)
			continue
		}
	}
}

func (c *Connection) handleSubscribe(m *proto.Subscribe) {
	if m.Header.QosLevel != proto.QosAtLeastOnce {
		// protocol error, disconnect
		return
	}
	suback := &proto.SubAck{
		MessageId: m.MessageId,
		TopicsQos: make([]proto.QosLevel, len(m.Topics)),
	}
	for i, tq := range m.Topics {
		// TODO: Handle varying QoS correctly
		c.broker.Subscribe(tq.Topic, c)
		suback.TopicsQos[i] = proto.QosAtMostOnce
	}
	c.submit(suback)
	// Process retained messages.
	for _, tq := range m.Topics {
		if pubmsg, ok := c.broker.storage.GetRetain(tq.Topic); ok {
			c.submit(pubmsg)
		}
	}
}

func (c *Connection) handleUnsubscribe(m *proto.Unsubscribe) {
	for _, t := range m.Topics {
		c.broker.Unsubscribe(t, c)
	}
	ack := &proto.UnsubAck{MessageId: m.MessageId}
	c.submit(ack)
}

func (c *Connection) handleConnect(m *proto.Connect) {
	rc := proto.RetCodeAccepted
	if m.ProtocolName != "MQIsdp" ||
		m.ProtocolVersion != 3 {
		log.Print("reader: reject connection from ", m.ProtocolName, " version ", m.ProtocolVersion)
		rc = proto.RetCodeUnacceptableProtocolVersion
	}

	// Check client id.
	if len(m.ClientId) < 1 || len(m.ClientId) > 23 {
		rc = proto.RetCodeIdentifierRejected
	}
	c.clientid = m.ClientId

	currrent_c, err := c.storage.MergeClient(c.clientid, c)
	if err != nil {
		c.storage.DeleteClient(c.clientid, c)
		return
	}

	// TODO: Last will
	connack := &proto.ConnAck{
		ReturnCode: rc,
	}

	currrent_c.submit(connack)

	// close connection if it was a bad connect
	if rc != proto.RetCodeAccepted {
		log.Printf("Connection refused for %v: %v", currrent_c.conn.RemoteAddr(), ConnectionErrors[rc])
		return
	}

	// Log in mosquitto format.
	clean := 0
	if m.CleanSession {
		clean = 1
	}
	log.Printf("New client connected from %v as %v (c%v, k%v).", currrent_c.conn.RemoteAddr(), currrent_c.clientid, clean, m.KeepAliveTimer)
}

func (c *Connection) handlePublish(m *proto.Publish) {
	// TODO: Proper QoS support
	if m.Header.QosLevel != proto.QosAtMostOnce {
		log.Printf("reader: no support for QoS %v yet", m.Header.QosLevel)
		return
	}
	c.broker.Publish(m)

	if m.Header.Retain {
		c.broker.UpdateRetain(m)
		log.Printf("Publish msg retained: %s", m.TopicName)
	}
	c.submit(&proto.PubAck{MessageId: m.MessageId})
}

// Queue a message; no notification of sending is done.
func (c *Connection) submit(m proto.Message) {
	j := job{m: m}
	select {
	case c.jobs <- j:
	default:
		log.Print(c, ": failed to submit message")
	}
	return
}

// Queue a message, returns a channel that will be readable
// when the message is sent.
func (c *Connection) submitSync(m proto.Message) receipt {
	j := job{m: m, r: make(receipt)}
	c.jobs <- j
	return j.r
}

func (c *Connection) writer() {

	// Close connection on exit in order to cause reader to exit.
	defer func() {
		c.conn.Close()
		c.storage.DeleteClient(c.clientid, c)
		//		c.svr.subs.unsubAll(c)
	}()

	for job := range c.jobs {
		// TODO: write timeout
		err := job.m.Encode(c.conn)
		if job.r != nil {
			// notifiy the sender that this message is sent
			close(job.r)
		}
		if err != nil {
			// This one is not interesting; it happens when clients
			// disappear before we send their acks.
			if err.Error() == "use of closed network connection" {
				return
			}
			log.Print("writer: ", err)
			return
		}
		//		c.svr.stats.messageSend()

		if _, ok := job.m.(*proto.Disconnect); ok {
			log.Print("writer: sent disconnect message")
			return
		}
	}
}

func (c *Connection) Start() {
	go c.handleConnection()
	go c.writer()

	/*
		defer c.conn.Close()
		reader := bufio.NewReader(c.conn)
		for {
			_, err := reader.ReadByte()
			if err != nil {
				return
			}
		}
	*/
}

func NewConnection(b *Broker, conn net.Conn) *Connection {
	c := &Connection{
		broker:  b,
		conn:    conn,
		storage: b.storage,
		jobs:    make(chan job, b.conf.Queue.SendingQueueLength),
		Status:  ClientAvailable,
		//		out:      make(chan job, clientQueueLength),
		//		Incoming: make(chan *proto.Publish, clientQueueLength),
		//		done:     make(chan struct{}),
		//		connack:  make(chan *proto.ConnAck),
		//		suback:   make(chan *proto.SubAck),
	}
	return c
}
<commit_msg>implement Disconnect<commit_after>package main

import (
	//	"bufio"
	"errors"
	//	proto "github.com/shirou/mqtt"
	proto "github.com/huin/mqtt"
	"io"
	"log"
	"net"
	"strings"
)

// ConnectionErrors is an array of errors corresponding to the
// Connect return codes specified in the specification.
var ConnectionErrors = [6]error{
	nil, // Connection Accepted (not an error)
	errors.New("Connection Refused: unacceptable protocol version"),
	errors.New("Connection Refused: identifier rejected"),
	errors.New("Connection Refused: server unavailable"),
	errors.New("Connection Refused: bad user name or password"),
	errors.New("Connection Refused: not authorized"),
}

const (
	ClientAvailable   int = 0
	ClientUnAvailable     // no PINGACK, no DISCONNECT
)

type Connection struct {
	broker    *Broker
	conn      net.Conn
	clientid  string
	storage   Storage
	jobs      chan job
	Done      chan struct{}
	Status    int
	TopicList []string // Subscribed topic list
}

type job struct {
	m proto.Message
	r receipt
}

type receipt chan struct{}

// Wait for the receipt to indicate that the job is done.
func (r receipt) wait() {
	// TODO: timeout
	<-r
}

func (c *Connection) handleConnection() {
	defer func() {
		c.conn.Close()
		c.broker.stats.clientDisconnect(c)
		close(c.jobs)
	}()

	for {
		m, err := proto.DecodeOneMessage(c.conn, nil)
		if err != nil {
			if err == io.EOF {
				return
			}
			if strings.HasSuffix(err.Error(), "use of closed network connection") {
				return
			}
			log.Print("reader: ", err)
			return
		}
		log.Printf("incoming: %T", m)
		switch m := m.(type) {
		case *proto.Connect:
			c.handleConnect(m)
		case *proto.Publish:
			c.handlePublish(m)
		case *proto.PingReq:
			c.submit(&proto.PingResp{})
		case *proto.Disconnect:
			c.handleDisconnect(m)
			return
		case *proto.Subscribe:
			c.handleSubscribe(m)
		case *proto.Unsubscribe:
			c.handleUnsubscribe(m)
		default:
			log.Printf("reader: unknown msg type %T, continue anyway", m)
			continue
		}
	}
}

func (c *Connection) handleSubscribe(m *proto.Subscribe) {
	if m.Header.QosLevel != proto.QosAtLeastOnce {
		// protocol error, disconnect
		return
	}
	suback := &proto.SubAck{
		MessageId: m.MessageId,
		TopicsQos: make([]proto.QosLevel, len(m.Topics)),
	}
	for i, tq := range m.Topics {
		// TODO: Handle varying QoS correctly
		c.broker.Subscribe(tq.Topic, c)
		suback.TopicsQos[i] = proto.QosAtMostOnce

		c.TopicList = append(c.TopicList, tq.Topic)
	}
	c.submit(suback)
	// Process retained messages.
	for _, tq := range m.Topics {
		if pubmsg, ok := c.broker.storage.GetRetain(tq.Topic); ok {
			c.submit(pubmsg)
		}
	}
}

func (c *Connection) handleUnsubscribe(m *proto.Unsubscribe) {
	for _, topic := range m.Topics {
		c.broker.Unsubscribe(topic, c)
	}
	ack := &proto.UnsubAck{MessageId: m.MessageId}
	c.submit(ack)
}

func (c *Connection) handleConnect(m *proto.Connect) {
	rc := proto.RetCodeAccepted
	if m.ProtocolName != "MQIsdp" ||
		m.ProtocolVersion != 3 {
		log.Print("reader: reject connection from ", m.ProtocolName, " version ", m.ProtocolVersion)
		rc = proto.RetCodeUnacceptableProtocolVersion
	}

	// Check client id.
	if len(m.ClientId) < 1 || len(m.ClientId) > 23 {
		rc = proto.RetCodeIdentifierRejected
	}
	c.clientid = m.ClientId

	currrent_c, err := c.storage.MergeClient(c.clientid, c)
	if err != nil {
		c.storage.DeleteClient(c.clientid, c)
		return
	}

	// TODO: Last will
	connack := &proto.ConnAck{
		ReturnCode: rc,
	}

	currrent_c.submit(connack)

	// close connection if it was a bad connect
	if rc != proto.RetCodeAccepted {
		log.Printf("Connection refused for %v: %v", currrent_c.conn.RemoteAddr(), ConnectionErrors[rc])
		return
	}

	// Log in mosquitto format.
	clean := 0
	if m.CleanSession {
		clean = 1
	}
	log.Printf("New client connected from %v as %v (c%v, k%v).", currrent_c.conn.RemoteAddr(), currrent_c.clientid, clean, m.KeepAliveTimer)
}

func (c *Connection) handleDisconnect(m *proto.Disconnect) {
	for _, topic := range c.TopicList {
		c.broker.Unsubscribe(topic, c)
	}
}

func (c *Connection) handlePublish(m *proto.Publish) {
	// TODO: Proper QoS support
	if m.Header.QosLevel != proto.QosAtMostOnce {
		log.Printf("reader: no support for QoS %v yet", m.Header.QosLevel)
		return
	}
	c.broker.Publish(m)

	if m.Header.Retain {
		c.broker.UpdateRetain(m)
		log.Printf("Publish msg retained: %s", m.TopicName)
	}
	c.submit(&proto.PubAck{MessageId: m.MessageId})
}

// Queue a message; no notification of sending is done.
func (c *Connection) submit(m proto.Message) {
	j := job{m: m}
	select {
	case c.jobs <- j:
	default:
		log.Print(c, ": failed to submit message")
	}
	return
}

// Queue a message, returns a channel that will be readable
// when the message is sent.
func (c *Connection) submitSync(m proto.Message) receipt {
	j := job{m: m, r: make(receipt)}
	c.jobs <- j
	return j.r
}

func (c *Connection) writer() {

	// Close connection on exit in order to cause reader to exit.
	defer func() {
		c.conn.Close()
		c.storage.DeleteClient(c.clientid, c)
		//		c.svr.subs.unsubAll(c)
	}()

	for job := range c.jobs {
		// TODO: write timeout
		err := job.m.Encode(c.conn)
		if job.r != nil {
			// notifiy the sender that this message is sent
			close(job.r)
		}
		if err != nil {
			// This one is not interesting; it happens when clients
			// disappear before we send their acks.
			if err.Error() == "use of closed network connection" {
				return
			}
			log.Print("writer: ", err)
			return
		}
		//		c.svr.stats.messageSend()

		if _, ok := job.m.(*proto.Disconnect); ok {
			log.Print("writer: sent disconnect message")
			return
		}
	}
}

func (c *Connection) Start() {
	go c.handleConnection()
	go c.writer()

	/*
		defer c.conn.Close()
		reader := bufio.NewReader(c.conn)
		for {
			_, err := reader.ReadByte()
			if err != nil {
				return
			}
		}
	*/
}

func NewConnection(b *Broker, conn net.Conn) *Connection {
	c := &Connection{
		broker:  b,
		conn:    conn,
		storage: b.storage,
		jobs:    make(chan job, b.conf.Queue.SendingQueueLength),
		Status:  ClientAvailable,
		//		out:      make(chan job, clientQueueLength),
		//		Incoming: make(chan *proto.Publish, clientQueueLength),
		//		done:     make(chan struct{}),
		//		connack:  make(chan *proto.ConnAck),
		//		suback:   make(chan *proto.SubAck),
	}
	return c
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxc/config"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/gnuflag"
	"github.com/lxc/lxd/shared/i18n"
)

type profileList []string

var configMap map[string]string

func (f *profileList) String() string {
	return fmt.Sprint(*f)
}

type configList []string

func (f *configList) String() string {
	return fmt.Sprint(configMap)
}

func (f *configList) Set(value string) error {
	if value == "" {
		return fmt.Errorf(i18n.G("Invalid configuration key"))
	}

	items := strings.SplitN(value, "=", 2)
	if len(items) < 2 {
		return fmt.Errorf(i18n.G("Invalid configuration key"))
	}

	if configMap == nil {
		configMap = map[string]string{}
	}

	configMap[items[0]] = items[1]

	return nil
}

func (f *profileList) Set(value string) error {
	if value == "" {
		initRequestedEmptyProfiles = true
		return nil
	}
	if f == nil {
		*f = make(profileList, 1)
	} else {
		*f = append(*f, value)
	}
	return nil
}

var initRequestedEmptyProfiles bool

type initCmd struct {
	profArgs     profileList
	confArgs     configList
	ephem        bool
	network      string
	storagePool  string
	instanceType string
}

func (c *initCmd) showByDefault() bool {
	return false
}

func (c *initCmd) usage() string {
	return i18n.G(
		`Usage: lxc init [<remote>:]<image> [<remote>:][<name>] [--ephemeral|-e] [--profile|-p <profile>...] [--config|-c <key=value>...] [--network|-n <network>] [--storage|-s <pool>] [--type|-t <instance type>]

Create containers from images.

Not specifying -p will result in the default profile.
Specifying "-p" with no argument will result in no profile.

Examples:
    lxc init ubuntu:16.04 u1`)
}

func (c *initCmd) is_ephem(s string) bool {
	switch s {
	case "-e":
		return true
	case "--ephemeral":
		return true
	}
	return false
}

func (c *initCmd) is_profile(s string) bool {
	switch s {
	case "-p":
		return true
	case "--profile":
		return true
	}
	return false
}

func (c *initCmd) massage_args() {
	l := len(os.Args)
	if l < 2 {
		return
	}

	if c.is_profile(os.Args[l-1]) {
		initRequestedEmptyProfiles = true
		os.Args = os.Args[0 : l-1]
		return
	}

	if l < 3 {
		return
	}

	/* catch "lxc init ubuntu -p -e */
	if c.is_ephem(os.Args[l-1]) && c.is_profile(os.Args[l-2]) {
		initRequestedEmptyProfiles = true
		newargs := os.Args[0 : l-2]
		newargs = append(newargs, os.Args[l-1])
		os.Args = newargs
		return
	}
}

func (c *initCmd) flags() {
	c.massage_args()
	gnuflag.Var(&c.confArgs, "config", i18n.G("Config key/value to apply to the new container"))
	gnuflag.Var(&c.confArgs, "c", i18n.G("Config key/value to apply to the new container"))
	gnuflag.Var(&c.profArgs, "profile", i18n.G("Profile to apply to the new container"))
	gnuflag.Var(&c.profArgs, "p", i18n.G("Profile to apply to the new container"))
	gnuflag.BoolVar(&c.ephem, "ephemeral", false, i18n.G("Ephemeral container"))
	gnuflag.BoolVar(&c.ephem, "e", false, i18n.G("Ephemeral container"))
	gnuflag.StringVar(&c.network, "network", "", i18n.G("Network name"))
	gnuflag.StringVar(&c.network, "n", "", i18n.G("Network name"))
	gnuflag.StringVar(&c.storagePool, "storage", "", i18n.G("Storage pool name"))
	gnuflag.StringVar(&c.storagePool, "s", "", i18n.G("Storage pool name"))
	gnuflag.StringVar(&c.instanceType, "t", "", i18n.G("Instance type"))
}

func (c *initCmd) run(conf *config.Config, args []string) error {
	_, _, err := c.create(conf, args)
	return err
}

func (c *initCmd) create(conf *config.Config, args []string) (lxd.ContainerServer, string, error) {
	if len(args) > 2 || len(args) < 1 {
		return nil, "", errArgs
	}

	iremote, image, err := conf.ParseRemote(args[0])
	if err != nil {
		return nil, "", err
	}

	var name string
	var remote string
	if len(args) == 2 {
		remote, name, err = conf.ParseRemote(args[1])
		if err != nil {
			return nil, "", err
		}
	} else {
		remote, name, err = conf.ParseRemote("")
		if err != nil {
			return nil, "", err
		}
	}

	d, err := conf.GetContainerServer(remote)
	if err != nil {
		return nil, "", err
	}

	/*
	 * initRequestedEmptyProfiles means user requested empty
	 * !initRequestedEmptyProfiles but len(profArgs) == 0 means use profile default
	 */
	profiles := []string{}
	for _, p := range c.profArgs {
		profiles = append(profiles, p)
	}

	if name == "" {
		fmt.Printf(i18n.G("Creating the container") + "\n")
	} else {
		fmt.Printf(i18n.G("Creating %s")+"\n", name)
	}

	devicesMap := map[string]map[string]string{}
	if c.network != "" {
		network, _, err := d.GetNetwork(c.network)
		if err != nil {
			return nil, "", err
		}

		if network.Type == "bridge" {
			devicesMap[c.network] = map[string]string{"type": "nic", "nictype": "bridged", "parent": c.network}
		} else {
			devicesMap[c.network] = map[string]string{"type": "nic", "nictype": "macvlan", "parent": c.network}
		}
	}

	// Check if the specified storage pool exists.
	if c.storagePool != "" {
		_, _, err := d.GetStoragePool(c.storagePool)
		if err != nil {
			return nil, "", err
		}

		devicesMap["root"] = map[string]string{
			"type": "disk",
			"path": "/",
			"pool": c.storagePool,
		}
	}

	// Get the image server and image info
	iremote, image = c.guessImage(conf, d, remote, iremote, image)
	var imgRemote lxd.ImageServer
	var imgInfo *api.Image

	// Connect to the image server
	if iremote == remote {
		imgRemote = d
	} else {
		imgRemote, err = conf.GetImageServer(iremote)
		if err != nil {
			return nil, "", err
		}
	}

	// Deal with the default image
	if image == "" {
		image = "default"
	}

	// Setup container creation request
	req := api.ContainersPost{
		Name:         name,
		InstanceType: c.instanceType,
	}
	req.Config = configMap
	req.Devices = devicesMap
	if !initRequestedEmptyProfiles && len(profiles) == 0 {
		req.Profiles = nil
	} else {
		req.Profiles = profiles
	}
	req.Ephemeral = c.ephem

	// Optimisation for simplestreams
	if conf.Remotes[iremote].Protocol == "simplestreams" {
		imgInfo = &api.Image{}
		imgInfo.Fingerprint = image
		imgInfo.Public = true
		req.Source.Alias = image
	} else {
		// Attempt to resolve an image alias
		alias, _, err := imgRemote.GetImageAlias(image)
		if err == nil {
			req.Source.Alias = image
			image = alias.Target
		}

		// Get the image info
		imgInfo, _, err = imgRemote.GetImage(image)
		if err != nil {
			return nil, "", err
		}
	}

	// Create the container
	op, err := d.CreateContainerFromImage(imgRemote, *imgInfo, req)
	if err != nil {
		return nil, "", err
	}

	// Watch the background operation
	progress := ProgressRenderer{Format: i18n.G("Retrieving image: %s")}
	_, err = op.AddHandler(progress.UpdateOp)
	if err != nil {
		progress.Done("")
		return nil, "", err
	}

	err = cancelableWait(op, &progress)
	if err != nil {
		progress.Done("")
		return nil, "", err
	}
	progress.Done("")

	// Extract the container name
	opInfo, err := op.GetTarget()
	if err != nil {
		return nil, "", err
	}

	containers, ok := opInfo.Resources["containers"]
	if !ok || len(containers) == 0 {
		return nil, "", fmt.Errorf(i18n.G("didn't get any affected image, container or snapshot from server"))
	}

	if len(containers) == 1 && name == "" {
		fields := strings.Split(containers[0], "/")
		name = fields[len(fields)-1]
		fmt.Printf(i18n.G("Container name is: %s")+"\n", name)
	}

	// Validate the network setup
	c.checkNetwork(d, name)

	return d, name, nil
}

func (c *initCmd) guessImage(conf *config.Config, d lxd.ContainerServer, remote string, iremote string, image string) (string, string) {
	if remote != iremote {
		return iremote, image
	}

	_, ok := conf.Remotes[image]
	if !ok {
		return iremote, image
	}

	_, _, err := d.GetImageAlias(image)
	if err == nil {
		return iremote, image
	}

	_, _, err = d.GetImage(image)
	if err == nil {
		return iremote, image
	}

	fmt.Fprintf(os.Stderr, i18n.G("The local image '%s' couldn't be found, trying '%s:' instead.")+"\n", image, image)
	return image, "default"
}

func (c *initCmd) checkNetwork(d lxd.ContainerServer, name string) {
	ct, _, err := d.GetContainer(name)
	if err != nil {
		return
	}

	for _, d := range ct.ExpandedDevices {
		if d["type"] == "nic" {
			return
		}
	}

	fmt.Fprintf(os.Stderr, "\n"+i18n.G("The container you are starting doesn't have any network attached to it.")+"\n")
	fmt.Fprintf(os.Stderr, "  "+i18n.G("To create a new network, use: lxc network create")+"\n")
	fmt.Fprintf(os.Stderr, "  "+i18n.G("To attach a network to a container, use: lxc network attach")+"\n\n")
}
<commit_msg>lxc: Make the <remote>/<alias> syntax work<commit_after>package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxc/config"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/gnuflag"
	"github.com/lxc/lxd/shared/i18n"
)

type profileList []string

var configMap map[string]string

func (f *profileList) String() string {
	return fmt.Sprint(*f)
}

type configList []string

func (f *configList) String() string {
	return fmt.Sprint(configMap)
}

func (f *configList) Set(value string) error {
	if value == "" {
		return fmt.Errorf(i18n.G("Invalid configuration key"))
	}

	items := strings.SplitN(value, "=", 2)
	if len(items) < 2 {
		return fmt.Errorf(i18n.G("Invalid configuration key"))
	}

	if configMap == nil {
		configMap = map[string]string{}
	}

	configMap[items[0]] = items[1]

	return nil
}

func (f *profileList) Set(value string) error {
	if value == "" {
		initRequestedEmptyProfiles = true
		return nil
	}
	if f == nil {
		*f = make(profileList, 1)
	} else {
		*f = append(*f, value)
	}
	return nil
}

var initRequestedEmptyProfiles bool

type initCmd struct {
	profArgs     profileList
	confArgs     configList
	ephem        bool
	network      string
	storagePool  string
	instanceType string
}

func (c *initCmd) showByDefault() bool {
	return false
}

func (c *initCmd) usage() string {
	return i18n.G(
		`Usage: lxc init [<remote>:]<image> [<remote>:][<name>] [--ephemeral|-e] [--profile|-p <profile>...] [--config|-c <key=value>...] [--network|-n <network>] [--storage|-s <pool>] [--type|-t <instance type>]

Create containers from images.

Not specifying -p will result in the default profile.
Specifying "-p" with no argument will result in no profile.

Examples:
    lxc init ubuntu:16.04 u1`)
}

func (c *initCmd) is_ephem(s string) bool {
	switch s {
	case "-e":
		return true
	case "--ephemeral":
		return true
	}
	return false
}

func (c *initCmd) is_profile(s string) bool {
	switch s {
	case "-p":
		return true
	case "--profile":
		return true
	}
	return false
}

func (c *initCmd) massage_args() {
	l := len(os.Args)
	if l < 2 {
		return
	}

	if c.is_profile(os.Args[l-1]) {
		initRequestedEmptyProfiles = true
		os.Args = os.Args[0 : l-1]
		return
	}

	if l < 3 {
		return
	}

	/* catch "lxc init ubuntu -p -e */
	if c.is_ephem(os.Args[l-1]) && c.is_profile(os.Args[l-2]) {
		initRequestedEmptyProfiles = true
		newargs := os.Args[0 : l-2]
		newargs = append(newargs, os.Args[l-1])
		os.Args = newargs
		return
	}
}

func (c *initCmd) flags() {
	c.massage_args()
	gnuflag.Var(&c.confArgs, "config", i18n.G("Config key/value to apply to the new container"))
	gnuflag.Var(&c.confArgs, "c", i18n.G("Config key/value to apply to the new container"))
	gnuflag.Var(&c.profArgs, "profile", i18n.G("Profile to apply to the new container"))
	gnuflag.Var(&c.profArgs, "p", i18n.G("Profile to apply to the new container"))
	gnuflag.BoolVar(&c.ephem, "ephemeral", false, i18n.G("Ephemeral container"))
	gnuflag.BoolVar(&c.ephem, "e", false, i18n.G("Ephemeral container"))
	gnuflag.StringVar(&c.network, "network", "", i18n.G("Network name"))
	gnuflag.StringVar(&c.network, "n", "", i18n.G("Network name"))
	gnuflag.StringVar(&c.storagePool, "storage", "", i18n.G("Storage pool name"))
	gnuflag.StringVar(&c.storagePool, "s", "", i18n.G("Storage pool name"))
	gnuflag.StringVar(&c.instanceType, "t", "", i18n.G("Instance type"))
}

func (c *initCmd) run(conf *config.Config, args []string) error {
	_, _, err := c.create(conf, args)
	return err
}

func (c *initCmd) create(conf *config.Config, args []string) (lxd.ContainerServer, string, error) {
	if len(args) > 2 || len(args) < 1 {
		return nil, "", errArgs
	}

	iremote, image, err := conf.ParseRemote(args[0])
	if err != nil {
		return nil, "", err
	}

	var name string
	var remote string
	if len(args) == 2 {
		remote, name, err = conf.ParseRemote(args[1])
		if err != nil {
			return nil, "", err
		}
	} else {
		remote, name, err = conf.ParseRemote("")
		if err != nil {
			return nil, "", err
		}
	}

	d, err := conf.GetContainerServer(remote)
	if err != nil {
		return nil, "", err
	}

	/*
	 * initRequestedEmptyProfiles means user requested empty
	 * !initRequestedEmptyProfiles but len(profArgs) == 0 means use profile default
	 */
	profiles := []string{}
	for _, p := range c.profArgs {
		profiles = append(profiles, p)
	}

	if name == "" {
		fmt.Printf(i18n.G("Creating the container") + "\n")
	} else {
		fmt.Printf(i18n.G("Creating %s")+"\n", name)
	}

	devicesMap := map[string]map[string]string{}
	if c.network != "" {
		network, _, err := d.GetNetwork(c.network)
		if err != nil {
			return nil, "", err
		}

		if network.Type == "bridge" {
			devicesMap[c.network] = map[string]string{"type": "nic", "nictype": "bridged", "parent": c.network}
		} else {
			devicesMap[c.network] = map[string]string{"type": "nic", "nictype": "macvlan", "parent": c.network}
		}
	}

	// Check if the specified storage pool exists.
	if c.storagePool != "" {
		_, _, err := d.GetStoragePool(c.storagePool)
		if err != nil {
			return nil, "", err
		}

		devicesMap["root"] = map[string]string{
			"type": "disk",
			"path": "/",
			"pool": c.storagePool,
		}
	}

	// Get the image server and image info
	iremote, image = c.guessImage(conf, d, remote, iremote, image)
	var imgRemote lxd.ImageServer
	var imgInfo *api.Image

	// Connect to the image server
	if iremote == remote {
		imgRemote = d
	} else {
		imgRemote, err = conf.GetImageServer(iremote)
		if err != nil {
			return nil, "", err
		}
	}

	// Deal with the default image
	if image == "" {
		image = "default"
	}

	// Setup container creation request
	req := api.ContainersPost{
		Name:         name,
		InstanceType: c.instanceType,
	}
	req.Config = configMap
	req.Devices = devicesMap
	if !initRequestedEmptyProfiles && len(profiles) == 0 {
		req.Profiles = nil
	} else {
		req.Profiles = profiles
	}
	req.Ephemeral = c.ephem

	// Optimisation for simplestreams
	if conf.Remotes[iremote].Protocol == "simplestreams" {
		imgInfo = &api.Image{}
		imgInfo.Fingerprint = image
		imgInfo.Public = true
		req.Source.Alias = image
	} else {
		// Attempt to resolve an image alias
		alias, _, err := imgRemote.GetImageAlias(image)
		if err == nil {
			req.Source.Alias = image
			image = alias.Target
		}

		// Get the image info
		imgInfo, _, err = imgRemote.GetImage(image)
		if err != nil {
			return nil, "", err
		}
	}

	// Create the container
	op, err := d.CreateContainerFromImage(imgRemote, *imgInfo, req)
	if err != nil {
		return nil, "", err
	}

	// Watch the background operation
	progress := ProgressRenderer{Format: i18n.G("Retrieving image: %s")}
	_, err = op.AddHandler(progress.UpdateOp)
	if err != nil {
		progress.Done("")
		return nil, "", err
	}

	err = cancelableWait(op, &progress)
	if err != nil {
		progress.Done("")
		return nil, "", err
	}
	progress.Done("")

	// Extract the container name
	opInfo, err := op.GetTarget()
	if err != nil {
		return nil, "", err
	}

	containers, ok := opInfo.Resources["containers"]
	if !ok || len(containers) == 0 {
		return nil, "", fmt.Errorf(i18n.G("didn't get any affected image, container or snapshot from server"))
	}

	if len(containers) == 1 && name == "" {
		fields := strings.Split(containers[0], "/")
		name = fields[len(fields)-1]
		fmt.Printf(i18n.G("Container name is: %s")+"\n", name)
	}

	// Validate the network setup
	c.checkNetwork(d, name)

	return d, name, nil
}

func (c *initCmd) guessImage(conf *config.Config, d lxd.ContainerServer, remote string, iremote string, image string) (string, string) {
	if remote != iremote {
		return iremote, image
	}

	fields := strings.SplitN(image, "/", 2)
	_, ok := conf.Remotes[fields[0]]
	if !ok {
		return iremote, image
	}

	_, _, err := d.GetImageAlias(image)
	if err == nil {
		return iremote, image
	}

	_, _, err = d.GetImage(image)
	if err == nil {
		return iremote, image
	}

	if len(fields) == 1 {
		fmt.Fprintf(os.Stderr, i18n.G("The local image '%s' couldn't be found, trying '%s:' instead.")+"\n", image, fields[0])
		return fields[0], "default"
	}

	fmt.Fprintf(os.Stderr, i18n.G("The local image '%s' couldn't be found, trying '%s:%s' instead.")+"\n", image, fields[0], fields[1])
	return fields[0], fields[1]
}

func (c *initCmd) checkNetwork(d lxd.ContainerServer, name string) {
	ct, _, err := d.GetContainer(name)
	if err != nil {
		return
	}

	for _, d := range ct.ExpandedDevices {
		if d["type"] == "nic" {
			return
		}
	}

	fmt.Fprintf(os.Stderr, "\n"+i18n.G("The container you are starting doesn't have any network attached to it.")+"\n")
	fmt.Fprintf(os.Stderr, "  "+i18n.G("To create a new network, use: lxc network create")+"\n")
	fmt.Fprintf(os.Stderr, "  "+i18n.G("To attach a network to a container, use: lxc network attach")+"\n\n")
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	//	proto "github.com/shirou/mqtt"
	proto "github.com/huin/mqtt"
	"log"
	"net"
	"sync"
	"time"
)

// ConnectionErrors is an array of errors corresponding to the
// Connect return codes specified in the specification.
var ConnectionErrors = [6]error{
	nil, // Connection Accepted (not an error)
	errors.New("Connection Refused: unacceptable protocol version"),
	errors.New("Connection Refused: identifier rejected"),
	errors.New("Connection Refused: server unavailable"),
	errors.New("Connection Refused: bad user name or password"),
	errors.New("Connection Refused: not authorized"),
}

const (
	ClientAvailable   uint8 = iota
	ClientUnAvailable       // no PINGACK, no DISCONNECT
	ClientDisconnectedNormally
)

type Connection struct {
	broker            *Broker
	conn              net.Conn
	clientid          string
	storage           Storage
	jobs              chan job
	Done              chan struct{}
	Status            uint8
	TopicList         []string // Subscribed topic list
	LastUpdated       time.Time
	SendingMsgs       *StoredQueue // msgs which not sent
	SentMsgs          *StoredQueue // msgs which already sent
	WillMsg           *proto.Publish
	KeepAliveTimer    uint16
	lastKeepAliveTime time.Time
	Username          string
}

type job struct {
	m           proto.Message
	r           receipt
	storedmsgid string
}

type receipt chan struct{}

// Wait for the receipt to indicate that the job is done.
func (r receipt) wait() {
	// TODO: timeout
	<-r
}

func (c *Connection) handleConnection() {
	defer func() {
		c.conn.Close()
		close(c.jobs)
	}()

	for {
		m, err := proto.DecodeOneMessage(c.conn, nil)
		if err != nil {
			log.Printf("disconnected unexpectedly (%s): %s", c.clientid, err)

			if c.WillMsg != nil {
				log.Printf("Send Will message of %s", c.clientid)
				c.handlePublish(c.WillMsg)
			}

			c.Status = ClientUnAvailable
			return
		}
		log.Printf("incoming: %T from %v", m, c.clientid)
		switch m := m.(type) {
		case *proto.Connect:
			c.handleConnect(m)
		case *proto.Publish:
			c.handlePublish(m)
		case *proto.PubRel:
			c.handlePubRel(m)
		case *proto.PubRec:
			c.handlePubRec(m)
		case *proto.PubComp:
			c.handlePubComp(m)
		case *proto.PingReq:
			c.submit(&proto.PingResp{})
		case *proto.Disconnect:
			c.handleDisconnect(m)
			c.Status = ClientDisconnectedNormally
			return
		case *proto.Subscribe:
			c.handleSubscribe(m)
		case *proto.Unsubscribe:
			c.handleUnsubscribe(m)
		default:
			log.Printf("reader: unknown msg type %T, continue anyway", m)
		}
		continue // loop until Disconnect comes.
	}
}

func (c *Connection) handleSubscribe(m *proto.Subscribe) {
	if m.Header.QosLevel != proto.QosAtLeastOnce {
		// protocol error, silent discarded(not disconnect)
		return
	}
	suback := &proto.SubAck{
		MessageId: m.MessageId,
		TopicsQos: make([]proto.QosLevel, len(m.Topics)),
	}
	for i, tq := range m.Topics {
		// TODO: Handle varying QoS correctly
		c.broker.Subscribe(tq.Topic, c)
		suback.TopicsQos[i] = proto.QosAtMostOnce

		c.TopicList = append(c.TopicList, tq.Topic)
	}
	c.submit(suback)

	// Process retained messages.
	for _, tq := range m.Topics {
		if pubmsg, ok := c.broker.storage.GetRetain(tq.Topic); ok {
			c.submit(pubmsg)
		}
	}
}

func (c *Connection) handleUnsubscribe(m *proto.Unsubscribe) {
	for _, topic := range m.Topics {
		c.broker.Unsubscribe(topic, c)
	}
	ack := &proto.UnsubAck{MessageId: m.MessageId}
	c.submit(ack)
}

func (c *Connection) handleConnect(m *proto.Connect) {
	rc := proto.RetCodeAccepted
	if m.ProtocolName != "MQIsdp" ||
		m.ProtocolVersion != 3 {
		log.Print("reader: reject connection from ", m.ProtocolName, " version ", m.ProtocolVersion)
		rc = proto.RetCodeUnacceptableProtocolVersion
	}

	if m.UsernameFlag {
		if c.broker.Auth(m.Username, m.Password) == false {
			log.Printf("Auth failed: %s", m.Username)
			rc = proto.RetCodeNotAuthorized
		} else {
			c.Username = m.Username
		}

	}

	// Check client id.
	if len(m.ClientId) < 1 || len(m.ClientId) > 23 {
		rc = proto.RetCodeIdentifierRejected
	}
	c.clientid = m.ClientId

	clean := 0
	if m.CleanSession {
		clean = 1
	}

	currrent_c, err := c.storage.MergeClient(c.clientid, c, clean)
	if err != nil {
		c.storage.DeleteClient(c.clientid, c)
		return
	}

	if m.WillFlag {
		header := proto.Header{
			DupFlag:  false,
			QosLevel: m.WillQos,
			Retain:   m.WillRetain,
		}

		c.WillMsg = &proto.Publish{
			Header:    header,
			TopicName: m.WillTopic,
			Payload:   newStringPayload(m.WillMessage),
		}
	}

	connack := &proto.ConnAck{
		ReturnCode: rc,
	}

	currrent_c.submit(connack)

	// close connection if it was a bad connect
	if rc != proto.RetCodeAccepted {
		log.Printf("Connection refused for %v: %v", currrent_c.conn.RemoteAddr(), ConnectionErrors[rc])
		return
	}

	log.Printf("New client connected from %v as %v (c%v, k%v).", currrent_c.conn.RemoteAddr(), currrent_c.clientid, clean, m.KeepAliveTimer)
}

func (c *Connection) handleDisconnect(m *proto.Disconnect) {
	for _, topic := range c.TopicList {
		c.broker.Unsubscribe(topic, c)
	}
	c.storage.DeleteClient(c.clientid, c)
	c.broker.stats.clientDisconnect()
}

func (c *Connection) handlePublish(m *proto.Publish) {
	c.broker.Publish(m)

	if m.Header.Retain {
		c.broker.UpdateRetain(m)
		log.Printf("Publish msg retained: %s", m.TopicName)
	}

	switch m.Header.QosLevel {
	case proto.QosAtLeastOnce:
		// do nothing
	case proto.QosAtMostOnce:
		c.submit(&proto.PubAck{MessageId: m.MessageId})
	case proto.QosExactlyOnce:
		c.submit(&proto.PubRec{MessageId: m.MessageId})
	default:
		log.Printf("Wrong QosLevel on Publish")
	}

	c.broker.stats.messageRecv()
}

func (c *Connection) handlePubRel(m *proto.PubRel) {
	c.submit(&proto.PubComp{MessageId: m.MessageId})
	log.Printf("PubComp sent")
}

func (c *Connection) handlePubRec(m *proto.PubRec) {
	c.submit(&proto.PubRel{MessageId: m.MessageId})
	log.Printf("PubRel sent")
}
func (c *Connection) handlePubComp(m *proto.PubComp) {
	// TODO:
}

// Queue a message; no notification of sending is done.
func (c *Connection) submit(m proto.Message) {
	storedMsgId := ""
	switch pubm := m.(type) {
	case *proto.Publish:
		storedMsgId = c.broker.storage.StoreMsg(c.clientid, pubm)
		log.Printf("msg stored: %s", storedMsgId)
		c.SendingMsgs.Put(storedMsgId)
	}

	log.Printf("%s, %d", c.clientid, c.Status)
	if c.Status != ClientAvailable {
		log.Printf("msg sent to not available client, msg stored: %s", c.clientid)
		return
	}

	j := job{m: m, storedmsgid: storedMsgId}
	select {
	case c.jobs <- j:
	default:
		log.Print(c, ": failed to submit message")
	}
	return
}

// Queue a message, returns a channel that will be readable
// when the message is sent.
func (c *Connection) submitSync(m proto.Message) receipt {
	j := job{m: m, r: make(receipt)}
	c.jobs <- j
	return j.r
}

func (c *Connection) writer() {
	defer func() {
		log.Printf("writer close: %s", c.clientid)
		c.conn.Close()
	}()

	for job := range c.jobs {
		log.Printf("writer begin: %T, %s", job.m, c.clientid)

		// Disconnect msg is used for shutdown writer goroutine.
		if _, ok := job.m.(*proto.Disconnect); ok {
			log.Print("writer: sent disconnect message")
			return
		}

		// TODO: write timeout
		err := job.m.Encode(c.conn)

		if err != nil {
			log.Print("writer: ", err)
			continue // Error does not shutdown Connection, wait re-connect
		}
		// if storedmsgid is set, (QoS 1 or 2,) move to sentQueue
		if job.storedmsgid != "" {
			c.SendingMsgs.Get() // TODO: it ssumes Queue is FIFO
			c.SentMsgs.Put(job.storedmsgid)
			log.Printf("msg %s is moved to SentMsgs", job.storedmsgid)
		}

		if job.r != nil {
			close(job.r)
		}
	}
}

func (c *Connection) Start() {
	go c.handleConnection()
	go c.writer()
}

func NewConnection(b *Broker, conn net.Conn) *Connection {
	c := &Connection{
		broker:      b,
		conn:        conn,
		storage:     b.storage,
		jobs:        make(chan job, b.conf.Queue.SendingQueueLength),
		Status:      ClientAvailable,
		LastUpdated: time.Now(),
		SendingMsgs: NewStoredQueue(b.conf.Queue.SendingQueueLength),
		SentMsgs:    NewStoredQueue(b.conf.Queue.SentQueueLength),
		//		out:      make(chan job, clientQueueLength),
		//		Incoming: make(chan *proto.Publish, clientQueueLength),
		//		done:     make(chan struct{}),
		//		connack:  make(chan *proto.ConnAck),
		//		suback:   make(chan *proto.SubAck),
	}
	return c
}

//
// StoredQueue is a fixed length queue to store messages in a connection.
//
// XXX: should be usecontainer/list ?

type storedQueueNode struct {
	storedMsgId string
	next        *storedQueueNode
}

type StoredQueue struct {
	head  *storedQueueNode
	tail  *storedQueueNode
	count int
	max   int
	lock  *sync.Mutex
}

func NewStoredQueue(max int) *StoredQueue {
	return &StoredQueue{
		lock: &sync.Mutex{},
		max:  max,
	}
}

func (q *storedQueueNode) Next() *storedQueueNode {
	return q.Next()
}

func (q *StoredQueue) Len() int {
	q.lock.Lock()
	defer q.lock.Unlock()

	return q.count
}

func (q *StoredQueue) Put(storedMsgId string) {
	q.lock.Lock()

	n := &storedQueueNode{storedMsgId: storedMsgId}

	if q.tail == nil {
		q.tail = n
		q.head = n
	} else {
		q.tail.next = n
		q.tail = n
	}
	q.count++

	if q.count > q.max {
		q.lock.Unlock()
		q.Get()
		return
	}
	q.lock.Unlock()
}
func (q *StoredQueue) Get() string {
	q.lock.Lock()
	defer q.lock.Unlock()

	n := q.head
	if n == nil {
		return ""
	}

	q.head = n.next

	if q.head == nil {
		q.tail = nil
	}
	q.count--

	return n.storedMsgId
}
<commit_msg>protocol check temporary disabled due to the difference between 3.1 and 3.1.1<commit_after>package main

import (
	"errors"
	//	proto "github.com/shirou/mqtt"
	proto "github.com/huin/mqtt"
	"log"
	"net"
	"sync"
	"time"
)

// ConnectionErrors is an array of errors corresponding to the
// Connect return codes specified in the specification.
var ConnectionErrors = [6]error{
	nil, // Connection Accepted (not an error)
	errors.New("Connection Refused: unacceptable protocol version"),
	errors.New("Connection Refused: identifier rejected"),
	errors.New("Connection Refused: server unavailable"),
	errors.New("Connection Refused: bad user name or password"),
	errors.New("Connection Refused: not authorized"),
}

const (
	ClientAvailable   uint8 = iota
	ClientUnAvailable       // no PINGACK, no DISCONNECT
	ClientDisconnectedNormally
)

type Connection struct {
	broker            *Broker
	conn              net.Conn
	clientid          string
	storage           Storage
	jobs              chan job
	Done              chan struct{}
	Status            uint8
	TopicList         []string // Subscribed topic list
	LastUpdated       time.Time
	SendingMsgs       *StoredQueue // msgs which not sent
	SentMsgs          *StoredQueue // msgs which already sent
	WillMsg           *proto.Publish
	KeepAliveTimer    uint16
	lastKeepAliveTime time.Time
	Username          string
}

type job struct {
	m           proto.Message
	r           receipt
	storedmsgid string
}

type receipt chan struct{}

// Wait for the receipt to indicate that the job is done.
func (r receipt) wait() {
	// TODO: timeout
	<-r
}

func (c *Connection) handleConnection() {
	defer func() {
		c.conn.Close()
		close(c.jobs)
	}()

	for {
		m, err := proto.DecodeOneMessage(c.conn, nil)
		if err != nil {
			log.Printf("disconnected unexpectedly (%s): %s", c.clientid, err)

			if c.WillMsg != nil {
				log.Printf("Send Will message of %s", c.clientid)
				c.handlePublish(c.WillMsg)
			}

			c.Status = ClientUnAvailable
			return
		}
		log.Printf("incoming: %T from %v", m, c.clientid)
		switch m := m.(type) {
		case *proto.Connect:
			c.handleConnect(m)
		case *proto.Publish:
			c.handlePublish(m)
		case *proto.PubRel:
			c.handlePubRel(m)
		case *proto.PubRec:
			c.handlePubRec(m)
		case *proto.PubComp:
			c.handlePubComp(m)
		case *proto.PingReq:
			c.submit(&proto.PingResp{})
		case *proto.Disconnect:
			c.handleDisconnect(m)
			c.Status = ClientDisconnectedNormally
			return
		case *proto.Subscribe:
			c.handleSubscribe(m)
		case *proto.Unsubscribe:
			c.handleUnsubscribe(m)
		default:
			log.Printf("reader: unknown msg type %T, continue anyway", m)
		}
		continue // loop until Disconnect comes.
	}
}

func (c *Connection) handleSubscribe(m *proto.Subscribe) {
	if m.Header.QosLevel != proto.QosAtLeastOnce {
		// protocol error, silent discarded(not disconnect)
		return
	}
	suback := &proto.SubAck{
		MessageId: m.MessageId,
		TopicsQos: make([]proto.QosLevel, len(m.Topics)),
	}
	for i, tq := range m.Topics {
		// TODO: Handle varying QoS correctly
		c.broker.Subscribe(tq.Topic, c)
		suback.TopicsQos[i] = proto.QosAtMostOnce

		c.TopicList = append(c.TopicList, tq.Topic)
	}
	c.submit(suback)

	// Process retained messages.
	for _, tq := range m.Topics {
		if pubmsg, ok := c.broker.storage.GetRetain(tq.Topic); ok {
			c.submit(pubmsg)
		}
	}
}

func (c *Connection) handleUnsubscribe(m *proto.Unsubscribe) {
	for _, topic := range m.Topics {
		c.broker.Unsubscribe(topic, c)
	}
	ack := &proto.UnsubAck{MessageId: m.MessageId}
	c.submit(ack)
}

func (c *Connection) handleConnect(m *proto.Connect) {
	// Protocol check disabled due to difference between 3.1 and 3.1.1
	//	if m.ProtocolName != "MQIsdp" || // should be MQTT?
	//		m.ProtocolVersion != 3 {
	//		log.Print("reader: reject connection from ", m.ProtocolName, " version ", m.ProtocolVersion)
	//      connack := &proto.ConnAck{
	//		    ReturnCode: proto.RetCodeUnacceptableProtocolVersio,
	//	    }
	//		c.submit(connack)
	//		return
	//	}

	if m.UsernameFlag {
		if c.broker.Auth(m.Username, m.Password) == false {
			log.Printf("Auth failed: %s, %s", m.Username, c.conn.RemoteAddr())
			connack := &proto.ConnAck{
				ReturnCode: proto.RetCodeNotAuthorized,
			}
			c.submit(connack)
			return
		} else {
			c.Username = m.Username
		}

	}

	// Check client id.
	if len(m.ClientId) < 1 || len(m.ClientId) > 23 {
		connack := &proto.ConnAck{
			ReturnCode: proto.RetCodeIdentifierRejected,
		}
		c.submit(connack)
		return
	}
	c.clientid = m.ClientId

	clean := 0
	if m.CleanSession {
		clean = 1
	}

	currrent_c, err := c.storage.MergeClient(c.clientid, c, clean)
	if err != nil {
		c.storage.DeleteClient(c.clientid, c)
		return
	}

	if m.WillFlag {
		header := proto.Header{
			DupFlag:  false,
			QosLevel: m.WillQos,
			Retain:   m.WillRetain,
		}

		c.WillMsg = &proto.Publish{
			Header:    header,
			TopicName: m.WillTopic,
			Payload:   newStringPayload(m.WillMessage),
		}
	}

	connack := &proto.ConnAck{
		ReturnCode: proto.RetCodeAccepted,
	}
	currrent_c.submit(connack)

	log.Printf("New client connected from %v as %v (c%v, k%v).", currrent_c.conn.RemoteAddr(), currrent_c.clientid, clean, m.KeepAliveTimer)
}

func (c *Connection) handleDisconnect(m *proto.Disconnect) {
	for _, topic := range c.TopicList {
		c.broker.Unsubscribe(topic, c)
	}
	c.storage.DeleteClient(c.clientid, c)
	c.broker.stats.clientDisconnect()
}

func (c *Connection) handlePublish(m *proto.Publish) {
	c.broker.Publish(m)

	if m.Header.Retain {
		c.broker.UpdateRetain(m)
		log.Printf("Publish msg retained: %s", m.TopicName)
	}

	switch m.Header.QosLevel {
	case proto.QosAtLeastOnce:
		// do nothing
	case proto.QosAtMostOnce:
		c.submit(&proto.PubAck{MessageId: m.MessageId})
	case proto.QosExactlyOnce:
		c.submit(&proto.PubRec{MessageId: m.MessageId})
	default:
		log.Printf("Wrong QosLevel on Publish")
	}

	c.broker.stats.messageRecv()
}

func (c *Connection) handlePubRel(m *proto.PubRel) {
	c.submit(&proto.PubComp{MessageId: m.MessageId})
	log.Printf("PubComp sent")
}

func (c *Connection) handlePubRec(m *proto.PubRec) {
	c.submit(&proto.PubRel{MessageId: m.MessageId})
	log.Printf("PubRel sent")
}
func (c *Connection) handlePubComp(m *proto.PubComp) {
	// TODO:
}

// Queue a message; no notification of sending is done.
func (c *Connection) submit(m proto.Message) {
	storedMsgId := ""
	switch pubm := m.(type) {
	case *proto.Publish:
		storedMsgId = c.broker.storage.StoreMsg(c.clientid, pubm)
		log.Printf("msg stored: %s", storedMsgId)
		c.SendingMsgs.Put(storedMsgId)
	}

	log.Printf("%s, %d", c.clientid, c.Status)
	if c.Status != ClientAvailable {
		log.Printf("msg sent to not available client, msg stored: %s", c.clientid)
		return
	}

	j := job{m: m, storedmsgid: storedMsgId}
	select {
	case c.jobs <- j:
	default:
		log.Print(c, ": failed to submit message")
	}
	return
}

// Queue a message, returns a channel that will be readable
// when the message is sent.
func (c *Connection) submitSync(m proto.Message) receipt {
	j := job{m: m, r: make(receipt)}
	c.jobs <- j
	return j.r
}

func (c *Connection) writer() {
	defer func() {
		log.Printf("writer close: %s", c.clientid)
		c.conn.Close()
	}()

	for job := range c.jobs {
		log.Printf("writer begin: %T, %s", job.m, c.clientid)

		// Disconnect msg is used for shutdown writer goroutine.
		if _, ok := job.m.(*proto.Disconnect); ok {
			log.Print("writer: sent disconnect message")
			return
		}

		// TODO: write timeout
		err := job.m.Encode(c.conn)

		if err != nil {
			log.Print("writer: ", err)
			continue // Error does not shutdown Connection, wait re-connect
		}
		// if storedmsgid is set, (QoS 1 or 2,) move to sentQueue
		if job.storedmsgid != "" {
			c.SendingMsgs.Get() // TODO: it ssumes Queue is FIFO
			c.SentMsgs.Put(job.storedmsgid)
			log.Printf("msg %s is moved to SentMsgs", job.storedmsgid)
		}

		if job.r != nil {
			close(job.r)
		}
	}
}

func (c *Connection) Start() {
	go c.handleConnection()
	go c.writer()
}

func NewConnection(b *Broker, conn net.Conn) *Connection {
	c := &Connection{
		broker:      b,
		conn:        conn,
		storage:     b.storage,
		jobs:        make(chan job, b.conf.Queue.SendingQueueLength),
		Status:      ClientAvailable,
		LastUpdated: time.Now(),
		SendingMsgs: NewStoredQueue(b.conf.Queue.SendingQueueLength),
		SentMsgs:    NewStoredQueue(b.conf.Queue.SentQueueLength),
		//		out:      make(chan job, clientQueueLength),
		//		Incoming: make(chan *proto.Publish, clientQueueLength),
		//		done:     make(chan struct{}),
		//		connack:  make(chan *proto.ConnAck),
		//		suback:   make(chan *proto.SubAck),
	}
	return c
}

//
// StoredQueue is a fixed length queue to store messages in a connection.
//
// XXX: should be usecontainer/list ?

type storedQueueNode struct {
	storedMsgId string
	next        *storedQueueNode
}

type StoredQueue struct {
	head  *storedQueueNode
	tail  *storedQueueNode
	count int
	max   int
	lock  *sync.Mutex
}

func NewStoredQueue(max int) *StoredQueue {
	return &StoredQueue{
		lock: &sync.Mutex{},
		max:  max,
	}
}

func (q *storedQueueNode) Next() *storedQueueNode {
	return q.Next()
}

func (q *StoredQueue) Len() int {
	q.lock.Lock()
	defer q.lock.Unlock()

	return q.count
}

func (q *StoredQueue) Put(storedMsgId string) {
	q.lock.Lock()

	n := &storedQueueNode{storedMsgId: storedMsgId}

	if q.tail == nil {
		q.tail = n
		q.head = n
	} else {
		q.tail.next = n
		q.tail = n
	}
	q.count++

	if q.count > q.max {
		q.lock.Unlock()
		q.Get()
		return
	}
	q.lock.Unlock()
}
func (q *StoredQueue) Get() string {
	q.lock.Lock()
	defer q.lock.Unlock()

	n := q.head
	if n == nil {
		return ""
	}

	q.head = n.next

	if q.head == nil {
		q.tail = nil
	}
	q.count--

	return n.storedMsgId
}
<|endoftext|>


<commit_before>package cloudflare

import (
	"os"
	"testing"

	. "github.com/motain/gocheck"
	"github.com/pearkes/cloudflare/testutil"
)

type S struct {
	client *Client
}

var _ = Suite(&S{})

var testServer = testutil.NewHTTPServer()

func (s *S) SetUpSuite(c *C) {
	testServer.Start()
	var err error
	s.client, err = NewClient("foobar", "foobar")
	s.client.URL = "http://localhost:4444"
	if err != nil {
		panic(err)
	}
}

func (s *S) TearDownTest(c *C) {
	testServer.Flush()
}

func makeClient(t *testing.T) *Client {
	client, err := NewClient("foobaremail", "foobartoken")

	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if client.Token != "foobartoken" {
		t.Fatalf("token not set on client: %s", client.Token)
	}

	if client.Email != "foobaremail" {
		t.Fatalf("email not set on client: %s", client.Token)
	}

	return client
}

func Test_NewClient_env(t *testing.T) {
	os.Setenv("CLOUDFLARE_TOKEN", "bar")
	os.Setenv("CLOUDFLARE_EMAIL", "bar")
	client, err := NewClient("", "")

	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if client.Token != "bar" {
		t.Fatalf("token not set on client: %s", client.Token)
	}

	if client.Email != "bar" {
		t.Fatalf("email not set on client: %s", client.Email)
	}
}

func TestClient_NewRequest(t *testing.T) {
	c := makeClient(t)

	params := map[string]string{
		"foo": "bar",
		"baz": "bar",
	}
	req, err := c.NewRequest(params, "POST", "baz")
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	encoded := req.URL.Query()
	if encoded.Get("foo") != "bar" {
		t.Fatalf("bad: %v", encoded)
	}

	if encoded.Get("baz") != "bar" {
		t.Fatalf("bad: %v", encoded)
	}

	if encoded.Get("baz") != "bar" {
		t.Fatalf("bad: %v", encoded)
	}
	expected := "https://www.cloudflare.com/api_json.html?a=baz&baz=bar&email=foobaremail&foo=bar&tkn=foobartoken"
	if req.URL.String() != expected {
		t.Fatalf("bad base url: %v\n\nexpected: %v", req.URL.String(), expected)
	}

	if req.Method != "POST" {
		t.Fatalf("bad method: %v", req.Method)
	}
}
<commit_msg>api: fix NewRequest action test case<commit_after>package cloudflare

import (
	"os"
	"testing"

	. "github.com/motain/gocheck"
	"github.com/pearkes/cloudflare/testutil"
)

type S struct {
	client *Client
}

var _ = Suite(&S{})

var testServer = testutil.NewHTTPServer()

func (s *S) SetUpSuite(c *C) {
	testServer.Start()
	var err error
	s.client, err = NewClient("foobar", "foobar")
	s.client.URL = "http://localhost:4444"
	if err != nil {
		panic(err)
	}
}

func (s *S) TearDownTest(c *C) {
	testServer.Flush()
}

func makeClient(t *testing.T) *Client {
	client, err := NewClient("foobaremail", "foobartoken")

	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if client.Token != "foobartoken" {
		t.Fatalf("token not set on client: %s", client.Token)
	}

	if client.Email != "foobaremail" {
		t.Fatalf("email not set on client: %s", client.Token)
	}

	return client
}

func Test_NewClient_env(t *testing.T) {
	os.Setenv("CLOUDFLARE_TOKEN", "bar")
	os.Setenv("CLOUDFLARE_EMAIL", "bar")
	client, err := NewClient("", "")

	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if client.Token != "bar" {
		t.Fatalf("token not set on client: %s", client.Token)
	}

	if client.Email != "bar" {
		t.Fatalf("email not set on client: %s", client.Email)
	}
}

func TestClient_NewRequest(t *testing.T) {
	c := makeClient(t)

	params := map[string]string{
		"foo": "bar",
		"baz": "bar",
	}
	req, err := c.NewRequest(params, "POST", "baz")
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	encoded := req.URL.Query()
	if encoded.Get("foo") != "bar" {
		t.Fatalf("bad: %v", encoded)
	}

	if encoded.Get("baz") != "bar" {
		t.Fatalf("bad: %v", encoded)
	}

	if encoded.Get("a") != "baz" {
		t.Fatalf("bad: %v", encoded)
	}
	expected := "https://www.cloudflare.com/api_json.html?a=baz&baz=bar&email=foobaremail&foo=bar&tkn=foobartoken"
	if req.URL.String() != expected {
		t.Fatalf("bad base url: %v\n\nexpected: %v", req.URL.String(), expected)
	}

	if req.Method != "POST" {
		t.Fatalf("bad method: %v", req.Method)
	}
}
<|endoftext|>


<commit_before>package consistent

import (
	"errors"
	"hash/crc32"
	"sort"
	"sync"
)

var ErrNodeNotFound = errors.New("node not found")

type Ring struct {
	Nodes          Nodes
	sortedNodekeys uint32
	sync.Mutex
}

func NewRing() *Ring {
	return &Ring{Nodes: []*Node{}}
}

func (r *Ring) AddNode(id string) {
	r.Lock()
	defer r.Unlock()

	node := &Node{
		Id: id, HashId: hashId(id),
	}

	r.Nodes = append(r.Nodes, node)
	sort.Sort(r.Nodes)
}

func (r *Ring) RemoveNode(id string) error {
	r.Lock()
	defer r.Unlock()

	h := hashId(id)
	i := r.search(h, id)

	if i > r.Nodes.Len() || r.Nodes[i].HashId != h {
		return ErrNodeNotFound
	}

	r.Nodes = append(r.Nodes[:i], r.Nodes[i+1:]...)

	return nil
}

func (r *Ring) Get(id string) string {
	i := r.search(hashId(id), id)
	if i >= r.Nodes.Len() {
		i = 0
	}

	return r.Nodes[i].Id
}

func (r *Ring) search(h uint32, id string) int {
	searchfn := func(i int) bool { return r.Nodes[i].HashId >= h }
	return sort.Search(r.Nodes.Len(), searchfn)
}

//----------------------------------------------------------
// Node
//----------------------------------------------------------

type Node struct {
	Id     string
	HashId uint32
}

type Nodes []*Node

func (n Nodes) Len() int           { return len(n) }
func (n Nodes) Swap(i, j int)      { n[i], n[j] = n[j], n[i] }
func (n Nodes) Less(i, j int) bool { return n[i].HashId < n[j].HashId }

//----------------------------------------------------------
// Helpers
//----------------------------------------------------------

func hashId(key string) uint32 {
	return crc32.ChecksumIEEE([]byte(key))
}
<commit_msg>many fixes<commit_after>package consistent

import (
	"errors"
	"hash/crc32"
	"sort"
	"sync"
)

var ErrNodeNotFound = errors.New("node not found")

type Ring struct {
	Nodes Nodes
	sync.Mutex
}

func NewRing() *Ring {
	return &Ring{Nodes: Nodes{}}
}

func (r *Ring) AddNode(id string) {
	r.Lock()
	defer r.Unlock()

	node := NewNode(id)
	r.Nodes = append(r.Nodes, node)

	sort.Sort(r.Nodes)
}

func (r *Ring) RemoveNode(id string) error {
	r.Lock()
	defer r.Unlock()

	i := r.search(id)
	if i >= r.Nodes.Len() || r.Nodes[i].Id != id {
		return ErrNodeNotFound
	}

	r.Nodes = append(r.Nodes[:i], r.Nodes[i+1:]...)

	return nil
}

func (r *Ring) Get(id string) string {
	i := r.search(id)
	if i >= r.Nodes.Len() {
		i = 0
	}

	return r.Nodes[i].Id
}

func (r *Ring) search(id string) int {
	searchfn := func(i int) bool {
		return r.Nodes[i].HashId >= hashId(id)
	}

	return sort.Search(r.Nodes.Len(), searchfn)
}

//----------------------------------------------------------
// Node
//----------------------------------------------------------

type Node struct {
	Id     string
	HashId uint32
}

func NewNode(id string) *Node {
	return &Node{
		Id:     id,
		HashId: hashId(id),
	}
}

type Nodes []*Node

func (n Nodes) Len() int           { return len(n) }
func (n Nodes) Swap(i, j int)      { n[i], n[j] = n[j], n[i] }
func (n Nodes) Less(i, j int) bool { return n[i].HashId < n[j].HashId }

//----------------------------------------------------------
// Helpers
//----------------------------------------------------------

func hashId(key string) uint32 {
	return crc32.ChecksumIEEE([]byte(key))
}
<|endoftext|>


<commit_before>package consul

import (
	"fmt"
	"github.com/hashicorp/consul/consul/structs"
	"github.com/ugorji/go/codec"
	"io"
	"math/rand"
	"net"
)

type RPCType byte

const (
	rpcConsul RPCType = iota
	rpcRaft
)

// listen is used to listen for incoming RPC connections
func (s *Server) listen() {
	for {
		// Accept a connection
		conn, err := s.rpcListener.Accept()
		if err != nil {
			if s.shutdown {
				return
			}
			s.logger.Printf("[ERR] consul.rpc: failed to accept RPC conn: %v", err)
			continue
		}

		// Track this client
		s.rpcClientLock.Lock()
		s.rpcClients[conn] = struct{}{}
		s.rpcClientLock.Unlock()

		go s.handleConn(conn)
	}
}

// handleConn is used to determine if this is a Raft or
// Consul type RPC connection and invoke the correct handler
func (s *Server) handleConn(conn net.Conn) {
	// Read a single byte
	buf := make([]byte, 1)
	if _, err := conn.Read(buf); err != nil {
		s.logger.Printf("[ERR] consul.rpc: failed to read byte: %v", err)
		conn.Close()
		return
	}

	// Switch on the byte
	switch RPCType(buf[0]) {
	case rpcConsul:
		s.handleConsulConn(conn)

	case rpcRaft:
		s.raftLayer.Handoff(conn)

	default:
		s.logger.Printf("[ERR] consul.rpc: unrecognized RPC byte: %v", buf[0])
		conn.Close()
		return
	}
}

// handleConsulConn is used to service a single Consul RPC connection
func (s *Server) handleConsulConn(conn net.Conn) {
	defer func() {
		conn.Close()
		s.rpcClientLock.Lock()
		delete(s.rpcClients, conn)
		s.rpcClientLock.Unlock()
	}()

	rpcCodec := codec.GoRpc.ServerCodec(conn, &codec.MsgpackHandle{})
	for !s.shutdown {
		if err := s.rpcServer.ServeRequest(rpcCodec); err != nil {
			if err != io.EOF {
				s.logger.Printf("[ERR] consul.rpc: RPC error: %v (%v)", err, conn)
			}
			return
		}
	}
}

// forward is used to forward to a remote DC or to forward to the local leader
// Returns a bool of if forwarding was performed, as well as any error
func (s *Server) forward(method, dc string, args interface{}, reply interface{}) (bool, error) {
	// Handle DC forwarding
	if dc != s.config.Datacenter {
		err := s.forwardDC(method, dc, args, reply)
		return true, err
	}

	// Handle leader forwarding
	if !s.IsLeader() {
		err := s.forwardLeader(method, args, reply)
		return true, err
	}
	return false, nil
}

// forwardLeader is used to forward an RPC call to the leader, or fail if no leader
func (s *Server) forwardLeader(method string, args interface{}, reply interface{}) error {
	leader := s.raft.Leader()
	if leader == nil {
		return structs.ErrNoLeader
	}
	return s.connPool.RPC(leader, method, args, reply)
}

// forwardDC is used to forward an RPC call to a remote DC, or fail if no servers
func (s *Server) forwardDC(method, dc string, args interface{}, reply interface{}) error {
	// Bail if we can't find any servers
	s.remoteLock.RLock()
	servers := s.remoteConsuls[dc]
	if len(servers) == 0 {
		s.remoteLock.RUnlock()
		s.logger.Printf("[WARN] consul.rpc: RPC request for DC '%s', no path found", dc)
		return structs.ErrNoDCPath
	}

	// Select a random addr
	offset := rand.Int31() % int32(len(servers))
	server := servers[offset]
	s.remoteLock.RUnlock()

	// Forward to remote Consul
	return s.connPool.RPC(server, method, args, reply)
}

// raftApply is used to encode a message, run it through raft, and return
// the FSM response along with any errors
func (s *Server) raftApply(t structs.MessageType, msg interface{}) (interface{}, error) {
	buf, err := structs.Encode(t, msg)
	if err != nil {
		return nil, fmt.Errorf("Failed to encode request: %v", err)
	}

	future := s.raft.Apply(buf, 0)
	if err := future.Error(); err != nil {
		return nil, err
	}

	return future.Response(), nil
}
<commit_msg>consul: blockingRPC is a helper method for queries that need to block<commit_after>package consul

import (
	"fmt"
	"github.com/hashicorp/consul/consul/structs"
	"github.com/ugorji/go/codec"
	"io"
	"math/rand"
	"net"
	"time"
)

type RPCType byte

const (
	rpcConsul RPCType = iota
	rpcRaft
)

const (
	maxQueryTime = 600 * time.Second
)

// listen is used to listen for incoming RPC connections
func (s *Server) listen() {
	for {
		// Accept a connection
		conn, err := s.rpcListener.Accept()
		if err != nil {
			if s.shutdown {
				return
			}
			s.logger.Printf("[ERR] consul.rpc: failed to accept RPC conn: %v", err)
			continue
		}

		// Track this client
		s.rpcClientLock.Lock()
		s.rpcClients[conn] = struct{}{}
		s.rpcClientLock.Unlock()

		go s.handleConn(conn)
	}
}

// handleConn is used to determine if this is a Raft or
// Consul type RPC connection and invoke the correct handler
func (s *Server) handleConn(conn net.Conn) {
	// Read a single byte
	buf := make([]byte, 1)
	if _, err := conn.Read(buf); err != nil {
		s.logger.Printf("[ERR] consul.rpc: failed to read byte: %v", err)
		conn.Close()
		return
	}

	// Switch on the byte
	switch RPCType(buf[0]) {
	case rpcConsul:
		s.handleConsulConn(conn)

	case rpcRaft:
		s.raftLayer.Handoff(conn)

	default:
		s.logger.Printf("[ERR] consul.rpc: unrecognized RPC byte: %v", buf[0])
		conn.Close()
		return
	}
}

// handleConsulConn is used to service a single Consul RPC connection
func (s *Server) handleConsulConn(conn net.Conn) {
	defer func() {
		conn.Close()
		s.rpcClientLock.Lock()
		delete(s.rpcClients, conn)
		s.rpcClientLock.Unlock()
	}()

	rpcCodec := codec.GoRpc.ServerCodec(conn, &codec.MsgpackHandle{})
	for !s.shutdown {
		if err := s.rpcServer.ServeRequest(rpcCodec); err != nil {
			if err != io.EOF {
				s.logger.Printf("[ERR] consul.rpc: RPC error: %v (%v)", err, conn)
			}
			return
		}
	}
}

// forward is used to forward to a remote DC or to forward to the local leader
// Returns a bool of if forwarding was performed, as well as any error
func (s *Server) forward(method, dc string, args interface{}, reply interface{}) (bool, error) {
	// Handle DC forwarding
	if dc != s.config.Datacenter {
		err := s.forwardDC(method, dc, args, reply)
		return true, err
	}

	// Handle leader forwarding
	if !s.IsLeader() {
		err := s.forwardLeader(method, args, reply)
		return true, err
	}
	return false, nil
}

// forwardLeader is used to forward an RPC call to the leader, or fail if no leader
func (s *Server) forwardLeader(method string, args interface{}, reply interface{}) error {
	leader := s.raft.Leader()
	if leader == nil {
		return structs.ErrNoLeader
	}
	return s.connPool.RPC(leader, method, args, reply)
}

// forwardDC is used to forward an RPC call to a remote DC, or fail if no servers
func (s *Server) forwardDC(method, dc string, args interface{}, reply interface{}) error {
	// Bail if we can't find any servers
	s.remoteLock.RLock()
	servers := s.remoteConsuls[dc]
	if len(servers) == 0 {
		s.remoteLock.RUnlock()
		s.logger.Printf("[WARN] consul.rpc: RPC request for DC '%s', no path found", dc)
		return structs.ErrNoDCPath
	}

	// Select a random addr
	offset := rand.Int31() % int32(len(servers))
	server := servers[offset]
	s.remoteLock.RUnlock()

	// Forward to remote Consul
	return s.connPool.RPC(server, method, args, reply)
}

// raftApply is used to encode a message, run it through raft, and return
// the FSM response along with any errors
func (s *Server) raftApply(t structs.MessageType, msg interface{}) (interface{}, error) {
	buf, err := structs.Encode(t, msg)
	if err != nil {
		return nil, fmt.Errorf("Failed to encode request: %v", err)
	}

	future := s.raft.Apply(buf, 0)
	if err := future.Error(); err != nil {
		return nil, err
	}

	return future.Response(), nil
}

// blockingRPC is used for queries that need to wait for a
// minimum index. This is used to block and wait for changes.
func (s *Server) blockingRPC(b *structs.BlockingQuery, tables MDBTables, run func() (uint64, error)) error {
	var timeout <-chan time.Time
	var notifyCh chan struct{}

	// Fast path non-blocking
	if b.MinQueryIndex == 0 {
		goto RUN_QUERY
	}

	// Restrict the max query time
	if b.MaxQueryTime > maxQueryTime {
		b.MaxQueryTime = maxQueryTime
	}

	// Ensure a time limit is set if we have an index
	if b.MinQueryIndex > 0 && b.MaxQueryTime == 0 {
		b.MaxQueryTime = maxQueryTime
	}

	// Setup a query timeout
	if b.MaxQueryTime > 0 {
		timeout = time.After(b.MaxQueryTime)
	}

	// Setup a notification channel for changes
SETUP_NOTIFY:
	if b.MinQueryIndex > 0 {
		notifyCh = make(chan struct{}, 1)
		s.fsm.State().Watch(tables, notifyCh)
	}

	// Run the query function
RUN_QUERY:
	idx, err := run()

	// Check for minimum query time
	if err == nil && idx <= b.MinQueryIndex {
		select {
		case <-notifyCh:
			goto SETUP_NOTIFY
		case <-timeout:
		}
	}
	return err
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package version

// Base version information.
//
// This is the fallback data used when version information from git is not
// provided via go ldflags. It provides an approximation of the Kubernetes
// version for ad-hoc builds (e.g. `go build`) that cannot get the version
// information from git.
//
// If you are looking at these fields in the git tree, they look
// strange. They are modified on the fly by the build process. The
// in-tree values are dummy values used for "git archive", which also
// works for GitHub tar downloads.
//
// When releasing a new Kubernetes version, this file is updated by
// build/mark_new_version.sh to reflect the new version, and then a
// git annotated tag (using format vX.Y where X == Major version and Y
// == Minor version) is created to point to the commit that updates
// component-base/version/base.go
var (
	// TODO: Deprecate gitMajor and gitMinor, use only gitVersion
	// instead. First step in deprecation, keep the fields but make
	// them irrelevant. (Next we'll take it out, which may muck with
	// scripts consuming the kubectl version output - but most of
	// these should be looking at gitVersion already anyways.)
	gitMajor string // major version, always numeric
	gitMinor string // minor version, numeric possibly followed by "+"

	// semantic version, derived by build scripts (see
	// https://github.com/kubernetes/community/blob/master/contributors/design-proposals/release/versioning.md
	// for a detailed discussion of this field)
	//
	// TODO: This field is still called "gitVersion" for legacy
	// reasons. For prerelease versions, the build metadata on the
	// semantic version is a git hash, but the version itself is no
	// longer the direct output of "git describe", but a slight
	// translation to be semver compliant.

	// NOTE: The $Format strings are replaced during 'git archive' thanks to the
	// companion .gitattributes file containing 'export-subst' in this same
	// directory.  See also https://git-scm.com/docs/gitattributes
	gitVersion   = "v0.0.0-master+$Format:%h$"
	gitCommit    = "$Format:%H$" // sha1 from git, output of $(git rev-parse HEAD)
	gitTreeState = ""            // state of git tree, either "clean" or "dirty"

	buildDate = "1970-01-01T00:00:00Z" // build date in ISO8601 format, output of $(date -u +'%Y-%m-%dT%H:%M:%SZ')
)
<commit_msg>Switch gitVersion format to non-abbreviated hash<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package version

// Base version information.
//
// This is the fallback data used when version information from git is not
// provided via go ldflags. It provides an approximation of the Kubernetes
// version for ad-hoc builds (e.g. `go build`) that cannot get the version
// information from git.
//
// If you are looking at these fields in the git tree, they look
// strange. They are modified on the fly by the build process. The
// in-tree values are dummy values used for "git archive", which also
// works for GitHub tar downloads.
//
// When releasing a new Kubernetes version, this file is updated by
// build/mark_new_version.sh to reflect the new version, and then a
// git annotated tag (using format vX.Y where X == Major version and Y
// == Minor version) is created to point to the commit that updates
// component-base/version/base.go
var (
	// TODO: Deprecate gitMajor and gitMinor, use only gitVersion
	// instead. First step in deprecation, keep the fields but make
	// them irrelevant. (Next we'll take it out, which may muck with
	// scripts consuming the kubectl version output - but most of
	// these should be looking at gitVersion already anyways.)
	gitMajor string // major version, always numeric
	gitMinor string // minor version, numeric possibly followed by "+"

	// semantic version, derived by build scripts (see
	// https://github.com/kubernetes/community/blob/master/contributors/design-proposals/release/versioning.md
	// for a detailed discussion of this field)
	//
	// TODO: This field is still called "gitVersion" for legacy
	// reasons. For prerelease versions, the build metadata on the
	// semantic version is a git hash, but the version itself is no
	// longer the direct output of "git describe", but a slight
	// translation to be semver compliant.

	// NOTE: The $Format strings are replaced during 'git archive' thanks to the
	// companion .gitattributes file containing 'export-subst' in this same
	// directory.  See also https://git-scm.com/docs/gitattributes
	gitVersion   = "v0.0.0-master+$Format:%H$"
	gitCommit    = "$Format:%H$" // sha1 from git, output of $(git rev-parse HEAD)
	gitTreeState = ""            // state of git tree, either "clean" or "dirty"

	buildDate = "1970-01-01T00:00:00Z" // build date in ISO8601 format, output of $(date -u +'%Y-%m-%dT%H:%M:%SZ')
)
<|endoftext|>


<commit_before>package redis

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"github.com/cloudfoundry-incubator/cf-lager"
	"github.com/pivotal-cf/cf-redis-broker/brokerconfig"
	"github.com/pivotal-cf/cf-redis-broker/redis/client"
	"github.com/pivotal-cf/cf-redis-broker/redisconf"
	"github.com/pivotal-cf/cf-redis-broker/s3bucket"
	"github.com/pivotal-golang/lager"
)

type Backup struct {
	Config *brokerconfig.Config
}

func (backup Backup) Create(instanceID string) error {
	s3Client := s3bucket.NewClient(
		backup.Config.RedisConfiguration.BackupConfiguration.EndpointUrl,
		backup.Config.RedisConfiguration.BackupConfiguration.S3Region,
		backup.Config.RedisConfiguration.BackupConfiguration.AccessKeyId,
		backup.Config.RedisConfiguration.BackupConfiguration.SecretAccessKey,
	)

	bucket, err := s3Client.GetOrCreate(backup.Config.RedisConfiguration.BackupConfiguration.BucketName)
	if err != nil {
		log.Fatal(err)
	}

	return backup.backupInstance(instanceID, bucket)
}

func (backup Backup) backupInstance(instanceID string, bucket s3bucket.Bucket) error {
	logger := cf_lager.New("backup")

	pathToInstanceDirectory := filepath.Join(backup.Config.RedisConfiguration.InstanceDataDirectory, instanceID)
	if !fileExists(pathToInstanceDirectory) {
		logger.Info("instance directory not found, skipping instance backup", lager.Data{
			"Local file": pathToInstanceDirectory,
		})
		return nil
	}

	err := backup.saveAndWaitUntilFinished(instanceID)
	if err != nil {
		return err
	}

	pathToRdbFile := filepath.Join(backup.Config.RedisConfiguration.InstanceDataDirectory, instanceID, "db", "dump.rdb")
	if !fileExists(pathToRdbFile) {
		logger.Info("dump.rb not found, skipping instance backup", lager.Data{
			"Local file": pathToRdbFile,
		})
		return nil
	}

	rdbBytes, err := ioutil.ReadFile(pathToRdbFile)
	if err != nil {
		return err
	}

	remotePath := fmt.Sprintf("%s/%s", backup.Config.RedisConfiguration.BackupConfiguration.Path, instanceID)

	logger.Info("Backing up instance", lager.Data{
		"Local file":  pathToRdbFile,
		"Remote file": remotePath,
	})

	return bucket.Upload(rdbBytes, remotePath)
}

func fileExists(path string) bool {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return false
	}
	return true
}

func (backup Backup) saveAndWaitUntilFinished(instanceID string) error {
	client, err := backup.buildRedisClient(instanceID)
	if err != nil {
		return err
	}

	return client.CreateSnapshot(backup.Config.RedisConfiguration.BackupConfiguration.BGSaveTimeoutSeconds)
}

func (backup Backup) buildRedisClient(instanceID string) (*client.Client, error) {

	localRepo := LocalRepository{RedisConf: backup.Config.RedisConfiguration}
	instance, err := localRepo.FindByID(instanceID)
	if err != nil {
		return nil, err
	}

	instanceConf, err := redisconf.Load(localRepo.InstanceConfigPath(instanceID))
	if err != nil {
		return nil, err
	}

	return client.Connect(instance.Host, uint(instance.Port), instance.Password, instanceConf)
}
<commit_msg>Extracting private methods to make the code a little more readable<commit_after>package redis

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"github.com/cloudfoundry-incubator/cf-lager"
	"github.com/pivotal-cf/cf-redis-broker/brokerconfig"
	"github.com/pivotal-cf/cf-redis-broker/redis/client"
	"github.com/pivotal-cf/cf-redis-broker/redisconf"
	"github.com/pivotal-cf/cf-redis-broker/s3bucket"
	"github.com/pivotal-golang/lager"
)

type Backup struct {
	Config *brokerconfig.Config
}

func (backup Backup) Create(instanceID string) error {
	logger := cf_lager.New("backup")

	bucket := backup.createBucket()

	if !backup.validateInstanceDirectoryIsPresentFor(instanceID, logger) {
		return nil
	}

	err := backup.createSnapshot(instanceID)
	if err != nil {
		return err
	}

	pathToRdbFile := filepath.Join(backup.Config.RedisConfiguration.InstanceDataDirectory, instanceID, "db", "dump.rdb")

	if !backup.validateBackupFileCreatedFor(pathToRdbFile, logger) {
		return nil
	}

	return backup.uploadToS3(instanceID, pathToRdbFile, bucket, logger)
}

func fileExists(path string) bool {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return false
	}
	return true
}

func (backup Backup) createSnapshot(instanceID string) error {
	client, err := backup.buildRedisClient(instanceID)
	if err != nil {
		return err
	}

	return client.CreateSnapshot(backup.Config.RedisConfiguration.BackupConfiguration.BGSaveTimeoutSeconds)
}

func (backup Backup) buildRedisClient(instanceID string) (*client.Client, error) {

	localRepo := LocalRepository{RedisConf: backup.Config.RedisConfiguration}
	instance, err := localRepo.FindByID(instanceID)
	if err != nil {
		return nil, err
	}

	instanceConf, err := redisconf.Load(localRepo.InstanceConfigPath(instanceID))
	if err != nil {
		return nil, err
	}

	return client.Connect(instance.Host, uint(instance.Port), instance.Password, instanceConf)
}

func (backup Backup) validateInstanceDirectoryIsPresentFor(instanceID string, logger lager.Logger) bool {
	pathToInstanceDirectory := filepath.Join(backup.Config.RedisConfiguration.InstanceDataDirectory, instanceID)
	if !fileExists(pathToInstanceDirectory) {
		logger.Info("instance directory not found, skipping instance backup", lager.Data{
			"Local file": pathToInstanceDirectory,
		})
		return false
	}
	return true
}

func (backup Backup) validateBackupFileCreatedFor(pathToRdbFile string, logger lager.Logger) bool {
	if !fileExists(pathToRdbFile) {
		logger.Info("dump.rb not found, skipping instance backup", lager.Data{
			"Local file": pathToRdbFile,
		})
		return false
	}
	return true
}

func (backup Backup) uploadToS3(instanceID, pathToRdbFile string, bucket s3bucket.Bucket, logger lager.Logger) error {
	rdbBytes, err := ioutil.ReadFile(pathToRdbFile)
	if err != nil {
		return err
	}

	remotePath := fmt.Sprintf("%s/%s", backup.Config.RedisConfiguration.BackupConfiguration.Path, instanceID)

	logger.Info("Backing up instance", lager.Data{
		"Local file":  pathToRdbFile,
		"Remote file": remotePath,
	})

	return bucket.Upload(rdbBytes, remotePath)
}

func (backup Backup) createBucket() s3bucket.Bucket {
	s3Client := s3bucket.NewClient(
		backup.Config.RedisConfiguration.BackupConfiguration.EndpointUrl,
		backup.Config.RedisConfiguration.BackupConfiguration.S3Region,
		backup.Config.RedisConfiguration.BackupConfiguration.AccessKeyId,
		backup.Config.RedisConfiguration.BackupConfiguration.SecretAccessKey,
	)

	bucket, err := s3Client.GetOrCreate(backup.Config.RedisConfiguration.BackupConfiguration.BucketName)
	if err != nil {
		log.Fatal(err)
	}
	return bucket
}
<|endoftext|>