<commit_before>// Copyright 2017 Google Inc.
//
// 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
//
//     https://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 server

import (
	"bytes"
	"context"
	"crypto"
	"crypto/x509"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"sort"
	"time"

	log "github.com/golang/glog"
	"github.com/golang/protobuf/proto"

	"github.com/google/fleetspeak/fleetspeak/src/common"
	"github.com/google/fleetspeak/fleetspeak/src/server/authorizer"
	"github.com/google/fleetspeak/fleetspeak/src/server/comms"
	"github.com/google/fleetspeak/fleetspeak/src/server/db"
	"github.com/google/fleetspeak/fleetspeak/src/server/internal/signatures"
	"github.com/google/fleetspeak/fleetspeak/src/server/stats"

	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
)

const maxMessagesPerContact = 100
const processingChunkSize = 10

type commsContext struct {
	s *Server
}

func (c commsContext) InitializeConnection(ctx context.Context, addr net.Addr, key crypto.PublicKey, wcd *fspb.WrappedContactData) (*comms.ConnectionInfo, *fspb.ContactData, error) {
	id, err := common.MakeClientID(key)
	if err != nil {
		return nil, nil, err
	}

	contactInfo := authorizer.ContactInfo{
		ID:           id,
		ContactSize:  len(wcd.ContactData),
		ClientLabels: wcd.ClientLabels,
	}
	if !c.s.authorizer.Allow2(addr, contactInfo) {
		return nil, nil, comms.NotAuthorizedError
	}

	ci, err := c.getClientInfo(ctx, id)
	if err != nil {
		return nil, nil, err
	}

	res := comms.ConnectionInfo{
		Addr: addr,
	}
	if ci == nil {
		res.AuthClientInfo.New = true
	} else {
		res.Client = *ci
		res.AuthClientInfo.Labels = ci.Labels
	}

	if !c.s.authorizer.Allow3(addr, contactInfo, res.AuthClientInfo) {
		return nil, nil, comms.NotAuthorizedError
	}

	if ci == nil {
		if err := c.addClient(ctx, id, key); err != nil {
			return nil, nil, err
		}
		res.Client.ID = id
		res.Client.Key = key
		// Set initial labels for the client according to the contact data. Going
		// forward, labels will be adjusted when the client sends a ClientInfo
		// message (see system_service.go).
		for _, l := range wcd.ClientLabels {
			cl := &fspb.Label{ServiceName: "client", Label: l}

			// Ignore errors - if this fails, the first ClientInfo message will try again
			// in a context where we can retry easily.
			c.s.dataStore.AddClientLabel(ctx, id, cl)

			res.Client.Labels = append(res.Client.Labels, cl)
		}
		res.AuthClientInfo.Labels = res.Client.Labels
	}

	sigs, err := signatures.ValidateWrappedContactData(id, wcd)
	if err != nil {
		return nil, nil, err
	}
	accept, validationInfo := c.s.authorizer.Allow4(
		addr,
		contactInfo,
		res.AuthClientInfo,
		sigs)
	if !accept {
		return nil, nil, errors.New("contact not authorized")
	}

	var cd fspb.ContactData
	if err = proto.Unmarshal(wcd.ContactData, &cd); err != nil {
		return nil, nil, fmt.Errorf("unable to parse contact_data: %v", err)
	}
	if len(cd.Messages) > maxMessagesPerContact {
		return nil, nil, fmt.Errorf("contact_data contains %d messages, only %d allowed", len(cd.Messages), maxMessagesPerContact)
	}
	res.NonceReceived = cd.SequencingNonce
	toSend := fspb.ContactData{SequencingNonce: uint64(rand.Int63())}
	res.NonceSent = toSend.SequencingNonce
	res.ContactID, err = c.s.dataStore.RecordClientContact(ctx,
		db.ContactData{
			ClientID:      id,
			NonceSent:     toSend.SequencingNonce,
			NonceReceived: cd.SequencingNonce,
			Addr:          addr.String(),
			ClientClock:   cd.ClientClock,
		})
	if err != nil {
		return nil, nil, err
	}

	err = c.handleMessagesFromClient(ctx, &res.Client, res.ContactID, &cd, validationInfo)
	if err != nil {
		return nil, nil, err
	}

	toSend.Messages, err = c.FindMessagesForClient(ctx, &res.Client, res.ContactID, 100)
	if err != nil {
		return nil, nil, err
	}

	res.AuthClientInfo.New = false
	return &res, &toSend, nil
}

// getClientInfo loads basic information about a client. Returns nil if the client does
// not exist in the datastore.
func (c commsContext) getClientInfo(ctx context.Context, id common.ClientID) (*comms.ClientInfo, error) {
	cld, cacheHit, err := c.s.clientCache.GetOrRead(ctx, id, c.s.dataStore)
	if err != nil {
		if c.s.dataStore.IsNotFound(err) {
			return nil, nil
		}
		return nil, err
	}

	k, err := x509.ParsePKIXPublicKey(cld.Key)
	if err != nil {
		return nil, err
	}
	return &comms.ClientInfo{
		ID:          id,
		Key:         k,
		Labels:      cld.Labels,
		Blacklisted: cld.Blacklisted,
		Cached:      cacheHit}, nil
}

func (c commsContext) HandleMessagesFromClient(ctx context.Context, info *comms.ConnectionInfo, wcd *fspb.WrappedContactData) error {
	sigs, err := signatures.ValidateWrappedContactData(info.Client.ID, wcd)
	if err != nil {
		return err
	}

	accept, validationInfo := c.s.authorizer.Allow4(
		info.Addr,
		authorizer.ContactInfo{
			ID:           info.Client.ID,
			ContactSize:  len(wcd.ContactData),
			ClientLabels: wcd.ClientLabels,
		},
		info.AuthClientInfo,
		sigs)
	if !accept {
		return errors.New("contact not authorized")
	}

	var cd fspb.ContactData
	if err = proto.Unmarshal(wcd.ContactData, &cd); err != nil {
		return fmt.Errorf("unable to parse contact_data: %v", err)
	}
	if len(cd.Messages) > maxMessagesPerContact {
		return fmt.Errorf("contact_data contains %d messages, only %d allowed", len(cd.Messages), maxMessagesPerContact)
	}

	err = c.handleMessagesFromClient(ctx, &info.Client, info.ContactID, &cd, validationInfo)
	if err != nil {
		return err
	}
	return nil
}

func (c commsContext) GetMessagesForClient(ctx context.Context, info *comms.ConnectionInfo) (*fspb.ContactData, error) {
	toSend := fspb.ContactData{
		SequencingNonce: info.NonceSent,
	}
	var err error
	toSend.Messages, err = c.FindMessagesForClient(ctx, &info.Client, info.ContactID, 100)
	if err != nil || len(toSend.Messages) == 0 {
		return nil, err
	}
	return &toSend, nil
}

// addClient adds a new client to the system.
func (c commsContext) addClient(ctx context.Context, id common.ClientID, key crypto.PublicKey) error {
	k, err := x509.MarshalPKIXPublicKey(key)
	if err != nil {
		return err
	}
	if err := c.s.dataStore.AddClient(ctx, id, &db.ClientData{Key: k}); err != nil {
		return err
	}
	return nil
}

// FindMessagesForClient finds unprocessed messages for a given client and
// reserves them for processing.
func (c commsContext) FindMessagesForClient(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID, maxMessages int) ([]*fspb.Message, error) {
	if info.Blacklisted {
		log.Warningf("Contact from blacklisted id [%v], creating RekeyRequest.", info.ID)
		m, err := c.MakeBlacklistMessage(ctx, info, contactID)
		return []*fspb.Message{m}, err
	}
	msgs, err := c.s.dataStore.ClientMessagesForProcessing(ctx, info.ID, maxMessages)
	if err != nil {
		if len(msgs) == 0 {
			return nil, err
		}
		log.Warning("Got %v messages along with error, continuing: %v", len(msgs), err)
	}

	// If the client recently contacted us, the broadcast situation is unlikely to
	// have changed, so we skip checking for broadcasts. To keep this from delaying
	// broadcast distribution, the broadcast manager clears the client cache when it
	// finds more broadcasts.
	if !info.Cached {
		bms, err := c.s.broadcastManager.MakeBroadcastMessagesForClient(ctx, info.ID, info.Labels)
		if err != nil {
			return nil, err
		}
		msgs = append(msgs, bms...)
	}

	if len(msgs) == 0 {
		return msgs, nil
	}

	mids := make([]common.MessageID, 0, len(msgs))
	for _, m := range msgs {
		id, err := common.BytesToMessageID(m.MessageId)
		if err != nil {
			return nil, err
		}
		mids = append(mids, id)
	}
	err = c.s.dataStore.LinkMessagesToContact(ctx, contactID, mids)
	if err != nil {
		return nil, err
	}
	return msgs, nil
}

func (c commsContext) MakeBlacklistMessage(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID) (*fspb.Message, error) {
	mid, err := common.RandomMessageID()
	if err != nil {
		return nil, fmt.Errorf("unable to create message id: %v", err)
	}
	msg := &fspb.Message{
		MessageId: mid.Bytes(),
		Source: &fspb.Address{
			ServiceName: "system",
		},
		Destination: &fspb.Address{
			ServiceName: "system",
			ClientId:    info.ID.Bytes(),
		},
		MessageType:  "RekeyRequest",
		CreationTime: db.NowProto(),
	}
	if err = c.s.dataStore.StoreMessages(ctx, []*fspb.Message{msg}, contactID); err != nil {
		return nil, fmt.Errorf("unable to store RekeyRequest: %v", err)
	}
	return msg, nil
}

func (c commsContext) validateMessageFromClient(id common.ClientID, m *fspb.Message, validationInfo *fspb.ValidationInfo) error {
	if m.Destination == nil {
		return fmt.Errorf("message must have Destination")
	}
	if m.Destination.ClientId != nil {
		return fmt.Errorf("cannot send a message directly to another client [%v]", m.Destination.ClientId)
	}
	if m.Source == nil || m.Source.ServiceName == "" {
		return fmt.Errorf("message must have a source with a ServiceName, got: %v", m.Source)
	}
	if m.SourceMessageId == nil {
		return fmt.Errorf("source message id cannot be empty")
	}

	m.Source.ClientId = id.Bytes()
	m.ValidationInfo = validationInfo
	m.MessageId = common.MakeMessageID(m.Source, m.SourceMessageId).Bytes()
	return nil
}

// handleMessagesFromClient processes a block of messages from a particular
// client. It saves them to the database, associates them with the contact
// identified by contactTime, and processes them.
func (c commsContext) handleMessagesFromClient(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID, received *fspb.ContactData, validationInfo *fspb.ValidationInfo) error {
	msgs := make([]*fspb.Message, 0, len(received.Messages))
	for _, m := range received.Messages {
		err := c.validateMessageFromClient(info.ID, m, validationInfo)
		if err != nil {
			log.Errorf("Dropping invalid message from [%v]: %v", info.ID, err)
			continue
		}
		msgs = append(msgs, m)
	}
	if len(msgs) == 0 {
		return nil
	}

	sort.Slice(msgs, func(a, b int) bool {
		return bytes.Compare(msgs[a].MessageId, msgs[b].MessageId) == -1
	})

	for {
		if len(msgs) <= processingChunkSize {
			return c.s.serviceConfig.HandleNewMessages(ctx, msgs, contactID)
		}

		if err := c.s.serviceConfig.HandleNewMessages(ctx, msgs[:processingChunkSize], contactID); err != nil {
			return err
		}
		msgs = msgs[processingChunkSize:]
	}
}

// ReadFile returns the data and modification time of file. Caller is
// responsible for closing data.
//
// Calls to data are permitted to fail if ctx is canceled or expired.
func (c commsContext) ReadFile(ctx context.Context, service, name string) (data db.ReadSeekerCloser, modtime time.Time, err error) {
	return c.s.dataStore.ReadFile(ctx, service, name)
}

// IsNotFound returns whether an error returned by ReadFile indicates that the
// file was not found.
func (c commsContext) IsNotFound(err error) bool {
	return c.s.dataStore.IsNotFound(err)
}

// StatsCollector returns the stats.Collector used by the Fleetspeak
// system. Access is provided to allow collection of stats relating to the
// client communication.
func (c commsContext) StatsCollector() stats.Collector {
	return c.s.statsCollector
}

func (c commsContext) Authorizer() authorizer.Authorizer {
	return c.s.authorizer
}
<commit_msg>Review fixup - NotAuthorizedError.<commit_after>// Copyright 2017 Google Inc.
//
// 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
//
//     https://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 server

import (
	"bytes"
	"context"
	"crypto"
	"crypto/x509"
	"fmt"
	"math/rand"
	"net"
	"sort"
	"time"

	log "github.com/golang/glog"
	"github.com/golang/protobuf/proto"

	"github.com/google/fleetspeak/fleetspeak/src/common"
	"github.com/google/fleetspeak/fleetspeak/src/server/authorizer"
	"github.com/google/fleetspeak/fleetspeak/src/server/comms"
	"github.com/google/fleetspeak/fleetspeak/src/server/db"
	"github.com/google/fleetspeak/fleetspeak/src/server/internal/signatures"
	"github.com/google/fleetspeak/fleetspeak/src/server/stats"

	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
)

const maxMessagesPerContact = 100
const processingChunkSize = 10

type commsContext struct {
	s *Server
}

func (c commsContext) InitializeConnection(ctx context.Context, addr net.Addr, key crypto.PublicKey, wcd *fspb.WrappedContactData) (*comms.ConnectionInfo, *fspb.ContactData, error) {
	id, err := common.MakeClientID(key)
	if err != nil {
		return nil, nil, err
	}

	contactInfo := authorizer.ContactInfo{
		ID:           id,
		ContactSize:  len(wcd.ContactData),
		ClientLabels: wcd.ClientLabels,
	}
	if !c.s.authorizer.Allow2(addr, contactInfo) {
		return nil, nil, comms.NotAuthorizedError
	}

	ci, err := c.getClientInfo(ctx, id)
	if err != nil {
		return nil, nil, err
	}

	res := comms.ConnectionInfo{
		Addr: addr,
	}
	if ci == nil {
		res.AuthClientInfo.New = true
	} else {
		res.Client = *ci
		res.AuthClientInfo.Labels = ci.Labels
	}

	if !c.s.authorizer.Allow3(addr, contactInfo, res.AuthClientInfo) {
		return nil, nil, comms.NotAuthorizedError
	}

	if ci == nil {
		if err := c.addClient(ctx, id, key); err != nil {
			return nil, nil, err
		}
		res.Client.ID = id
		res.Client.Key = key
		// Set initial labels for the client according to the contact data. Going
		// forward, labels will be adjusted when the client sends a ClientInfo
		// message (see system_service.go).
		for _, l := range wcd.ClientLabels {
			cl := &fspb.Label{ServiceName: "client", Label: l}

			// Ignore errors - if this fails, the first ClientInfo message will try again
			// in a context where we can retry easily.
			c.s.dataStore.AddClientLabel(ctx, id, cl)

			res.Client.Labels = append(res.Client.Labels, cl)
		}
		res.AuthClientInfo.Labels = res.Client.Labels
	}

	sigs, err := signatures.ValidateWrappedContactData(id, wcd)
	if err != nil {
		return nil, nil, err
	}
	accept, validationInfo := c.s.authorizer.Allow4(
		addr,
		contactInfo,
		res.AuthClientInfo,
		sigs)
	if !accept {
		return nil, nil, comms.NotAuthorizedError
	}

	var cd fspb.ContactData
	if err = proto.Unmarshal(wcd.ContactData, &cd); err != nil {
		return nil, nil, fmt.Errorf("unable to parse contact_data: %v", err)
	}
	if len(cd.Messages) > maxMessagesPerContact {
		return nil, nil, fmt.Errorf("contact_data contains %d messages, only %d allowed", len(cd.Messages), maxMessagesPerContact)
	}
	res.NonceReceived = cd.SequencingNonce
	toSend := fspb.ContactData{SequencingNonce: uint64(rand.Int63())}
	res.NonceSent = toSend.SequencingNonce
	res.ContactID, err = c.s.dataStore.RecordClientContact(ctx,
		db.ContactData{
			ClientID:      id,
			NonceSent:     toSend.SequencingNonce,
			NonceReceived: cd.SequencingNonce,
			Addr:          addr.String(),
			ClientClock:   cd.ClientClock,
		})
	if err != nil {
		return nil, nil, err
	}

	err = c.handleMessagesFromClient(ctx, &res.Client, res.ContactID, &cd, validationInfo)
	if err != nil {
		return nil, nil, err
	}

	toSend.Messages, err = c.FindMessagesForClient(ctx, &res.Client, res.ContactID, 100)
	if err != nil {
		return nil, nil, err
	}

	res.AuthClientInfo.New = false
	return &res, &toSend, nil
}

// getClientInfo loads basic information about a client. Returns nil if the client does
// not exist in the datastore.
func (c commsContext) getClientInfo(ctx context.Context, id common.ClientID) (*comms.ClientInfo, error) {
	cld, cacheHit, err := c.s.clientCache.GetOrRead(ctx, id, c.s.dataStore)
	if err != nil {
		if c.s.dataStore.IsNotFound(err) {
			return nil, nil
		}
		return nil, err
	}

	k, err := x509.ParsePKIXPublicKey(cld.Key)
	if err != nil {
		return nil, err
	}
	return &comms.ClientInfo{
		ID:          id,
		Key:         k,
		Labels:      cld.Labels,
		Blacklisted: cld.Blacklisted,
		Cached:      cacheHit}, nil
}

func (c commsContext) HandleMessagesFromClient(ctx context.Context, info *comms.ConnectionInfo, wcd *fspb.WrappedContactData) error {
	sigs, err := signatures.ValidateWrappedContactData(info.Client.ID, wcd)
	if err != nil {
		return err
	}

	accept, validationInfo := c.s.authorizer.Allow4(
		info.Addr,
		authorizer.ContactInfo{
			ID:           info.Client.ID,
			ContactSize:  len(wcd.ContactData),
			ClientLabels: wcd.ClientLabels,
		},
		info.AuthClientInfo,
		sigs)
	if !accept {
		return comms.NotAuthorizedError
	}

	var cd fspb.ContactData
	if err = proto.Unmarshal(wcd.ContactData, &cd); err != nil {
		return fmt.Errorf("unable to parse contact_data: %v", err)
	}
	if len(cd.Messages) > maxMessagesPerContact {
		return fmt.Errorf("contact_data contains %d messages, only %d allowed", len(cd.Messages), maxMessagesPerContact)
	}

	err = c.handleMessagesFromClient(ctx, &info.Client, info.ContactID, &cd, validationInfo)
	if err != nil {
		return err
	}
	return nil
}

func (c commsContext) GetMessagesForClient(ctx context.Context, info *comms.ConnectionInfo) (*fspb.ContactData, error) {
	toSend := fspb.ContactData{
		SequencingNonce: info.NonceSent,
	}
	var err error
	toSend.Messages, err = c.FindMessagesForClient(ctx, &info.Client, info.ContactID, 100)
	if err != nil || len(toSend.Messages) == 0 {
		return nil, err
	}
	return &toSend, nil
}

// addClient adds a new client to the system.
func (c commsContext) addClient(ctx context.Context, id common.ClientID, key crypto.PublicKey) error {
	k, err := x509.MarshalPKIXPublicKey(key)
	if err != nil {
		return err
	}
	if err := c.s.dataStore.AddClient(ctx, id, &db.ClientData{Key: k}); err != nil {
		return err
	}
	return nil
}

// FindMessagesForClient finds unprocessed messages for a given client and
// reserves them for processing.
func (c commsContext) FindMessagesForClient(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID, maxMessages int) ([]*fspb.Message, error) {
	if info.Blacklisted {
		log.Warningf("Contact from blacklisted id [%v], creating RekeyRequest.", info.ID)
		m, err := c.MakeBlacklistMessage(ctx, info, contactID)
		return []*fspb.Message{m}, err
	}
	msgs, err := c.s.dataStore.ClientMessagesForProcessing(ctx, info.ID, maxMessages)
	if err != nil {
		if len(msgs) == 0 {
			return nil, err
		}
		log.Warning("Got %v messages along with error, continuing: %v", len(msgs), err)
	}

	// If the client recently contacted us, the broadcast situation is unlikely to
	// have changed, so we skip checking for broadcasts. To keep this from delaying
	// broadcast distribution, the broadcast manager clears the client cache when it
	// finds more broadcasts.
	if !info.Cached {
		bms, err := c.s.broadcastManager.MakeBroadcastMessagesForClient(ctx, info.ID, info.Labels)
		if err != nil {
			return nil, err
		}
		msgs = append(msgs, bms...)
	}

	if len(msgs) == 0 {
		return msgs, nil
	}

	mids := make([]common.MessageID, 0, len(msgs))
	for _, m := range msgs {
		id, err := common.BytesToMessageID(m.MessageId)
		if err != nil {
			return nil, err
		}
		mids = append(mids, id)
	}
	err = c.s.dataStore.LinkMessagesToContact(ctx, contactID, mids)
	if err != nil {
		return nil, err
	}
	return msgs, nil
}

func (c commsContext) MakeBlacklistMessage(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID) (*fspb.Message, error) {
	mid, err := common.RandomMessageID()
	if err != nil {
		return nil, fmt.Errorf("unable to create message id: %v", err)
	}
	msg := &fspb.Message{
		MessageId: mid.Bytes(),
		Source: &fspb.Address{
			ServiceName: "system",
		},
		Destination: &fspb.Address{
			ServiceName: "system",
			ClientId:    info.ID.Bytes(),
		},
		MessageType:  "RekeyRequest",
		CreationTime: db.NowProto(),
	}
	if err = c.s.dataStore.StoreMessages(ctx, []*fspb.Message{msg}, contactID); err != nil {
		return nil, fmt.Errorf("unable to store RekeyRequest: %v", err)
	}
	return msg, nil
}

func (c commsContext) validateMessageFromClient(id common.ClientID, m *fspb.Message, validationInfo *fspb.ValidationInfo) error {
	if m.Destination == nil {
		return fmt.Errorf("message must have Destination")
	}
	if m.Destination.ClientId != nil {
		return fmt.Errorf("cannot send a message directly to another client [%v]", m.Destination.ClientId)
	}
	if m.Source == nil || m.Source.ServiceName == "" {
		return fmt.Errorf("message must have a source with a ServiceName, got: %v", m.Source)
	}
	if m.SourceMessageId == nil {
		return fmt.Errorf("source message id cannot be empty")
	}

	m.Source.ClientId = id.Bytes()
	m.ValidationInfo = validationInfo
	m.MessageId = common.MakeMessageID(m.Source, m.SourceMessageId).Bytes()
	return nil
}

// handleMessagesFromClient processes a block of messages from a particular
// client. It saves them to the database, associates them with the contact
// identified by contactTime, and processes them.
func (c commsContext) handleMessagesFromClient(ctx context.Context, info *comms.ClientInfo, contactID db.ContactID, received *fspb.ContactData, validationInfo *fspb.ValidationInfo) error {
	msgs := make([]*fspb.Message, 0, len(received.Messages))
	for _, m := range received.Messages {
		err := c.validateMessageFromClient(info.ID, m, validationInfo)
		if err != nil {
			log.Errorf("Dropping invalid message from [%v]: %v", info.ID, err)
			continue
		}
		msgs = append(msgs, m)
	}
	if len(msgs) == 0 {
		return nil
	}

	sort.Slice(msgs, func(a, b int) bool {
		return bytes.Compare(msgs[a].MessageId, msgs[b].MessageId) == -1
	})

	for {
		if len(msgs) <= processingChunkSize {
			return c.s.serviceConfig.HandleNewMessages(ctx, msgs, contactID)
		}

		if err := c.s.serviceConfig.HandleNewMessages(ctx, msgs[:processingChunkSize], contactID); err != nil {
			return err
		}
		msgs = msgs[processingChunkSize:]
	}
}

// ReadFile returns the data and modification time of file. Caller is
// responsible for closing data.
//
// Calls to data are permitted to fail if ctx is canceled or expired.
func (c commsContext) ReadFile(ctx context.Context, service, name string) (data db.ReadSeekerCloser, modtime time.Time, err error) {
	return c.s.dataStore.ReadFile(ctx, service, name)
}

// IsNotFound returns whether an error returned by ReadFile indicates that the
// file was not found.
func (c commsContext) IsNotFound(err error) bool {
	return c.s.dataStore.IsNotFound(err)
}

// StatsCollector returns the stats.Collector used by the Fleetspeak
// system. Access is provided to allow collection of stats relating to the
// client communication.
func (c commsContext) StatsCollector() stats.Collector {
	return c.s.statsCollector
}

func (c commsContext) Authorizer() authorizer.Authorizer {
	return c.s.authorizer
}
<|endoftext|>


<commit_before>/*
Copyright 2017 WALLIX

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 console

import (
	"github.com/fatih/color"
	"github.com/wallix/awless/cloud"
)

var DefaultsColumnDefinitions = map[string][]ColumnDefinition{
	//EC2
	cloud.Instance: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "AvailabilityZone", Friendly: "Zone"},
		StringColumnDefinition{Prop: "Name"},
		ColoredValueColumnDefinition{
			StringColumnDefinition: StringColumnDefinition{Prop: "State"},
			ColoredValues:          map[string]color.Attribute{"running": color.FgGreen, "stopped": color.FgRed},
		},
		StringColumnDefinition{Prop: "Type"},
		StringColumnDefinition{Prop: "KeyName", Friendly: "Access Key"},
		StringColumnDefinition{Prop: "PublicIp", Friendly: "Public IP"},
		StringColumnDefinition{Prop: "PrivateIp", Friendly: "Private IP"},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "LaunchTime", Friendly: "Up Since"}},
	},
	cloud.Vpc: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name"},
		ColoredValueColumnDefinition{
			StringColumnDefinition: StringColumnDefinition{Prop: "IsDefault", Friendly: "Default"},
			ColoredValues:          map[string]color.Attribute{"true": color.FgGreen},
		},
		StringColumnDefinition{Prop: "State"},
		StringColumnDefinition{Prop: "CidrBlock"},
	},
	cloud.Subnet: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name"},
		StringColumnDefinition{Prop: "CidrBlock"},
		StringColumnDefinition{Prop: "AvailabilityZone", Friendly: "Zone"},
		ColoredValueColumnDefinition{
			StringColumnDefinition: StringColumnDefinition{Prop: "DefaultForAz", Friendly: "Default"},
			ColoredValues:          map[string]color.Attribute{"true": color.FgGreen},
		},
		StringColumnDefinition{Prop: "VpcId"},
		ColoredValueColumnDefinition{
			StringColumnDefinition: StringColumnDefinition{Prop: "MapPublicIpOnLaunch", Friendly: "Public"},
			ColoredValues:          map[string]color.Attribute{"true": color.FgYellow}},
		ColoredValueColumnDefinition{
			StringColumnDefinition: StringColumnDefinition{Prop: "State"},
			ColoredValues:          map[string]color.Attribute{"available": color.FgGreen}},
	},
	cloud.SecurityGroup: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "VpcId"},
		FirewallRulesColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "InboundRules", Friendly: "Inbound"}},
		FirewallRulesColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "OutboundRules", Friendly: "Outbound"}},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		StringColumnDefinition{Prop: "Description", DisableTruncate: true},
	},
	cloud.InternetGateway: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		StringColumnDefinition{Prop: "Vpcs", DisableTruncate: true},
	},
	cloud.RouteTable: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		StringColumnDefinition{Prop: "VpcId"},
		StringColumnDefinition{Prop: "Main"},
		RoutesColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "Routes"}},
	},
	cloud.Keypair: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "KeyFingerprint", DisableTruncate: true},
	},
	cloud.Volume: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		StringColumnDefinition{Prop: "VolumeType", Friendly: "Type"},
		StringColumnDefinition{Prop: "State"},
		StringColumnDefinition{Prop: "Size", Friendly: "Size (Gb)"},
		StringColumnDefinition{Prop: "Encrypted"},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateTime", Friendly: "Created"}},
		StringColumnDefinition{Prop: "AvailabilityZone", Friendly: "Zone"},
	},
	cloud.AvailabilityZone: {
		StringColumnDefinition{Prop: "Name"},
		StringColumnDefinition{Prop: "State"},
		StringColumnDefinition{Prop: "Region"},
		StringColumnDefinition{Prop: "Messages"},
	},
	// Loadbalancer
	cloud.LoadBalancer: {
		StringColumnDefinition{Prop: "Name"},
		StringColumnDefinition{Prop: "VpcId"},
		StringColumnDefinition{Prop: "State"},
		StringColumnDefinition{Prop: "DNSName"},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateTime", Friendly: "Created"}},
		StringColumnDefinition{Prop: "Scheme"},
	},
	cloud.TargetGroup: {
		StringColumnDefinition{Prop: "Name"},
		StringColumnDefinition{Prop: "VpcId"},
		StringColumnDefinition{Prop: "Matcher"},
		StringColumnDefinition{Prop: "Port"},
		StringColumnDefinition{Prop: "Protocol"},
		StringColumnDefinition{Prop: "HealthCheckIntervalSeconds", Friendly: "HCInterval"},
		StringColumnDefinition{Prop: "HealthCheckPath", Friendly: "HCPath"},
		StringColumnDefinition{Prop: "HealthCheckPort", Friendly: "HCPort"},
		StringColumnDefinition{Prop: "HealthCheckProtocol", Friendly: "HCProtocol"},
	},
	cloud.Listener: {
		StringColumnDefinition{Prop: "Id", DisableTruncate: true},
		StringColumnDefinition{Prop: "Actions"},
		StringColumnDefinition{Prop: "LoadBalancer"},
		StringColumnDefinition{Prop: "Port"},
		StringColumnDefinition{Prop: "Protocol"},
		StringColumnDefinition{Prop: "SslPolicy"},
	},
	//IAM
	cloud.User: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "PasswordLastUsedDate", Friendly: "PasswordLastUsed"}},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateDate"}},
	},
	cloud.Role: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateDate"}},
	},
	cloud.Policy: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateDate"}},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "UpdateDate"}},
	},
	cloud.Group: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateDate"}},
	},
	// S3
	cloud.Bucket: {
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		GrantsColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "Grants"}},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateDate"}},
	},
	cloud.Object: {
		StringColumnDefinition{Prop: "Key", TruncateRight: true},
		StringColumnDefinition{Prop: "BucketName", Friendly: "Bucket"},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "ModifiedDate"}},
		StringColumnDefinition{Prop: "OwnerId", TruncateRight: true},
		StringColumnDefinition{Prop: "Size"},
		StringColumnDefinition{Prop: "Class"},
	},
	//Notification
	cloud.Subscription: {
		StringColumnDefinition{Prop: "SubscriptionArn"},
		StringColumnDefinition{Prop: "TopicArn"},
		StringColumnDefinition{Prop: "Endpoint", DisableTruncate: true},
		StringColumnDefinition{Prop: "Protocol"},
		StringColumnDefinition{Prop: "Owner"},
	},
	cloud.Topic: {
		StringColumnDefinition{Prop: "TopicArn", DisableTruncate: true},
	},
	//Queue
	cloud.Queue: {
		StringColumnDefinition{Prop: "Id", Friendly: "URL", DisableTruncate: true},
		StringColumnDefinition{Prop: "ApproximateNumberOfMessages", Friendly: "~NbMsg"},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreatedTimestamp", Friendly: "Created"}},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "LastModifiedTimestamp", Friendly: "LastModif"}},
		StringColumnDefinition{Prop: "DelaySeconds", Friendly: "Delay(s)"},
	},
	// DNS
	cloud.Zone: {
		StringColumnDefinition{Prop: "Id", DisableTruncate: true},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		StringColumnDefinition{Prop: "Comment"},
		StringColumnDefinition{Prop: "IsPrivateZone"},
		StringColumnDefinition{Prop: "ResourceRecordSetCount"},
		StringColumnDefinition{Prop: "CallerReference", DisableTruncate: true},
	},
	cloud.Record: {
		StringColumnDefinition{Prop: "Id", Friendly: "AwlessId", DisableTruncate: true},
		StringColumnDefinition{Prop: "Type"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		SliceColumnDefinition{StringColumnDefinition{Prop: "Records"}},
		StringColumnDefinition{Prop: "TTL"},
	},
}
<commit_msg>Display: moving instance ACCESS KEY to get term space<commit_after>/*
Copyright 2017 WALLIX

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 console

import (
	"github.com/fatih/color"
	"github.com/wallix/awless/cloud"
)

var DefaultsColumnDefinitions = map[string][]ColumnDefinition{
	//EC2
	cloud.Instance: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "AvailabilityZone", Friendly: "Zone"},
		StringColumnDefinition{Prop: "Name"},
		ColoredValueColumnDefinition{
			StringColumnDefinition: StringColumnDefinition{Prop: "State"},
			ColoredValues:          map[string]color.Attribute{"running": color.FgGreen, "stopped": color.FgRed},
		},
		StringColumnDefinition{Prop: "Type"},
		StringColumnDefinition{Prop: "PublicIp", Friendly: "Public IP"},
		StringColumnDefinition{Prop: "PrivateIp", Friendly: "Private IP"},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "LaunchTime", Friendly: "Up Since"}},
		StringColumnDefinition{Prop: "KeyName", Friendly: "Access Key"},
	},
	cloud.Vpc: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name"},
		ColoredValueColumnDefinition{
			StringColumnDefinition: StringColumnDefinition{Prop: "IsDefault", Friendly: "Default"},
			ColoredValues:          map[string]color.Attribute{"true": color.FgGreen},
		},
		StringColumnDefinition{Prop: "State"},
		StringColumnDefinition{Prop: "CidrBlock"},
	},
	cloud.Subnet: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name"},
		StringColumnDefinition{Prop: "CidrBlock"},
		StringColumnDefinition{Prop: "AvailabilityZone", Friendly: "Zone"},
		ColoredValueColumnDefinition{
			StringColumnDefinition: StringColumnDefinition{Prop: "DefaultForAz", Friendly: "Default"},
			ColoredValues:          map[string]color.Attribute{"true": color.FgGreen},
		},
		StringColumnDefinition{Prop: "VpcId"},
		ColoredValueColumnDefinition{
			StringColumnDefinition: StringColumnDefinition{Prop: "MapPublicIpOnLaunch", Friendly: "Public"},
			ColoredValues:          map[string]color.Attribute{"true": color.FgYellow}},
		ColoredValueColumnDefinition{
			StringColumnDefinition: StringColumnDefinition{Prop: "State"},
			ColoredValues:          map[string]color.Attribute{"available": color.FgGreen}},
	},
	cloud.SecurityGroup: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "VpcId"},
		FirewallRulesColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "InboundRules", Friendly: "Inbound"}},
		FirewallRulesColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "OutboundRules", Friendly: "Outbound"}},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		StringColumnDefinition{Prop: "Description", DisableTruncate: true},
	},
	cloud.InternetGateway: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		StringColumnDefinition{Prop: "Vpcs", DisableTruncate: true},
	},
	cloud.RouteTable: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		StringColumnDefinition{Prop: "VpcId"},
		StringColumnDefinition{Prop: "Main"},
		RoutesColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "Routes"}},
	},
	cloud.Keypair: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "KeyFingerprint", DisableTruncate: true},
	},
	cloud.Volume: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		StringColumnDefinition{Prop: "VolumeType", Friendly: "Type"},
		StringColumnDefinition{Prop: "State"},
		StringColumnDefinition{Prop: "Size", Friendly: "Size (Gb)"},
		StringColumnDefinition{Prop: "Encrypted"},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateTime", Friendly: "Created"}},
		StringColumnDefinition{Prop: "AvailabilityZone", Friendly: "Zone"},
	},
	cloud.AvailabilityZone: {
		StringColumnDefinition{Prop: "Name"},
		StringColumnDefinition{Prop: "State"},
		StringColumnDefinition{Prop: "Region"},
		StringColumnDefinition{Prop: "Messages"},
	},
	// Loadbalancer
	cloud.LoadBalancer: {
		StringColumnDefinition{Prop: "Name"},
		StringColumnDefinition{Prop: "VpcId"},
		StringColumnDefinition{Prop: "State"},
		StringColumnDefinition{Prop: "DNSName"},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateTime", Friendly: "Created"}},
		StringColumnDefinition{Prop: "Scheme"},
	},
	cloud.TargetGroup: {
		StringColumnDefinition{Prop: "Name"},
		StringColumnDefinition{Prop: "VpcId"},
		StringColumnDefinition{Prop: "Matcher"},
		StringColumnDefinition{Prop: "Port"},
		StringColumnDefinition{Prop: "Protocol"},
		StringColumnDefinition{Prop: "HealthCheckIntervalSeconds", Friendly: "HCInterval"},
		StringColumnDefinition{Prop: "HealthCheckPath", Friendly: "HCPath"},
		StringColumnDefinition{Prop: "HealthCheckPort", Friendly: "HCPort"},
		StringColumnDefinition{Prop: "HealthCheckProtocol", Friendly: "HCProtocol"},
	},
	cloud.Listener: {
		StringColumnDefinition{Prop: "Id", DisableTruncate: true},
		StringColumnDefinition{Prop: "Actions"},
		StringColumnDefinition{Prop: "LoadBalancer"},
		StringColumnDefinition{Prop: "Port"},
		StringColumnDefinition{Prop: "Protocol"},
		StringColumnDefinition{Prop: "SslPolicy"},
	},
	//IAM
	cloud.User: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "PasswordLastUsedDate", Friendly: "PasswordLastUsed"}},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateDate"}},
	},
	cloud.Role: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateDate"}},
	},
	cloud.Policy: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateDate"}},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "UpdateDate"}},
	},
	cloud.Group: {
		StringColumnDefinition{Prop: "Id"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateDate"}},
	},
	// S3
	cloud.Bucket: {
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		GrantsColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "Grants"}},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreateDate"}},
	},
	cloud.Object: {
		StringColumnDefinition{Prop: "Key", TruncateRight: true},
		StringColumnDefinition{Prop: "BucketName", Friendly: "Bucket"},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "ModifiedDate"}},
		StringColumnDefinition{Prop: "OwnerId", TruncateRight: true},
		StringColumnDefinition{Prop: "Size"},
		StringColumnDefinition{Prop: "Class"},
	},
	//Notification
	cloud.Subscription: {
		StringColumnDefinition{Prop: "SubscriptionArn"},
		StringColumnDefinition{Prop: "TopicArn"},
		StringColumnDefinition{Prop: "Endpoint", DisableTruncate: true},
		StringColumnDefinition{Prop: "Protocol"},
		StringColumnDefinition{Prop: "Owner"},
	},
	cloud.Topic: {
		StringColumnDefinition{Prop: "TopicArn", DisableTruncate: true},
	},
	//Queue
	cloud.Queue: {
		StringColumnDefinition{Prop: "Id", Friendly: "URL", DisableTruncate: true},
		StringColumnDefinition{Prop: "ApproximateNumberOfMessages", Friendly: "~NbMsg"},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "CreatedTimestamp", Friendly: "Created"}},
		TimeColumnDefinition{StringColumnDefinition: StringColumnDefinition{Prop: "LastModifiedTimestamp", Friendly: "LastModif"}},
		StringColumnDefinition{Prop: "DelaySeconds", Friendly: "Delay(s)"},
	},
	// DNS
	cloud.Zone: {
		StringColumnDefinition{Prop: "Id", DisableTruncate: true},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		StringColumnDefinition{Prop: "Comment"},
		StringColumnDefinition{Prop: "IsPrivateZone"},
		StringColumnDefinition{Prop: "ResourceRecordSetCount"},
		StringColumnDefinition{Prop: "CallerReference", DisableTruncate: true},
	},
	cloud.Record: {
		StringColumnDefinition{Prop: "Id", Friendly: "AwlessId", DisableTruncate: true},
		StringColumnDefinition{Prop: "Type"},
		StringColumnDefinition{Prop: "Name", DisableTruncate: true},
		SliceColumnDefinition{StringColumnDefinition{Prop: "Records"}},
		StringColumnDefinition{Prop: "TTL"},
	},
}
<|endoftext|>


<commit_before>package container

import (
	"crypto/tls"
	"fmt"
	"time"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	defaultStopSignal = "SIGTERM"
	dryRunPrefix      = "DRY: "
)

// A Filter is a prototype for a function that can be used to filter the
// results from a call to the ListContainers() method on the Client.
type Filter func(Container) bool

// A Client is the interface through which Pumba interacts with the Docker API.
type Client interface {
	ListContainers(Filter) ([]Container, error)
	StopContainer(Container, time.Duration, bool) error
	KillContainer(Container, string, bool) error
	StartContainer(Container) error
	RenameContainer(Container, string) error
	RemoveImage(Container, bool, bool) error
	RemoveContainer(Container, bool, bool) error
	DisruptContainer(Container, string, bool) error
}

// NewClient returns a new Client instance which can be used to interact with
// the Docker API.
func NewClient(dockerHost string, tlsConfig *tls.Config, pullImages bool) Client {
	docker, err := dockerclient.NewDockerClient(dockerHost, tlsConfig)

	if err != nil {
		log.Fatalf("Error instantiating Docker client: %s", err)
	}

	return dockerClient{api: docker, pullImages: pullImages}
}

type dockerClient struct {
	api        dockerclient.Client
	pullImages bool
}

func (client dockerClient) ListContainers(fn Filter) ([]Container, error) {
	cs := []Container{}

	log.Debug("Retrieving running containers")

	runningContainers, err := client.api.ListContainers(false, false, "")
	if err != nil {
		return nil, err
	}
	for _, runningContainer := range runningContainers {
		containerInfo, err := client.api.InspectContainer(runningContainer.Id)
		if err != nil {
			return nil, err
		}
		log.Debugf("Running container: %s - (%s)", containerInfo.Name, containerInfo.Id)

		imageInfo, err := client.api.InspectImage(containerInfo.Image)
		if err != nil {
			return nil, err
		}

		c := Container{containerInfo: containerInfo, imageInfo: imageInfo}
		if fn(c) {
			cs = append(cs, c)
		}
	}

	return cs, nil
}

func (client dockerClient) KillContainer(c Container, signal string, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sKilling %s (%s) with signal %s", prefix, c.Name(), c.ID(), signal)
	if !dryrun {
		if err := client.api.KillContainer(c.ID(), signal); err != nil {
			return err
		}
	}
	return nil
}

func (client dockerClient) StopContainer(c Container, timeout time.Duration, dryrun bool) error {
	signal := c.StopSignal()
	if signal == "" {
		signal = defaultStopSignal
	}
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sStopping %s (%s) with %s", prefix, c.Name(), c.ID(), signal)

	if !dryrun {
		if err := client.api.KillContainer(c.ID(), signal); err != nil {
			return err
		}

		// Wait for container to exit, but proceed anyway after the timeout elapses
		if err := client.waitForStop(c, timeout); err != nil {
			log.Debugf("Error waiting for container %s (%s) to stop: ''%s'", c.Name(), c.ID(), err.Error())
		}

		log.Debugf("Removing container %s", c.ID())

		if err := client.api.RemoveContainer(c.ID(), true, false); err != nil {
			return err
		}

		// Wait for container to be removed. In this case an error is a good thing
		if err := client.waitForStop(c, timeout); err == nil {
			return fmt.Errorf("Container %s (%s) could not be removed", c.Name(), c.ID())
		}
	}

	return nil
}

func (client dockerClient) StartContainer(c Container) error {
	config := c.runtimeConfig()
	hostConfig := c.hostConfig()
	name := c.Name()

	log.Infof("Starting %s", name)

	newContainerID, err := client.api.CreateContainer(config, name, nil)
	if err != nil {
		return err
	}

	log.Debugf("Starting container %s (%s)", name, newContainerID)

	return client.api.StartContainer(newContainerID, hostConfig)
}

func (client dockerClient) RenameContainer(c Container, newName string) error {
	log.Debugf("Renaming container %s (%s) to %s", c.Name(), c.ID(), newName)
	return client.api.RenameContainer(c.ID(), newName)
}

func (client dockerClient) RemoveImage(c Container, force bool, dryrun bool) error {
	imageID := c.ImageID()
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sRemoving image %s", prefix, imageID)
	if !dryrun {
		_, err := client.api.RemoveImage(imageID, force)
		return err
	}
	return nil
}

func (client dockerClient) RemoveContainer(c Container, force bool, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sRemoving container %s", prefix, c.ID())
	if !dryrun {
		return client.api.RemoveContainer(c.ID(), force, true)
	}
	return nil
}

func (client dockerClient) DisruptContainer(c Container, netemCmd string, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sDisrupting container %s with netem cmd %s", prefix, c.ID(), netemCmd)
	if !dryrun {
		// use dockerclient ExecStart to run Traffic Control:
		// 'tc qdisc add dev eth0 root netem delay 100ms'
		// http://www.linuxfoundation.org/collaborate/workgroups/networking/netem
		netemBase := strings.Split("tc qdisc add dev eth0 root netem", " ")
		netemCommand := strings.Split(strings.ToLower(netemCmd), " ")
		netemMerge := append(netemBase, netemSplit)
		execConfig := &dockerclient.ExecConfig{
			Cmd: netemMerge,
			Container: c.ID(),
		}
		_id, err := client.api.ExecCreate(execConfig)
		if err != nil {
				return err
			}

		log.Debugf("Starting Exec %s (%s)", netemMerge, _id)
		return client.api.ExecStart(_id, execConfig)
	}
	return nil
}

func (client dockerClient) waitForStop(c Container, waitTime time.Duration) error {
	timeout := time.After(waitTime)

	for {
		select {
		case <-timeout:
			return nil
		default:
			if ci, err := client.api.InspectContainer(c.ID()); err != nil {
				return err
			} else if !ci.State.Running {
				return nil
			}
		}

		time.Sleep(1 * time.Second)
	}
}
<commit_msg>remove append syntax?<commit_after>package container

import (
	"crypto/tls"
	"fmt"
	"time"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	defaultStopSignal = "SIGTERM"
	dryRunPrefix      = "DRY: "
)

// A Filter is a prototype for a function that can be used to filter the
// results from a call to the ListContainers() method on the Client.
type Filter func(Container) bool

// A Client is the interface through which Pumba interacts with the Docker API.
type Client interface {
	ListContainers(Filter) ([]Container, error)
	StopContainer(Container, time.Duration, bool) error
	KillContainer(Container, string, bool) error
	StartContainer(Container) error
	RenameContainer(Container, string) error
	RemoveImage(Container, bool, bool) error
	RemoveContainer(Container, bool, bool) error
	DisruptContainer(Container, string, bool) error
}

// NewClient returns a new Client instance which can be used to interact with
// the Docker API.
func NewClient(dockerHost string, tlsConfig *tls.Config, pullImages bool) Client {
	docker, err := dockerclient.NewDockerClient(dockerHost, tlsConfig)

	if err != nil {
		log.Fatalf("Error instantiating Docker client: %s", err)
	}

	return dockerClient{api: docker, pullImages: pullImages}
}

type dockerClient struct {
	api        dockerclient.Client
	pullImages bool
}

func (client dockerClient) ListContainers(fn Filter) ([]Container, error) {
	cs := []Container{}

	log.Debug("Retrieving running containers")

	runningContainers, err := client.api.ListContainers(false, false, "")
	if err != nil {
		return nil, err
	}
	for _, runningContainer := range runningContainers {
		containerInfo, err := client.api.InspectContainer(runningContainer.Id)
		if err != nil {
			return nil, err
		}
		log.Debugf("Running container: %s - (%s)", containerInfo.Name, containerInfo.Id)

		imageInfo, err := client.api.InspectImage(containerInfo.Image)
		if err != nil {
			return nil, err
		}

		c := Container{containerInfo: containerInfo, imageInfo: imageInfo}
		if fn(c) {
			cs = append(cs, c)
		}
	}

	return cs, nil
}

func (client dockerClient) KillContainer(c Container, signal string, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sKilling %s (%s) with signal %s", prefix, c.Name(), c.ID(), signal)
	if !dryrun {
		if err := client.api.KillContainer(c.ID(), signal); err != nil {
			return err
		}
	}
	return nil
}

func (client dockerClient) StopContainer(c Container, timeout time.Duration, dryrun bool) error {
	signal := c.StopSignal()
	if signal == "" {
		signal = defaultStopSignal
	}
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sStopping %s (%s) with %s", prefix, c.Name(), c.ID(), signal)

	if !dryrun {
		if err := client.api.KillContainer(c.ID(), signal); err != nil {
			return err
		}

		// Wait for container to exit, but proceed anyway after the timeout elapses
		if err := client.waitForStop(c, timeout); err != nil {
			log.Debugf("Error waiting for container %s (%s) to stop: ''%s'", c.Name(), c.ID(), err.Error())
		}

		log.Debugf("Removing container %s", c.ID())

		if err := client.api.RemoveContainer(c.ID(), true, false); err != nil {
			return err
		}

		// Wait for container to be removed. In this case an error is a good thing
		if err := client.waitForStop(c, timeout); err == nil {
			return fmt.Errorf("Container %s (%s) could not be removed", c.Name(), c.ID())
		}
	}

	return nil
}

func (client dockerClient) StartContainer(c Container) error {
	config := c.runtimeConfig()
	hostConfig := c.hostConfig()
	name := c.Name()

	log.Infof("Starting %s", name)

	newContainerID, err := client.api.CreateContainer(config, name, nil)
	if err != nil {
		return err
	}

	log.Debugf("Starting container %s (%s)", name, newContainerID)

	return client.api.StartContainer(newContainerID, hostConfig)
}

func (client dockerClient) RenameContainer(c Container, newName string) error {
	log.Debugf("Renaming container %s (%s) to %s", c.Name(), c.ID(), newName)
	return client.api.RenameContainer(c.ID(), newName)
}

func (client dockerClient) RemoveImage(c Container, force bool, dryrun bool) error {
	imageID := c.ImageID()
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sRemoving image %s", prefix, imageID)
	if !dryrun {
		_, err := client.api.RemoveImage(imageID, force)
		return err
	}
	return nil
}

func (client dockerClient) RemoveContainer(c Container, force bool, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sRemoving container %s", prefix, c.ID())
	if !dryrun {
		return client.api.RemoveContainer(c.ID(), force, true)
	}
	return nil
}

func (client dockerClient) DisruptContainer(c Container, netemCmd string, dryrun bool) error {
	prefix := ""
	if dryrun {
		prefix = dryRunPrefix
	}
	log.Infof("%sDisrupting container %s with netem cmd %s", prefix, c.ID(), netemCmd)
	if !dryrun {
		// use dockerclient ExecStart to run Traffic Control:
		// 'tc qdisc add dev eth0 root netem delay 100ms'
		// http://www.linuxfoundation.org/collaborate/workgroups/networking/netem
		netemBase := strings.Split("tc qdisc add dev eth0 root netem", " ")
		netemCommand := strings.Split(strings.ToLower(netemCmd), " ")
		netemMerge := append(netemBase, netemCommand)
		execConfig := &dockerclient.ExecConfig{
			Cmd: netemMerge,
			Container: c.ID(),
		}
		_id, err := client.api.ExecCreate(execConfig)
		if err != nil {
				return err
			}

		log.Debugf("Starting Exec %s (%s)", netemMerge, _id)
		return client.api.ExecStart(_id, execConfig)
	}
	return nil
}

func (client dockerClient) waitForStop(c Container, waitTime time.Duration) error {
	timeout := time.After(waitTime)

	for {
		select {
		case <-timeout:
			return nil
		default:
			if ci, err := client.api.InspectContainer(c.ID()); err != nil {
				return err
			} else if !ci.State.Running {
				return nil
			}
		}

		time.Sleep(1 * time.Second)
	}
}
<|endoftext|>


<commit_before>package hue

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os/user"
	"time"

	"context"

	log "github.com/sirupsen/logrus"

	"github.com/heatxsink/go-hue/groups"
	"github.com/heatxsink/go-hue/lights"
	"github.com/heatxsink/go-hue/sensors"
)

// Server : The implementation of the Philips Hue gRPC service
type Server struct{}

// Bridge : Connection and authentication information related to the Philips Hue Bridge
type Bridge struct {
	Bridge     string `json:"bridge"`
	Username   string `json:"username"`
	Devicetype string `json:"devicetype"`
}

var hueBridge Bridge

// LoadHueBridgeConfig : Loads the Philips Hue configuration file from file.
func LoadHueBridgeConfig() Bridge {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	content, err := ioutil.ReadFile(usr.HomeDir + "/.philips-hue.json")
	if err != nil {
		log.Fatal(err)
	}
	err = json.Unmarshal(content, &hueBridge)
	if err != nil {
		fmt.Print("Error:", err)
	}
	return hueBridge
}

/************************************************************************************
 * Start gRPC Service implementation
 */

// GetSensors : Returns all known sensors
func (s Server) GetSensors(ctx context.Context, in *SensorRequest) (*Sensors, error) {
	start := time.Now()
	ss := sensors.New(hueBridge.Bridge, hueBridge.Username)
	allSensors, err := ss.GetAllSensors()
	if err != nil {
		return nil, err
	}
	log.WithField("bridge_exec_time", time.Since(start)).
		WithField("call", "GetSensors").Print("Incoming gRPC call")

	sensorsResp := &Sensors{}
	sensorsResp.Sensors = make([]*Sensor, len(allSensors))
	for i, s := range allSensors {
		sensorsResp.Sensors[i] = &Sensor{
			ID:               int32(s.ID),
			UniqueID:         s.UniqueID,
			Name:             s.Name,
			Type:             s.Type,
			ModelID:          s.ModelID,
			SWVersion:        s.SWVersion,
			ManufacturerName: s.ManufacturerName,
			State: &State{
				ButtonEvent: int32(s.State.ButtonEvent),
				Dark:        s.State.Dark,
				Daylight:    s.State.Daylight,
				LastUpdated: s.State.LastUpdated,
				LightLevel:  int32(s.State.LightLevel),
				Presence:    s.State.Presence,
				Status:      int32(s.State.Status),
				Temperature: int32(s.State.Temperature),
			},
		}
	}
	return sensorsResp, nil
}

// GetGroups implements hue.Lights : Returns group of lights.
func (s Server) GetGroups(ctx context.Context, in *LightsRequest) (*Groups, error) {
	start := time.Now()
	gg := groups.New(hueBridge.Bridge, hueBridge.Username)
	allGroups, err := gg.GetAllGroups()
	if err != nil {
		return nil, err
	}
	log.WithField("bridge_exec_time", time.Since(start)).
		WithField("call", "GetGroups").Print("Incoming gRPC call")

	groups := &Groups{}
	groups.Groups = make([]*Group, len(allGroups))
	for i, g := range allGroups {
		groups.Groups[i] = &Group{ID: int32(g.ID), Name: g.Name, On: g.Action.On, Brightness: int32(g.Action.Bri)}
	}
	return groups, nil
}

// SwitchOn implements hue.Lights
func (s Server) SwitchOn(ctx context.Context, in *LightsRequest) (*LightsResponse, error) {
	log.WithField("message", in).
		WithField("call", "SwitchOn").Print("Incoming gRPC call")
	gg := groups.New(hueBridge.Bridge, hueBridge.Username)
	g, err := gg.GetGroup(int(in.GetGroup()))
	if err != nil {
		return nil, err
	}
	brightness := uint8(255 * in.GetBrightnessPercent())
	gg.SetGroupState(g.ID, lights.State{
		On:  true,
		Bri: brightness,
		Sat: 254,
		Hue: 10000,
	})
	return &LightsResponse{Success: true}, nil
}

// SwitchOff implements hue.Lights
func (s Server) SwitchOff(ctx context.Context, in *LightsRequest) (*LightsResponse, error) {
	log.WithField("message", in).
		WithField("call", "SwitchOff").Print("Incoming gRPC call")
	gg := groups.New(hueBridge.Bridge, hueBridge.Username)
	g, err := gg.GetGroup(int(in.GetGroup()))
	if err != nil {
		return nil, err
	}
	gg.SetGroupState(g.ID, lights.State{On: false})
	return &LightsResponse{Success: true}, nil
}

// Blink implement hue.Lights : It switches on and off the lights in a short time period.
// After done it resets to the previous state.
func (s Server) Blink(ctx context.Context, in *LightsRequest) (*LightsResponse, error) {
	log.WithField("message", in).
		WithField("call", "Blink").Print("Incoming gRPC call")
	gg := groups.New(hueBridge.Bridge, hueBridge.Username)
	g, err := gg.GetGroup(int(in.GetGroup()))
	if err != nil {
		return nil, err
	}

	oldState := g.Action

	brightness := uint8(255 * in.GetBrightnessPercent())
	gg.SetGroupState(g.ID, lights.State{On: true, Bri: brightness, Alert: "lselect"})
	gg.SetGroupState(g.ID, oldState)
	return &LightsResponse{Success: true}, nil
}
<commit_msg>Removing the hard coded color values on lights switch on logic<commit_after>package hue

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os/user"
	"time"

	"context"

	log "github.com/sirupsen/logrus"

	"github.com/heatxsink/go-hue/groups"
	"github.com/heatxsink/go-hue/lights"
	"github.com/heatxsink/go-hue/sensors"
)

// Server : The implementation of the Philips Hue gRPC service
type Server struct{}

// Bridge : Connection and authentication information related to the Philips Hue Bridge
type Bridge struct {
	Bridge     string `json:"bridge"`
	Username   string `json:"username"`
	Devicetype string `json:"devicetype"`
}

var hueBridge Bridge

// LoadHueBridgeConfig : Loads the Philips Hue configuration file from file.
func LoadHueBridgeConfig() Bridge {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	content, err := ioutil.ReadFile(usr.HomeDir + "/.philips-hue.json")
	if err != nil {
		log.Fatal(err)
	}
	err = json.Unmarshal(content, &hueBridge)
	if err != nil {
		fmt.Print("Error:", err)
	}
	return hueBridge
}

/************************************************************************************
 * Start gRPC Service implementation
 */

// GetSensors : Returns all known sensors
func (s Server) GetSensors(ctx context.Context, in *SensorRequest) (*Sensors, error) {
	start := time.Now()
	ss := sensors.New(hueBridge.Bridge, hueBridge.Username)
	allSensors, err := ss.GetAllSensors()
	if err != nil {
		return nil, err
	}
	log.WithField("bridge_exec_time", time.Since(start)).
		WithField("call", "GetSensors").Print("Incoming gRPC call")

	sensorsResp := &Sensors{}
	sensorsResp.Sensors = make([]*Sensor, len(allSensors))
	for i, s := range allSensors {
		sensorsResp.Sensors[i] = &Sensor{
			ID:               int32(s.ID),
			UniqueID:         s.UniqueID,
			Name:             s.Name,
			Type:             s.Type,
			ModelID:          s.ModelID,
			SWVersion:        s.SWVersion,
			ManufacturerName: s.ManufacturerName,
			State: &State{
				ButtonEvent: int32(s.State.ButtonEvent),
				Dark:        s.State.Dark,
				Daylight:    s.State.Daylight,
				LastUpdated: s.State.LastUpdated,
				LightLevel:  int32(s.State.LightLevel),
				Presence:    s.State.Presence,
				Status:      int32(s.State.Status),
				Temperature: int32(s.State.Temperature),
			},
		}
	}
	return sensorsResp, nil
}

// GetGroups implements hue.Lights : Returns group of lights.
func (s Server) GetGroups(ctx context.Context, in *LightsRequest) (*Groups, error) {
	start := time.Now()
	gg := groups.New(hueBridge.Bridge, hueBridge.Username)
	allGroups, err := gg.GetAllGroups()
	if err != nil {
		return nil, err
	}
	log.WithField("bridge_exec_time", time.Since(start)).
		WithField("call", "GetGroups").Print("Incoming gRPC call")

	groups := &Groups{}
	groups.Groups = make([]*Group, len(allGroups))
	for i, g := range allGroups {
		groups.Groups[i] = &Group{ID: int32(g.ID), Name: g.Name, On: g.Action.On, Brightness: int32(g.Action.Bri)}
	}
	return groups, nil
}

// SwitchOn implements hue.Lights
func (s Server) SwitchOn(ctx context.Context, in *LightsRequest) (*LightsResponse, error) {
	log.WithField("message", in).
		WithField("call", "SwitchOn").Print("Incoming gRPC call")
	gg := groups.New(hueBridge.Bridge, hueBridge.Username)
	g, err := gg.GetGroup(int(in.GetGroup()))
	if err != nil {
		return nil, err
	}
	brightness := uint8(255 * in.GetBrightnessPercent())
	gg.SetGroupState(g.ID, lights.State{
		On:  true,
		Bri: brightness,
	})
	return &LightsResponse{Success: true}, nil
}

// SwitchOff implements hue.Lights
func (s Server) SwitchOff(ctx context.Context, in *LightsRequest) (*LightsResponse, error) {
	log.WithField("message", in).
		WithField("call", "SwitchOff").Print("Incoming gRPC call")
	gg := groups.New(hueBridge.Bridge, hueBridge.Username)
	g, err := gg.GetGroup(int(in.GetGroup()))
	if err != nil {
		return nil, err
	}
	gg.SetGroupState(g.ID, lights.State{On: false})
	return &LightsResponse{Success: true}, nil
}

// Blink implement hue.Lights : It switches on and off the lights in a short time period.
// After done it resets to the previous state.
func (s Server) Blink(ctx context.Context, in *LightsRequest) (*LightsResponse, error) {
	log.WithField("message", in).
		WithField("call", "Blink").Print("Incoming gRPC call")
	gg := groups.New(hueBridge.Bridge, hueBridge.Username)
	g, err := gg.GetGroup(int(in.GetGroup()))
	if err != nil {
		return nil, err
	}

	oldState := g.Action

	brightness := uint8(255 * in.GetBrightnessPercent())
	gg.SetGroupState(g.ID, lights.State{On: true, Bri: brightness, Alert: "lselect"})
	gg.SetGroupState(g.ID, oldState)
	return &LightsResponse{Success: true}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"sync"

	appplugin "github.com/ava-labs/avalanchego/app/plugin"
	"github.com/ava-labs/avalanchego/config"
	"github.com/ava-labs/avalanchego/node"
	"github.com/ava-labs/avalanchego/utils/constants"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/hashicorp/go-hclog"
	"github.com/hashicorp/go-plugin"
	"github.com/spf13/viper"
)

// nodeProcess wraps a node client
type nodeProcess struct {
	log logging.Logger
	// Every nodeProcess has a unique ID
	processID int
	// [rawClient].Kill() should eventually be called be called
	// on each nodeProcess
	rawClient *plugin.Client
	node      *appplugin.Client
}

// Returns a channel that the node's exit code is sent on when the node is done
// This method does not block.
func (np *nodeProcess) start() chan int {
	exitCodeChan := make(chan int, 1)
	go func() {
		exitCode, err := np.node.Start()
		if err != nil {
			// This error could be from the subprocess shutting down
			np.log.Debug("node returned: %s", err)
		}
		exitCodeChan <- exitCode
	}()
	return exitCodeChan
}

// Stop should be called on each nodeProcess when we are done with it.
// Calls [Kill()] on the underlying client.
func (np *nodeProcess) stop() error {
	err := np.node.Stop()
	np.rawClient.Kill()
	return err
}

type nodeManager struct {
	// Path to the build directory, which should have this structure:
	// build
	// |_avalanchego-latest
	//   |_avalanchego-process (the binary from compiling the app directory)
	//   |_plugins
	//     |_evm
	// |_avalanchego-preupgrade
	//   |_avalanchego-process (the binary from compiling the app directory)
	//   |_plugins
	//     |_evm
	buildDirPath string
	log          logging.Logger
	// nodeProcess ID --> nodeProcess
	nodes         map[int]*nodeProcess
	nextProcessID int
	processLock   sync.Mutex
}

func (nm *nodeManager) latestNodeVersionPath() string {
	return fmt.Sprintf("%s/avalanchego-latest/avalanchego-process", nm.buildDirPath)
}

func (nm *nodeManager) preupgradeNodeVersionPath() string {
	return fmt.Sprintf("%s/avalanchego-preupgrade/avalanchego-process", nm.buildDirPath)
}

func (nm *nodeManager) shutdown() {
	nm.processLock.Lock()
	defer nm.processLock.Unlock()

	for _, node := range nm.nodes {
		nm.log.Info("stopping process %v", node.processID)
		if err := nm.stop(node.processID); err != nil {
			nm.log.Error("error stopping node: %s", err)
		}
		nm.log.Info("done stopping process %v", node.processID)
	}
}

// stop a node. Blocks until the node is done shutting down.
// Assumes [nm.processLock] is not held
func (nm *nodeManager) Stop(processID int) error {
	nm.processLock.Lock()
	defer nm.processLock.Unlock()
	return nm.stop(processID)
}

// stop a node. Blocks until the node is done shutting down.
// Assumes [nm.processLock] is held
func (nm *nodeManager) stop(processID int) error {
	nodeProcess, exists := nm.nodes[processID]
	if !exists {
		return nil
	}
	if err := nodeProcess.stop(); err != nil {
		return err
	}
	delete(nm.nodes, processID)
	return nil
}

func newNodeManager(path string, log logging.Logger) *nodeManager {
	return &nodeManager{
		buildDirPath: path,
		log:          log,
		nodes:        map[int]*nodeProcess{},
		processLock:  sync.Mutex{},
	}
}

// Return a wrapper around a node wunning the binary at [path] with args [args].
// The returned nodeProcess must eventually have [nodeProcess.rawClient.Kill] called on it.
func (nm *nodeManager) newNode(path string, args []string, printToStdOut bool) (*nodeProcess, error) {
	clientConfig := &plugin.ClientConfig{
		HandshakeConfig: appplugin.Handshake,
		Plugins:         appplugin.PluginMap,
		Cmd:             exec.Command(path, args...),
		AllowedProtocols: []plugin.Protocol{
			plugin.ProtocolNetRPC,
			plugin.ProtocolGRPC,
		},
		Logger: hclog.New(&hclog.LoggerOptions{Level: hclog.Warn}),
	}
	if printToStdOut {
		clientConfig.SyncStdout = os.Stdout
		clientConfig.SyncStderr = os.Stderr
	}
	client := plugin.NewClient(clientConfig)
	rpcClient, err := client.Client()
	if err != nil {
		client.Kill()
		return nil, fmt.Errorf("couldn't get Client: %w", err)
	}
	raw, err := rpcClient.Dispense("nodeProcess")
	if err != nil {
		client.Kill()
		return nil, fmt.Errorf("couldn't dispense plugin '%s': %w", "nodeProcess", err)
	}
	node, ok := raw.(*appplugin.Client)
	if !ok {
		client.Kill()
		return nil, fmt.Errorf("expected *node.NodeClient but got %T", raw)
	}
	nm.nextProcessID++
	np := &nodeProcess{
		log:       nm.log,
		node:      node,
		rawClient: client,
		processID: nm.nextProcessID,
	}
	nm.nodes[np.processID] = np
	return np, nil
}

// Start a node compatible with the previous database version
// Override the staking port, HTTP port and plugin directory of the node.
// Assumes the node binary path is [buildDir]/avalanchego-preupgrade/avalanchego-process
// Assumes the node's plugin path is [buildDir]/avalanchego-preupgrade/plugins
// Assumes the binary can be served as a plugin
func (nm *nodeManager) preDBUpgradeNode(v *viper.Viper) (*nodeProcess, error) {
	argsMap := v.AllSettings()
	delete(argsMap, config.FetchOnlyKey)
	argsMap[config.PluginModeKey] = true
	argsMap[config.PluginDirKey] = fmt.Sprintf("%s/avalanchego-preupgrade/plugins", nm.buildDirPath)
	args := []string{}
	for k, v := range argsMap { // Pass args to subprocess
		args = append(args, formatArgs(k, v))
	}
	binaryPath := nm.preupgradeNodeVersionPath()
	return nm.newNode(binaryPath, args, true)
}

// Run the latest node version
func (nm *nodeManager) latestVersionNodeFetchOnly(v *viper.Viper, nodeConfig node.Config) (*nodeProcess, error) {
	argsMap := v.AllSettings()
	argsMap[config.BootstrapIPsKey] = fmt.Sprintf("127.0.0.1:%d", int(nodeConfig.StakingIP.Port))
	argsMap[config.BootstrapIDsKey] = fmt.Sprintf("%s%s", constants.NodeIDPrefix, nodeConfig.NodeID)
	argsMap[config.FetchOnlyKey] = true
	argsMap[config.StakingPortKey] = 0
	argsMap[config.HTTPPortKey] = 0
	argsMap[config.PluginModeKey] = true
	// replace the last folder named daemon in path with fetch-only
	daemonLogDir := nodeConfig.LoggingConfig.Directory
	i := strings.LastIndex(daemonLogDir, "daemon")
	argsMap[config.LogsDirKey] = daemonLogDir[:i] + strings.Replace(daemonLogDir[i:], "daemon", "fetch-only", 1)

	var args []string
	for k, v := range argsMap {
		args = append(args, formatArgs(k, v))
	}
	binaryPath := nm.latestNodeVersionPath()
	return nm.newNode(binaryPath, args, false)
}

// Run the latest node version with the config given by [v].
// Returns the node's exit code.
func (nm *nodeManager) runNormal(v *viper.Viper) (int, error) {
	nm.log.Info("starting latest node version")
	argsMap := v.AllSettings()
	argsMap[config.FetchOnlyKey] = false // don't run in fetch only mode
	argsMap[config.PluginModeKey] = true // run as plugin
	args := []string{}
	for k, v := range argsMap {
		args = append(args, formatArgs(k, v))
	}
	binaryPath := nm.latestNodeVersionPath()
	node, err := nm.newNode(binaryPath, args, true)
	if err != nil {
		return 1, fmt.Errorf("couldn't create node: %w", err)
	}
	exitCode := <-node.start()
	return exitCode, nil
}

func formatArgs(k string, v interface{}) string {
	if k == config.CorethConfigKey {
		if val, ok := v.(string); ok && val != config.DefaultString {
			s, _ := json.Marshal(v)
			v = string(s)
		}
	}
	return fmt.Sprintf("--%s=%v", k, v)
}
<commit_msg>fixing coreth flag load<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"sync"

	appplugin "github.com/ava-labs/avalanchego/app/plugin"
	"github.com/ava-labs/avalanchego/config"
	"github.com/ava-labs/avalanchego/node"
	"github.com/ava-labs/avalanchego/utils/constants"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/hashicorp/go-hclog"
	"github.com/hashicorp/go-plugin"
	"github.com/spf13/viper"
)

// nodeProcess wraps a node client
type nodeProcess struct {
	log logging.Logger
	// Every nodeProcess has a unique ID
	processID int
	// [rawClient].Kill() should eventually be called be called
	// on each nodeProcess
	rawClient *plugin.Client
	node      *appplugin.Client
}

// Returns a channel that the node's exit code is sent on when the node is done
// This method does not block.
func (np *nodeProcess) start() chan int {
	exitCodeChan := make(chan int, 1)
	go func() {
		exitCode, err := np.node.Start()
		if err != nil {
			// This error could be from the subprocess shutting down
			np.log.Debug("node returned: %s", err)
		}
		exitCodeChan <- exitCode
	}()
	return exitCodeChan
}

// Stop should be called on each nodeProcess when we are done with it.
// Calls [Kill()] on the underlying client.
func (np *nodeProcess) stop() error {
	err := np.node.Stop()
	np.rawClient.Kill()
	return err
}

type nodeManager struct {
	// Path to the build directory, which should have this structure:
	// build
	// |_avalanchego-latest
	//   |_avalanchego-process (the binary from compiling the app directory)
	//   |_plugins
	//     |_evm
	// |_avalanchego-preupgrade
	//   |_avalanchego-process (the binary from compiling the app directory)
	//   |_plugins
	//     |_evm
	buildDirPath string
	log          logging.Logger
	// nodeProcess ID --> nodeProcess
	nodes         map[int]*nodeProcess
	nextProcessID int
	processLock   sync.Mutex
}

func (nm *nodeManager) latestNodeVersionPath() string {
	return fmt.Sprintf("%s/avalanchego-latest/avalanchego-process", nm.buildDirPath)
}

func (nm *nodeManager) preupgradeNodeVersionPath() string {
	return fmt.Sprintf("%s/avalanchego-preupgrade/avalanchego-process", nm.buildDirPath)
}

func (nm *nodeManager) shutdown() {
	nm.processLock.Lock()
	defer nm.processLock.Unlock()

	for _, node := range nm.nodes {
		nm.log.Info("stopping process %v", node.processID)
		if err := nm.stop(node.processID); err != nil {
			nm.log.Error("error stopping node: %s", err)
		}
		nm.log.Info("done stopping process %v", node.processID)
	}
}

// stop a node. Blocks until the node is done shutting down.
// Assumes [nm.processLock] is not held
func (nm *nodeManager) Stop(processID int) error {
	nm.processLock.Lock()
	defer nm.processLock.Unlock()
	return nm.stop(processID)
}

// stop a node. Blocks until the node is done shutting down.
// Assumes [nm.processLock] is held
func (nm *nodeManager) stop(processID int) error {
	nodeProcess, exists := nm.nodes[processID]
	if !exists {
		return nil
	}
	if err := nodeProcess.stop(); err != nil {
		return err
	}
	delete(nm.nodes, processID)
	return nil
}

func newNodeManager(path string, log logging.Logger) *nodeManager {
	return &nodeManager{
		buildDirPath: path,
		log:          log,
		nodes:        map[int]*nodeProcess{},
		processLock:  sync.Mutex{},
	}
}

// Return a wrapper around a node wunning the binary at [path] with args [args].
// The returned nodeProcess must eventually have [nodeProcess.rawClient.Kill] called on it.
func (nm *nodeManager) newNode(path string, args []string, printToStdOut bool) (*nodeProcess, error) {
	clientConfig := &plugin.ClientConfig{
		HandshakeConfig: appplugin.Handshake,
		Plugins:         appplugin.PluginMap,
		Cmd:             exec.Command(path, args...),
		AllowedProtocols: []plugin.Protocol{
			plugin.ProtocolNetRPC,
			plugin.ProtocolGRPC,
		},
		Logger: hclog.New(&hclog.LoggerOptions{Level: hclog.Warn}),
	}
	if printToStdOut {
		clientConfig.SyncStdout = os.Stdout
		clientConfig.SyncStderr = os.Stderr
	}
	client := plugin.NewClient(clientConfig)
	rpcClient, err := client.Client()
	if err != nil {
		client.Kill()
		return nil, fmt.Errorf("couldn't get Client: %w", err)
	}
	raw, err := rpcClient.Dispense("nodeProcess")
	if err != nil {
		client.Kill()
		return nil, fmt.Errorf("couldn't dispense plugin '%s': %w", "nodeProcess", err)
	}
	node, ok := raw.(*appplugin.Client)
	if !ok {
		client.Kill()
		return nil, fmt.Errorf("expected *node.NodeClient but got %T", raw)
	}
	nm.nextProcessID++
	np := &nodeProcess{
		log:       nm.log,
		node:      node,
		rawClient: client,
		processID: nm.nextProcessID,
	}
	nm.nodes[np.processID] = np
	return np, nil
}

// Start a node compatible with the previous database version
// Override the staking port, HTTP port and plugin directory of the node.
// Assumes the node binary path is [buildDir]/avalanchego-preupgrade/avalanchego-process
// Assumes the node's plugin path is [buildDir]/avalanchego-preupgrade/plugins
// Assumes the binary can be served as a plugin
func (nm *nodeManager) preDBUpgradeNode(v *viper.Viper) (*nodeProcess, error) {
	argsMap := v.AllSettings()
	delete(argsMap, config.FetchOnlyKey)
	argsMap[config.PluginModeKey] = true
	argsMap[config.PluginDirKey] = fmt.Sprintf("%s/avalanchego-preupgrade/plugins", nm.buildDirPath)
	args := []string{}
	for k, v := range argsMap { // Pass args to subprocess
		args = append(args, formatArgs(k, v))
	}
	binaryPath := nm.preupgradeNodeVersionPath()
	return nm.newNode(binaryPath, args, true)
}

// Run the latest node version
func (nm *nodeManager) latestVersionNodeFetchOnly(v *viper.Viper, nodeConfig node.Config) (*nodeProcess, error) {
	argsMap := v.AllSettings()
	argsMap[config.BootstrapIPsKey] = fmt.Sprintf("127.0.0.1:%d", int(nodeConfig.StakingIP.Port))
	argsMap[config.BootstrapIDsKey] = fmt.Sprintf("%s%s", constants.NodeIDPrefix, nodeConfig.NodeID)
	argsMap[config.FetchOnlyKey] = true
	argsMap[config.StakingPortKey] = 0
	argsMap[config.HTTPPortKey] = 0
	argsMap[config.PluginModeKey] = true
	// replace the last folder named daemon in path with fetch-only
	daemonLogDir := nodeConfig.LoggingConfig.Directory
	i := strings.LastIndex(daemonLogDir, "daemon")
	argsMap[config.LogsDirKey] = daemonLogDir[:i] + strings.Replace(daemonLogDir[i:], "daemon", "fetch-only", 1)

	var args []string
	for k, v := range argsMap {
		args = append(args, formatArgs(k, v))
	}
	binaryPath := nm.latestNodeVersionPath()
	return nm.newNode(binaryPath, args, false)
}

// Run the latest node version with the config given by [v].
// Returns the node's exit code.
func (nm *nodeManager) runNormal(v *viper.Viper) (int, error) {
	nm.log.Info("starting latest node version")
	argsMap := v.AllSettings()
	argsMap[config.FetchOnlyKey] = false // don't run in fetch only mode
	argsMap[config.PluginModeKey] = true // run as plugin
	args := []string{}
	for k, v := range argsMap {
		args = append(args, formatArgs(k, v))
	}
	binaryPath := nm.latestNodeVersionPath()
	node, err := nm.newNode(binaryPath, args, true)
	if err != nil {
		return 1, fmt.Errorf("couldn't create node: %w", err)
	}
	exitCode := <-node.start()
	return exitCode, nil
}

func formatArgs(k string, v interface{}) string {
	if k == config.CorethConfigKey {
		// it either is a string, "default" or a json
		if val, ok := v.(string); ok {
			return fmt.Sprintf("--%s=%s", k, val)
		}
		// or it's a loaded config from either the defaults or a config file
		s, _ := json.Marshal(v)
		v = string(s)
	}
	return fmt.Sprintf("--%s=%v", k, v)
}
<|endoftext|>


<commit_before>package marathon

import (
	"testing"

	"github.com/jarcoal/httpmock"
	"github.com/stretchr/testify/assert"
)

func TestCreateApp(t *testing.T) {
	httpmock.Activate()
	defer httpmock.DeactivateAndReset()

	httpmock.RegisterResponder("POST", "https://api.marathon.com/v2/apps",
		httpmock.NewStringResponder(200, `{"id":"/utopia","cmd":null,"args":["production"],"user":null,"env":{},"instances":2,"cpus":2,"mem":1024,"disk":0,"executor":"","constraints":[],"uris":["/root/.dockercfg"],"storeUrls":[],"ports":[10003],"requirePorts":false,"backoffSeconds":1,"backoffFactor":1.15,"container":{"type":"DOCKER","volumes":[{"containerPath":"/var/run/docker.sock","hostPath":"/var/run/docker.sock","mode":"RW"}],"docker":{"image":"docker-prd.itriagehealth.com/utopia:0.1.6","network":null,"portMappings":null}},"healthChecks":[{"path":"/heartbeat","protocol":"HTTP","portIndex":0,"command":null,"gracePeriodSeconds":10,"intervalSeconds":10,"timeoutSeconds":10,"maxConsecutiveFailures":10}],"dependencies":[],"upgradeStrategy":{"minimumHealthCapacity":1},"version":"2015-01-07T18:59:38.310Z"}`))
	health := createHealthCheck()

	docker := DockerProperties{
		Image: "docker-prd.itriagehealth.com/utopia:0.1.6",
	}
	container := Container{
		Type:   "DOCKER",
		Docker: docker,
	}
	testApp := App{
		ID:           "/utopia",
		Args:         []string{"production"},
		CMD:          "",
		Instances:    2,
		CPUS:         2,
		Mem:          1024,
		HealthChecks: health,
		Container:    container,
	}
	c := NewClient("https://api.marathon.com")
	app, err := c.App.Create(testApp)
	assert.Equal(t, err, nil, "err should be nil")
	assert.Equal(t, app.Instances, 2, "There should be 2 instances")
	assert.Equal(t, app.CPUS, 2, "There should be 2 cpus")
	assert.Equal(t, app.Mem, 1024, "There should be 1024mb")
}

func createHealthCheck() []HealthCheck {
	var healthCheck = HealthCheck{
		Protocol:               "HTTP",
		Path:                   "/heartbeat",
		GracePeriodSeconds:     5,
		IntervalSeconds:        10,
		TimeoutSeconds:         10,
		MaxConsecutiveFailures: 3,
		PortIndex:              0,
	}
	healthChecks := []HealthCheck{healthCheck}
	return healthChecks
}
<commit_msg>adds a test for deploy update<commit_after>package marathon

import (
	"testing"

	"github.com/jarcoal/httpmock"
	"github.com/stretchr/testify/assert"
)

func TestCreateApp(t *testing.T) {
	httpmock.Activate()
	defer httpmock.DeactivateAndReset()

	httpmock.RegisterResponder("POST", "https://api.marathon.com/v2/apps",
		httpmock.NewStringResponder(200, `{"id":"/utopia","cmd":null,"args":["production"],"user":null,"env":{},"instances":2,"cpus":2,"mem":1024,"disk":0,"executor":"","constraints":[],"uris":["/root/.dockercfg"],"storeUrls":[],"ports":[10003],"requirePorts":false,"backoffSeconds":1,"backoffFactor":1.15,"container":{"type":"DOCKER","volumes":[{"containerPath":"/var/run/docker.sock","hostPath":"/var/run/docker.sock","mode":"RW"}],"docker":{"image":"docker-prd.itriagehealth.com/utopia:0.1.6","network":null,"portMappings":null}},"healthChecks":[{"path":"/heartbeat","protocol":"HTTP","portIndex":0,"command":null,"gracePeriodSeconds":10,"intervalSeconds":10,"timeoutSeconds":10,"maxConsecutiveFailures":10}],"dependencies":[],"upgradeStrategy":{"minimumHealthCapacity":1},"version":"2015-01-07T18:59:38.310Z"}`))
	health := createHealthCheck()

	docker := DockerProperties{
		Image: "docker-prd.itriagehealth.com/utopia:0.1.6",
	}
	container := Container{
		Type:   "DOCKER",
		Docker: docker,
	}
	testApp := App{
		ID:           "/utopia",
		Args:         []string{"production"},
		CMD:          "",
		Instances:    2,
		CPUS:         2,
		Mem:          1024,
		HealthChecks: health,
		Container:    container,
	}
	c := NewClient("https://api.marathon.com")
	app, err := c.App.Create(testApp)
	assert.Equal(t, err, nil, "err should be nil")
	assert.Equal(t, app.Instances, 2, "There should be 2 instances")
	assert.Equal(t, app.CPUS, 2, "There should be 2 cpus")
	assert.Equal(t, app.Mem, 1024, "There should be 1024mb")
}

func TestUpdateApp(t *testing.T) {
	httpmock.Activate()
	defer httpmock.DeactivateAndReset()

	httpmock.RegisterResponder("PUT", "https://api.marathon.com/v2/apps/utopia",
		httpmock.NewStringResponder(200, `{"id":"/utopia","cmd":null,"args":["production"],"user":null,"env":{},"instances":1,"cpus":2,"mem":1024,"disk":0,"executor":"","constraints":[],"uris":["/root/.dockercfg"],"storeUrls":[],"ports":[10003],"requirePorts":false,"backoffSeconds":1,"backoffFactor":1.15,"container":{"type":"DOCKER","volumes":[{"containerPath":"/var/run/docker.sock","hostPath":"/var/run/docker.sock","mode":"RW"}],"docker":{"image":"docker-prd.itriagehealth.com/utopia:0.1.6","network":null,"portMappings":null}},"healthChecks":[{"path":"/heartbeat","protocol":"HTTP","portIndex":0,"command":null,"gracePeriodSeconds":10,"intervalSeconds":10,"timeoutSeconds":10,"maxConsecutiveFailures":10}],"dependencies":[],"upgradeStrategy":{"minimumHealthCapacity":1},"version":"2015-01-07T18:59:38.310Z"}`))

	testApp := App{
		Instances: 1,
	}
	c := NewClient("https://api.marathon.com")
	app, err := c.App.Update("utopia", testApp)
	assert.Equal(t, err, nil, "err should be nil")
	assert.Equal(t, app.Instances, 1, "There should be 2 instances")
	assert.Equal(t, app.Mem, 1024, "There should be 1024mb of ram")
}

func createHealthCheck() []HealthCheck {
	var healthCheck = HealthCheck{
		Protocol:               "HTTP",
		Path:                   "/heartbeat",
		GracePeriodSeconds:     5,
		IntervalSeconds:        10,
		TimeoutSeconds:         10,
		MaxConsecutiveFailures: 3,
		PortIndex:              0,
	}
	healthChecks := []HealthCheck{healthCheck}
	return healthChecks
}
<|endoftext|>


<commit_before>// Copyright 2016 tsuru-client authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package installer

import (
	"fmt"

	"github.com/docker/engine-api/types/swarm"
	"github.com/fsouza/go-dockerclient"
)

var swarmPort = 2377

type dockerEnpoint interface {
	dockerClient() (*docker.Client, error)
	GetNetwork() *docker.Network
}

type SwarmCluster struct {
	Manager *Machine
	Workers []*Machine
}

func (c *SwarmCluster) dockerClient() (*docker.Client, error) {
	return c.Manager.dockerClient()
}

func (c *SwarmCluster) GetNetwork() *docker.Network {
	return c.Manager.GetNetwork()
}

// NewSwarmCluster creates a Swarm Cluster using the first machine as a manager
// and the rest as workers and also creates an overlay network between the nodes.
func NewSwarmCluster(machines []*Machine) (*SwarmCluster, error) {
	swarmOpts := docker.InitSwarmOptions{
		InitRequest: swarm.InitRequest{
			ListenAddr:    fmt.Sprintf("0.0.0.0:%d", swarmPort),
			AdvertiseAddr: fmt.Sprintf("%s:%d", machines[0].IP, swarmPort),
		},
	}
	dockerClient, err := machines[0].dockerClient()
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve machine %s docker client: %s", machines[0].Name, err)
	}
	_, err = dockerClient.InitSwarm(swarmOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to init swarm: %s", err)
	}
	swarmInspect, err := dockerClient.InspectSwarm(nil)
	if err != nil {
		return nil, fmt.Errorf("failed to inspect swarm: %s", err)
	}
	createNetworkOpts := docker.CreateNetworkOptions{
		Name:           "tsuru",
		Driver:         "overlay",
		CheckDuplicate: true,
		IPAM: docker.IPAMOptions{
			Driver: "default",
			Config: []docker.IPAMConfig{
				{
					Subnet: "10.0.9.0/24",
				},
			},
		},
	}
	network, err := dockerClient.CreateNetwork(createNetworkOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to create overlay network: %s", err)
	}
	for i, m := range machines {
		m.network = network
		if i == 0 {
			continue
		}
		dockerClient, err = m.dockerClient()
		if err != nil {
			return nil, fmt.Errorf("failed to retrieve machine %s docker client: %s", m.Name, err)
		}
		opts := docker.JoinSwarmOptions{
			JoinRequest: swarm.JoinRequest{
				ListenAddr:    fmt.Sprintf("0.0.0.0:%d", swarmPort),
				AdvertiseAddr: fmt.Sprintf("%s:%d", m.IP, swarmPort),
				JoinToken:     swarmInspect.JoinTokens.Worker,
				RemoteAddrs:   []string{fmt.Sprintf("%s:%d", machines[0].IP, swarmPort)},
			},
		}
		err = dockerClient.JoinSwarm(opts)
		if err != nil {
			return nil, fmt.Errorf("machine %s failed to join swarm: %s", m.Name, err)
		}
	}
	return &SwarmCluster{
		Manager: machines[0],
		Workers: machines,
	}, nil
}

// ServiceTaskExec finds a container running a service task and runs exec on it
func (c *SwarmCluster) ServiceExec(service string, cmd []string, startOpts docker.StartExecOptions) error {
	mClient, err := c.dockerClient()
	if err != nil {
		return fmt.Errorf("failed to retrive swarm docker client: %s", err)
	}
	tasks, err := mClient.ListTasks(docker.ListTasksOptions{
		Filters: map[string][]string{
			"service":       {service},
			"desired-state": {"running"},
		},
	})
	if err != nil {
		return fmt.Errorf("failed to list tasks for service %s: %s", service, err)
	}
	if len(tasks) == 0 {
		return fmt.Errorf("no running task found for service %s", service)
	}
	node, err := mClient.InspectNode(tasks[0].NodeID)
	if err != nil {
		return fmt.Errorf("failed to inspect node %s: %s", tasks[0].NodeID, err)
	}
	nodeName := node.Description.Hostname
	var machine *Machine
	for _, m := range c.Workers {
		if m.Name == nodeName {
			machine = m
			break
		}
	}
	if machine == nil {
		return fmt.Errorf("machine %s not found", nodeName)
	}
	client, err := machine.dockerClient()
	if err != nil {
		return fmt.Errorf("failed to retrieve task node %s docker client: %s", machine.Name, err)
	}
	container := tasks[0].Status.ContainerStatus.ContainerID
	exec, err := client.CreateExec(docker.CreateExecOptions{
		Cmd:          cmd,
		AttachStdout: true,
		AttachStderr: true,
		AttachStdin:  true,
		Container:    container,
	})
	if err != nil {
		return fmt.Errorf("failed to exec in task container %s: %s", container, err)
	}
	return client.StartExec(exec.ID, startOpts)
}

// CreateService creates a service on the swarm cluster
func (c *SwarmCluster) CreateService(opts docker.CreateServiceOptions) error {
	client, err := c.dockerClient()
	if err != nil {
		return err
	}
	opts.Networks = []swarm.NetworkAttachmentConfig{
		{Target: c.GetNetwork().Name},
	}
	_, err = client.CreateService(opts)
	return err
}

func (c *SwarmCluster) ListNodes() ([]swarm.Node, error) {
	client, err := c.dockerClient()
	if err != nil {
		return nil, err
	}
	return client.ListNodes(docker.ListNodesOptions{})
}

func (c *SwarmCluster) ListServices() ([]swarm.Service, error) {
	client, err := c.dockerClient()
	if err != nil {
		return nil, err
	}
	return client.ListServices(docker.ListServicesOptions{})
}
<commit_msg>installer/swarm: fix comment<commit_after>// Copyright 2016 tsuru-client authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package installer

import (
	"fmt"

	"github.com/docker/engine-api/types/swarm"
	"github.com/fsouza/go-dockerclient"
)

var swarmPort = 2377

type dockerEnpoint interface {
	dockerClient() (*docker.Client, error)
	GetNetwork() *docker.Network
}

type SwarmCluster struct {
	Manager *Machine
	Workers []*Machine
}

func (c *SwarmCluster) dockerClient() (*docker.Client, error) {
	return c.Manager.dockerClient()
}

func (c *SwarmCluster) GetNetwork() *docker.Network {
	return c.Manager.GetNetwork()
}

// NewSwarmCluster creates a Swarm Cluster using the first machine as a manager
// and the rest as workers and also creates an overlay network between the nodes.
func NewSwarmCluster(machines []*Machine) (*SwarmCluster, error) {
	swarmOpts := docker.InitSwarmOptions{
		InitRequest: swarm.InitRequest{
			ListenAddr:    fmt.Sprintf("0.0.0.0:%d", swarmPort),
			AdvertiseAddr: fmt.Sprintf("%s:%d", machines[0].IP, swarmPort),
		},
	}
	dockerClient, err := machines[0].dockerClient()
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve machine %s docker client: %s", machines[0].Name, err)
	}
	_, err = dockerClient.InitSwarm(swarmOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to init swarm: %s", err)
	}
	swarmInspect, err := dockerClient.InspectSwarm(nil)
	if err != nil {
		return nil, fmt.Errorf("failed to inspect swarm: %s", err)
	}
	createNetworkOpts := docker.CreateNetworkOptions{
		Name:           "tsuru",
		Driver:         "overlay",
		CheckDuplicate: true,
		IPAM: docker.IPAMOptions{
			Driver: "default",
			Config: []docker.IPAMConfig{
				{
					Subnet: "10.0.9.0/24",
				},
			},
		},
	}
	network, err := dockerClient.CreateNetwork(createNetworkOpts)
	if err != nil {
		return nil, fmt.Errorf("failed to create overlay network: %s", err)
	}
	for i, m := range machines {
		m.network = network
		if i == 0 {
			continue
		}
		dockerClient, err = m.dockerClient()
		if err != nil {
			return nil, fmt.Errorf("failed to retrieve machine %s docker client: %s", m.Name, err)
		}
		opts := docker.JoinSwarmOptions{
			JoinRequest: swarm.JoinRequest{
				ListenAddr:    fmt.Sprintf("0.0.0.0:%d", swarmPort),
				AdvertiseAddr: fmt.Sprintf("%s:%d", m.IP, swarmPort),
				JoinToken:     swarmInspect.JoinTokens.Worker,
				RemoteAddrs:   []string{fmt.Sprintf("%s:%d", machines[0].IP, swarmPort)},
			},
		}
		err = dockerClient.JoinSwarm(opts)
		if err != nil {
			return nil, fmt.Errorf("machine %s failed to join swarm: %s", m.Name, err)
		}
	}
	return &SwarmCluster{
		Manager: machines[0],
		Workers: machines,
	}, nil
}

// ServiceExec finds a container running a service task and runs exec on it
func (c *SwarmCluster) ServiceExec(service string, cmd []string, startOpts docker.StartExecOptions) error {
	mClient, err := c.dockerClient()
	if err != nil {
		return fmt.Errorf("failed to retrive swarm docker client: %s", err)
	}
	tasks, err := mClient.ListTasks(docker.ListTasksOptions{
		Filters: map[string][]string{
			"service":       {service},
			"desired-state": {"running"},
		},
	})
	if err != nil {
		return fmt.Errorf("failed to list tasks for service %s: %s", service, err)
	}
	if len(tasks) == 0 {
		return fmt.Errorf("no running task found for service %s", service)
	}
	node, err := mClient.InspectNode(tasks[0].NodeID)
	if err != nil {
		return fmt.Errorf("failed to inspect node %s: %s", tasks[0].NodeID, err)
	}
	nodeName := node.Description.Hostname
	var machine *Machine
	for _, m := range c.Workers {
		if m.Name == nodeName {
			machine = m
			break
		}
	}
	if machine == nil {
		return fmt.Errorf("machine %s not found", nodeName)
	}
	client, err := machine.dockerClient()
	if err != nil {
		return fmt.Errorf("failed to retrieve task node %s docker client: %s", machine.Name, err)
	}
	container := tasks[0].Status.ContainerStatus.ContainerID
	exec, err := client.CreateExec(docker.CreateExecOptions{
		Cmd:          cmd,
		AttachStdout: true,
		AttachStderr: true,
		AttachStdin:  true,
		Container:    container,
	})
	if err != nil {
		return fmt.Errorf("failed to exec in task container %s: %s", container, err)
	}
	return client.StartExec(exec.ID, startOpts)
}

// CreateService creates a service on the swarm cluster
func (c *SwarmCluster) CreateService(opts docker.CreateServiceOptions) error {
	client, err := c.dockerClient()
	if err != nil {
		return err
	}
	opts.Networks = []swarm.NetworkAttachmentConfig{
		{Target: c.GetNetwork().Name},
	}
	_, err = client.CreateService(opts)
	return err
}

func (c *SwarmCluster) ListNodes() ([]swarm.Node, error) {
	client, err := c.dockerClient()
	if err != nil {
		return nil, err
	}
	return client.ListNodes(docker.ListNodesOptions{})
}

func (c *SwarmCluster) ListServices() ([]swarm.Service, error) {
	client, err := c.dockerClient()
	if err != nil {
		return nil, err
	}
	return client.ListServices(docker.ListServicesOptions{})
}
<|endoftext|>