<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package binlog

import (
	"fmt"
	"sync"

	log "github.com/golang/glog"
	mproto "github.com/youtube/vitess/go/mysql/proto"
	"github.com/youtube/vitess/go/stats"
	"github.com/youtube/vitess/go/sync2"
	"github.com/youtube/vitess/go/tb"
	"github.com/youtube/vitess/go/vt/binlog/proto"
	"github.com/youtube/vitess/go/vt/key"
	"github.com/youtube/vitess/go/vt/mysqlctl"
	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"

	pb "github.com/youtube/vitess/go/vt/proto/topodata"
)

/* API and config for UpdateStream Service */

const (
	usDisabled int64 = iota
	usEnabled
)

var usStateNames = map[int64]string{
	usEnabled:  "Enabled",
	usDisabled: "Disabled",
}

var (
	streamCount          = stats.NewCounters("UpdateStreamStreamCount")
	updateStreamErrors   = stats.NewCounters("UpdateStreamErrors")
	updateStreamEvents   = stats.NewCounters("UpdateStreamEvents")
	keyrangeStatements   = stats.NewInt("UpdateStreamKeyRangeStatements")
	keyrangeTransactions = stats.NewInt("UpdateStreamKeyRangeTransactions")
	tablesStatements     = stats.NewInt("UpdateStreamTablesStatements")
	tablesTransactions   = stats.NewInt("UpdateStreamTablesTransactions")
)

// UpdateStreamControl is the interface an UpdateStream service implements
// to bring it up or down.
type UpdateStreamControl interface {
	// Enable will allow any new RPC calls
	Enable()

	// Disable will interrupt all current calls, and disallow any new call
	Disable()

	// IsEnabled returns true iff the service is enabled
	IsEnabled() bool
}

// UpdateStreamControlMock is an implementation of UpdateStreamControl
// to be used in tests
type UpdateStreamControlMock struct {
	enabled bool
}

// NewUpdateStreamControlMock creates a new UpdateStreamControlMock
func NewUpdateStreamControlMock() *UpdateStreamControlMock {
	return &UpdateStreamControlMock{}
}

// Enable is part of UpdateStreamControl
func (m *UpdateStreamControlMock) Enable() {
	m.enabled = true
}

// Disable is part of UpdateStreamControl
func (m *UpdateStreamControlMock) Disable() {
	m.enabled = false
}

// IsEnabled is part of UpdateStreamControl
func (m *UpdateStreamControlMock) IsEnabled() bool {
	return m.enabled
}

// UpdateStream is the real implementation of proto.UpdateStream
// and UpdateStreamControl
type UpdateStream struct {
	// the following variables are set at construction time

	mysqld mysqlctl.MysqlDaemon
	dbname string

	// actionLock protects the following variables
	actionLock     sync.Mutex
	state          sync2.AtomicInt64
	stateWaitGroup sync.WaitGroup
	streams        streamList
}

type streamList struct {
	sync.Mutex
	streams map[*sync2.ServiceManager]bool
}

func (sl *streamList) Init() {
	sl.Lock()
	sl.streams = make(map[*sync2.ServiceManager]bool)
	sl.Unlock()
}

func (sl *streamList) Add(e *sync2.ServiceManager) {
	sl.Lock()
	sl.streams[e] = true
	sl.Unlock()
}

func (sl *streamList) Delete(e *sync2.ServiceManager) {
	sl.Lock()
	delete(sl.streams, e)
	sl.Unlock()
}

func (sl *streamList) Stop() {
	sl.Lock()
	for stream := range sl.streams {
		stream.Stop()
	}
	sl.Unlock()
}

// RegisterUpdateStreamServiceFunc is the type to use for delayed
// registration of RPC servers until we have all the objects
type RegisterUpdateStreamServiceFunc func(proto.UpdateStream)

// RegisterUpdateStreamServices is the list of all registration
// callbacks to invoke
var RegisterUpdateStreamServices []RegisterUpdateStreamServiceFunc

// NewUpdateStream returns a new UpdateStream object
func NewUpdateStream(mysqld mysqlctl.MysqlDaemon, dbname string) *UpdateStream {
	return &UpdateStream{
		mysqld: mysqld,
		dbname: dbname,
	}
}

// RegisterService needs to be called to publish stats, and to start listening
// to clients. Only once instance can call this in a process.
func (updateStream *UpdateStream) RegisterService() {
	// publish the stats
	stats.Publish("UpdateStreamState", stats.StringFunc(func() string {
		return usStateNames[updateStream.state.Get()]
	}))

	// and register all the RPC protocols
	for _, f := range RegisterUpdateStreamServices {
		f(updateStream)
	}
}

func logError() {
	if x := recover(); x != nil {
		log.Errorf("%s at\n%s", x.(error).Error(), tb.Stack(4))
	}
}

// Enable will allow connections to the service
func (updateStream *UpdateStream) Enable() {
	defer logError()
	updateStream.actionLock.Lock()
	defer updateStream.actionLock.Unlock()
	if updateStream.IsEnabled() {
		return
	}

	updateStream.state.Set(usEnabled)
	updateStream.streams.Init()
	log.Infof("Enabling update stream, dbname: %s", updateStream.dbname)
}

// Disable will disallow any connection to the service
func (updateStream *UpdateStream) Disable() {
	defer logError()
	updateStream.actionLock.Lock()
	defer updateStream.actionLock.Unlock()
	if !updateStream.IsEnabled() {
		return
	}

	updateStream.state.Set(usDisabled)
	updateStream.streams.Stop()
	updateStream.stateWaitGroup.Wait()
	log.Infof("Update Stream Disabled")
}

// IsEnabled returns true if UpdateStream is enabled
func (updateStream *UpdateStream) IsEnabled() bool {
	return updateStream.state.Get() == usEnabled
}

// ServeUpdateStream is part of the proto.UpdateStream interface
func (updateStream *UpdateStream) ServeUpdateStream(position string, sendReply func(reply *proto.StreamEvent) error) (err error) {
	pos, err := myproto.DecodeReplicationPosition(position)
	if err != nil {
		return err
	}

	updateStream.actionLock.Lock()
	if !updateStream.IsEnabled() {
		updateStream.actionLock.Unlock()
		log.Errorf("Unable to serve client request: update stream service is not enabled")
		return fmt.Errorf("update stream service is not enabled")
	}
	updateStream.stateWaitGroup.Add(1)
	updateStream.actionLock.Unlock()
	defer updateStream.stateWaitGroup.Done()

	streamCount.Add("Updates", 1)
	defer streamCount.Add("Updates", -1)
	log.Infof("ServeUpdateStream starting @ %#v", pos)

	evs := NewEventStreamer(updateStream.dbname, updateStream.mysqld, pos, func(reply *proto.StreamEvent) error {
		if reply.Category == "ERR" {
			updateStreamErrors.Add("UpdateStream", 1)
		} else {
			updateStreamEvents.Add(reply.Category, 1)
		}
		return sendReply(reply)
	})

	svm := &sync2.ServiceManager{}
	svm.Go(evs.Stream)
	updateStream.streams.Add(svm)
	defer updateStream.streams.Delete(svm)
	return svm.Join()
}

// StreamKeyRange is part of the proto.UpdateStream interface
func (updateStream *UpdateStream) StreamKeyRange(position string, keyspaceIDType key.KeyspaceIdType, keyRange *pb.KeyRange, charset *mproto.Charset, sendReply func(reply *proto.BinlogTransaction) error) (err error) {
	pos, err := myproto.DecodeReplicationPosition(position)
	if err != nil {
		return err
	}

	updateStream.actionLock.Lock()
	if !updateStream.IsEnabled() {
		updateStream.actionLock.Unlock()
		log.Errorf("Unable to serve client request: Update stream service is not enabled")
		return fmt.Errorf("update stream service is not enabled")
	}
	updateStream.stateWaitGroup.Add(1)
	updateStream.actionLock.Unlock()
	defer updateStream.stateWaitGroup.Done()

	streamCount.Add("KeyRange", 1)
	defer streamCount.Add("KeyRange", -1)
	log.Infof("ServeUpdateStream starting @ %#v", pos)

	// Calls cascade like this: BinlogStreamer->KeyRangeFilterFunc->func(*proto.BinlogTransaction)->sendReply
	f := KeyRangeFilterFunc(keyspaceIDType, keyRange, func(reply *proto.BinlogTransaction) error {
		keyrangeStatements.Add(int64(len(reply.Statements)))
		keyrangeTransactions.Add(1)
		return sendReply(reply)
	})
	bls := NewBinlogStreamer(updateStream.dbname, updateStream.mysqld, charset, pos, f)

	svm := &sync2.ServiceManager{}
	svm.Go(bls.Stream)
	updateStream.streams.Add(svm)
	defer updateStream.streams.Delete(svm)
	return svm.Join()
}

// StreamTables is part of the proto.UpdateStream interface
func (updateStream *UpdateStream) StreamTables(position string, tables []string, charset *mproto.Charset, sendReply func(reply *proto.BinlogTransaction) error) (err error) {
	pos, err := myproto.DecodeReplicationPosition(position)
	if err != nil {
		return err
	}

	updateStream.actionLock.Lock()
	if !updateStream.IsEnabled() {
		updateStream.actionLock.Unlock()
		log.Errorf("Unable to serve client request: Update stream service is not enabled")
		return fmt.Errorf("update stream service is not enabled")
	}
	updateStream.stateWaitGroup.Add(1)
	updateStream.actionLock.Unlock()
	defer updateStream.stateWaitGroup.Done()

	streamCount.Add("Tables", 1)
	defer streamCount.Add("Tables", -1)
	log.Infof("ServeUpdateStream starting @ %#v", pos)

	// Calls cascade like this: BinlogStreamer->TablesFilterFunc->func(*proto.BinlogTransaction)->sendReply
	f := TablesFilterFunc(tables, func(reply *proto.BinlogTransaction) error {
		keyrangeStatements.Add(int64(len(reply.Statements)))
		keyrangeTransactions.Add(1)
		return sendReply(reply)
	})
	bls := NewBinlogStreamer(updateStream.dbname, updateStream.mysqld, charset, pos, f)

	svm := &sync2.ServiceManager{}
	svm.Go(bls.Stream)
	updateStream.streams.Add(svm)
	defer updateStream.streams.Delete(svm)
	return svm.Join()
}

// HandlePanic is part of the proto.UpdateStream interface
func (updateStream *UpdateStream) HandlePanic(err *error) {
	if x := recover(); x != nil {
		log.Errorf("Uncaught panic:\n%v\n%s", x, tb.Stack(4))
		*err = fmt.Errorf("uncaught panic: %v", x)
	}
}
<commit_msg>Fix data race in UpdateStreamControlMock.<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package binlog

import (
	"fmt"
	"sync"

	log "github.com/golang/glog"
	mproto "github.com/youtube/vitess/go/mysql/proto"
	"github.com/youtube/vitess/go/stats"
	"github.com/youtube/vitess/go/sync2"
	"github.com/youtube/vitess/go/tb"
	"github.com/youtube/vitess/go/vt/binlog/proto"
	"github.com/youtube/vitess/go/vt/key"
	"github.com/youtube/vitess/go/vt/mysqlctl"
	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"

	pb "github.com/youtube/vitess/go/vt/proto/topodata"
)

/* API and config for UpdateStream Service */

const (
	usDisabled int64 = iota
	usEnabled
)

var usStateNames = map[int64]string{
	usEnabled:  "Enabled",
	usDisabled: "Disabled",
}

var (
	streamCount          = stats.NewCounters("UpdateStreamStreamCount")
	updateStreamErrors   = stats.NewCounters("UpdateStreamErrors")
	updateStreamEvents   = stats.NewCounters("UpdateStreamEvents")
	keyrangeStatements   = stats.NewInt("UpdateStreamKeyRangeStatements")
	keyrangeTransactions = stats.NewInt("UpdateStreamKeyRangeTransactions")
	tablesStatements     = stats.NewInt("UpdateStreamTablesStatements")
	tablesTransactions   = stats.NewInt("UpdateStreamTablesTransactions")
)

// UpdateStreamControl is the interface an UpdateStream service implements
// to bring it up or down.
type UpdateStreamControl interface {
	// Enable will allow any new RPC calls
	Enable()

	// Disable will interrupt all current calls, and disallow any new call
	Disable()

	// IsEnabled returns true iff the service is enabled
	IsEnabled() bool
}

// UpdateStreamControlMock is an implementation of UpdateStreamControl
// to be used in tests
type UpdateStreamControlMock struct {
	enabled bool
	sync.Mutex
}

// NewUpdateStreamControlMock creates a new UpdateStreamControlMock
func NewUpdateStreamControlMock() *UpdateStreamControlMock {
	return &UpdateStreamControlMock{}
}

// Enable is part of UpdateStreamControl
func (m *UpdateStreamControlMock) Enable() {
	m.Lock()
	m.enabled = true
	m.Unlock()
}

// Disable is part of UpdateStreamControl
func (m *UpdateStreamControlMock) Disable() {
	m.Lock()
	m.enabled = false
	m.Unlock()
}

// IsEnabled is part of UpdateStreamControl
func (m *UpdateStreamControlMock) IsEnabled() bool {
	m.Lock()
	defer m.Unlock()
	return m.enabled
}

// UpdateStream is the real implementation of proto.UpdateStream
// and UpdateStreamControl
type UpdateStream struct {
	// the following variables are set at construction time

	mysqld mysqlctl.MysqlDaemon
	dbname string

	// actionLock protects the following variables
	actionLock     sync.Mutex
	state          sync2.AtomicInt64
	stateWaitGroup sync.WaitGroup
	streams        streamList
}

type streamList struct {
	sync.Mutex
	streams map[*sync2.ServiceManager]bool
}

func (sl *streamList) Init() {
	sl.Lock()
	sl.streams = make(map[*sync2.ServiceManager]bool)
	sl.Unlock()
}

func (sl *streamList) Add(e *sync2.ServiceManager) {
	sl.Lock()
	sl.streams[e] = true
	sl.Unlock()
}

func (sl *streamList) Delete(e *sync2.ServiceManager) {
	sl.Lock()
	delete(sl.streams, e)
	sl.Unlock()
}

func (sl *streamList) Stop() {
	sl.Lock()
	for stream := range sl.streams {
		stream.Stop()
	}
	sl.Unlock()
}

// RegisterUpdateStreamServiceFunc is the type to use for delayed
// registration of RPC servers until we have all the objects
type RegisterUpdateStreamServiceFunc func(proto.UpdateStream)

// RegisterUpdateStreamServices is the list of all registration
// callbacks to invoke
var RegisterUpdateStreamServices []RegisterUpdateStreamServiceFunc

// NewUpdateStream returns a new UpdateStream object
func NewUpdateStream(mysqld mysqlctl.MysqlDaemon, dbname string) *UpdateStream {
	return &UpdateStream{
		mysqld: mysqld,
		dbname: dbname,
	}
}

// RegisterService needs to be called to publish stats, and to start listening
// to clients. Only once instance can call this in a process.
func (updateStream *UpdateStream) RegisterService() {
	// publish the stats
	stats.Publish("UpdateStreamState", stats.StringFunc(func() string {
		return usStateNames[updateStream.state.Get()]
	}))

	// and register all the RPC protocols
	for _, f := range RegisterUpdateStreamServices {
		f(updateStream)
	}
}

func logError() {
	if x := recover(); x != nil {
		log.Errorf("%s at\n%s", x.(error).Error(), tb.Stack(4))
	}
}

// Enable will allow connections to the service
func (updateStream *UpdateStream) Enable() {
	defer logError()
	updateStream.actionLock.Lock()
	defer updateStream.actionLock.Unlock()
	if updateStream.IsEnabled() {
		return
	}

	updateStream.state.Set(usEnabled)
	updateStream.streams.Init()
	log.Infof("Enabling update stream, dbname: %s", updateStream.dbname)
}

// Disable will disallow any connection to the service
func (updateStream *UpdateStream) Disable() {
	defer logError()
	updateStream.actionLock.Lock()
	defer updateStream.actionLock.Unlock()
	if !updateStream.IsEnabled() {
		return
	}

	updateStream.state.Set(usDisabled)
	updateStream.streams.Stop()
	updateStream.stateWaitGroup.Wait()
	log.Infof("Update Stream Disabled")
}

// IsEnabled returns true if UpdateStream is enabled
func (updateStream *UpdateStream) IsEnabled() bool {
	return updateStream.state.Get() == usEnabled
}

// ServeUpdateStream is part of the proto.UpdateStream interface
func (updateStream *UpdateStream) ServeUpdateStream(position string, sendReply func(reply *proto.StreamEvent) error) (err error) {
	pos, err := myproto.DecodeReplicationPosition(position)
	if err != nil {
		return err
	}

	updateStream.actionLock.Lock()
	if !updateStream.IsEnabled() {
		updateStream.actionLock.Unlock()
		log.Errorf("Unable to serve client request: update stream service is not enabled")
		return fmt.Errorf("update stream service is not enabled")
	}
	updateStream.stateWaitGroup.Add(1)
	updateStream.actionLock.Unlock()
	defer updateStream.stateWaitGroup.Done()

	streamCount.Add("Updates", 1)
	defer streamCount.Add("Updates", -1)
	log.Infof("ServeUpdateStream starting @ %#v", pos)

	evs := NewEventStreamer(updateStream.dbname, updateStream.mysqld, pos, func(reply *proto.StreamEvent) error {
		if reply.Category == "ERR" {
			updateStreamErrors.Add("UpdateStream", 1)
		} else {
			updateStreamEvents.Add(reply.Category, 1)
		}
		return sendReply(reply)
	})

	svm := &sync2.ServiceManager{}
	svm.Go(evs.Stream)
	updateStream.streams.Add(svm)
	defer updateStream.streams.Delete(svm)
	return svm.Join()
}

// StreamKeyRange is part of the proto.UpdateStream interface
func (updateStream *UpdateStream) StreamKeyRange(position string, keyspaceIDType key.KeyspaceIdType, keyRange *pb.KeyRange, charset *mproto.Charset, sendReply func(reply *proto.BinlogTransaction) error) (err error) {
	pos, err := myproto.DecodeReplicationPosition(position)
	if err != nil {
		return err
	}

	updateStream.actionLock.Lock()
	if !updateStream.IsEnabled() {
		updateStream.actionLock.Unlock()
		log.Errorf("Unable to serve client request: Update stream service is not enabled")
		return fmt.Errorf("update stream service is not enabled")
	}
	updateStream.stateWaitGroup.Add(1)
	updateStream.actionLock.Unlock()
	defer updateStream.stateWaitGroup.Done()

	streamCount.Add("KeyRange", 1)
	defer streamCount.Add("KeyRange", -1)
	log.Infof("ServeUpdateStream starting @ %#v", pos)

	// Calls cascade like this: BinlogStreamer->KeyRangeFilterFunc->func(*proto.BinlogTransaction)->sendReply
	f := KeyRangeFilterFunc(keyspaceIDType, keyRange, func(reply *proto.BinlogTransaction) error {
		keyrangeStatements.Add(int64(len(reply.Statements)))
		keyrangeTransactions.Add(1)
		return sendReply(reply)
	})
	bls := NewBinlogStreamer(updateStream.dbname, updateStream.mysqld, charset, pos, f)

	svm := &sync2.ServiceManager{}
	svm.Go(bls.Stream)
	updateStream.streams.Add(svm)
	defer updateStream.streams.Delete(svm)
	return svm.Join()
}

// StreamTables is part of the proto.UpdateStream interface
func (updateStream *UpdateStream) StreamTables(position string, tables []string, charset *mproto.Charset, sendReply func(reply *proto.BinlogTransaction) error) (err error) {
	pos, err := myproto.DecodeReplicationPosition(position)
	if err != nil {
		return err
	}

	updateStream.actionLock.Lock()
	if !updateStream.IsEnabled() {
		updateStream.actionLock.Unlock()
		log.Errorf("Unable to serve client request: Update stream service is not enabled")
		return fmt.Errorf("update stream service is not enabled")
	}
	updateStream.stateWaitGroup.Add(1)
	updateStream.actionLock.Unlock()
	defer updateStream.stateWaitGroup.Done()

	streamCount.Add("Tables", 1)
	defer streamCount.Add("Tables", -1)
	log.Infof("ServeUpdateStream starting @ %#v", pos)

	// Calls cascade like this: BinlogStreamer->TablesFilterFunc->func(*proto.BinlogTransaction)->sendReply
	f := TablesFilterFunc(tables, func(reply *proto.BinlogTransaction) error {
		keyrangeStatements.Add(int64(len(reply.Statements)))
		keyrangeTransactions.Add(1)
		return sendReply(reply)
	})
	bls := NewBinlogStreamer(updateStream.dbname, updateStream.mysqld, charset, pos, f)

	svm := &sync2.ServiceManager{}
	svm.Go(bls.Stream)
	updateStream.streams.Add(svm)
	defer updateStream.streams.Delete(svm)
	return svm.Join()
}

// HandlePanic is part of the proto.UpdateStream interface
func (updateStream *UpdateStream) HandlePanic(err *error) {
	if x := recover(); x != nil {
		log.Errorf("Uncaught panic:\n%v\n%s", x, tb.Stack(4))
		*err = fmt.Errorf("uncaught panic: %v", x)
	}
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package pull

import (
	"sync"
	"time"

	"github.com/hyperledger/fabric/gossip/comm"
	"github.com/hyperledger/fabric/gossip/common"
	"github.com/hyperledger/fabric/gossip/discovery"
	"github.com/hyperledger/fabric/gossip/gossip/algo"
	"github.com/hyperledger/fabric/gossip/util"
	proto "github.com/hyperledger/fabric/protos/gossip"
	"github.com/pkg/errors"
	"go.uber.org/zap/zapcore"
)

// Constants go here.
const (
	HelloMsgType MsgType = iota
	DigestMsgType
	RequestMsgType
	ResponseMsgType
)

// MsgType defines the type of a message that is sent to the PullStore
type MsgType int

// MessageHook defines a function that will run after a certain pull message is received
type MessageHook func(itemIDs []string, items []*proto.SignedGossipMessage, msg proto.ReceivedMessage)

// Sender sends messages to remote peers
type Sender interface {
	// Send sends a message to a list of remote peers
	Send(msg *proto.SignedGossipMessage, peers ...*comm.RemotePeer)
}

// MembershipService obtains membership information of alive peers
type MembershipService interface {
	// GetMembership returns the membership of
	GetMembership() []discovery.NetworkMember
}

// Config defines the configuration of the pull mediator
type Config struct {
	ID                string
	PullInterval      time.Duration // Duration between pull invocations
	Channel           common.ChainID
	PeerCountToSelect int // Number of peers to initiate pull with
	Tag               proto.GossipMessage_Tag
	MsgType           proto.PullMsgType
}

// IngressDigestFilter filters out entities in digests that are received from remote peers
type IngressDigestFilter func(digestMsg *proto.DataDigest) *proto.DataDigest

// EgressDigestFilter filters digests to be sent to a remote peer, that
// sent a hello with the following message
type EgressDigestFilter func(helloMsg proto.ReceivedMessage) func(digestItem string) bool

// byContext converts this EgressDigFilter to an algo.DigestFilter
func (df EgressDigestFilter) byContext() algo.DigestFilter {
	return func(context interface{}) func(digestItem string) bool {
		return func(digestItem string) bool {
			return df(context.(proto.ReceivedMessage))(digestItem)
		}
	}
}

// PullAdapter defines methods of the pullStore to interact
// with various modules of gossip
type PullAdapter struct {
	Sndr             Sender
	MemSvc           MembershipService
	IdExtractor      proto.IdentifierExtractor
	MsgCons          proto.MsgConsumer
	EgressDigFilter  EgressDigestFilter
	IngressDigFilter IngressDigestFilter
}

// Mediator is a component wrap a PullEngine and provides the methods
// it needs to perform pull synchronization.
// The specialization of a pull mediator to a certain type of message is
// done by the configuration, a IdentifierExtractor, IdentifierExtractor
// given at construction, and also hooks that can be registered for each
// type of pullMsgType (hello, digest, req, res).
type Mediator interface {
	// Stop stop the Mediator
	Stop()

	// RegisterMsgHook registers a message hook to a specific type of pull message
	RegisterMsgHook(MsgType, MessageHook)

	// Add adds a GossipMessage to the Mediator
	Add(*proto.SignedGossipMessage)

	// Remove removes a GossipMessage from the Mediator with a matching digest,
	// if such a message exits
	Remove(digest string)

	// HandleMessage handles a message from some remote peer
	HandleMessage(msg proto.ReceivedMessage)
}

// pullMediatorImpl is an implementation of Mediator
type pullMediatorImpl struct {
	sync.RWMutex
	*PullAdapter
	msgType2Hook map[MsgType][]MessageHook
	config       Config
	logger       util.Logger
	itemID2Msg   map[string]*proto.SignedGossipMessage
	engine       *algo.PullEngine
}

// NewPullMediator returns a new Mediator
func NewPullMediator(config Config, adapter *PullAdapter) Mediator {
	egressDigFilter := adapter.EgressDigFilter

	acceptAllFilter := func(_ proto.ReceivedMessage) func(string) bool {
		return func(_ string) bool {
			return true
		}
	}

	if egressDigFilter == nil {
		egressDigFilter = acceptAllFilter
	}

	p := &pullMediatorImpl{
		PullAdapter:  adapter,
		msgType2Hook: make(map[MsgType][]MessageHook),
		config:       config,
		logger:       util.GetLogger(util.PullLogger, config.ID),
		itemID2Msg:   make(map[string]*proto.SignedGossipMessage),
	}

	p.engine = algo.NewPullEngineWithFilter(p, config.PullInterval, egressDigFilter.byContext())

	if adapter.IngressDigFilter == nil {
		// Create accept all filter
		adapter.IngressDigFilter = func(digestMsg *proto.DataDigest) *proto.DataDigest {
			return digestMsg
		}
	}
	return p

}

func (p *pullMediatorImpl) HandleMessage(m proto.ReceivedMessage) {
	if m.GetGossipMessage() == nil || !m.GetGossipMessage().IsPullMsg() {
		return
	}
	msg := m.GetGossipMessage()
	msgType := msg.GetPullMsgType()
	if msgType != p.config.MsgType {
		return
	}

	p.logger.Debug(msg)

	itemIDs := []string{}
	items := []*proto.SignedGossipMessage{}
	var pullMsgType MsgType

	if helloMsg := msg.GetHello(); helloMsg != nil {
		pullMsgType = HelloMsgType
		p.engine.OnHello(helloMsg.Nonce, m)
	}
	if digest := msg.GetDataDig(); digest != nil {
		d := p.PullAdapter.IngressDigFilter(digest)
		itemIDs = util.BytesToStrings(d.Digests)
		pullMsgType = DigestMsgType
		p.engine.OnDigest(itemIDs, d.Nonce, m)
	}
	if req := msg.GetDataReq(); req != nil {
		itemIDs = util.BytesToStrings(req.Digests)
		pullMsgType = RequestMsgType
		p.engine.OnReq(itemIDs, req.Nonce, m)
	}
	if res := msg.GetDataUpdate(); res != nil {
		itemIDs = make([]string, len(res.Data))
		items = make([]*proto.SignedGossipMessage, len(res.Data))
		pullMsgType = ResponseMsgType
		for i, pulledMsg := range res.Data {
			msg, err := pulledMsg.ToGossipMessage()
			if err != nil {
				p.logger.Warningf("Data update contains an invalid message: %+v", errors.WithStack(err))
				return
			}
			p.MsgCons(msg)
			itemIDs[i] = p.IdExtractor(msg)
			items[i] = msg
			p.Lock()
			p.itemID2Msg[itemIDs[i]] = msg
			p.Unlock()
		}
		p.engine.OnRes(itemIDs, res.Nonce)
	}

	// Invoke hooks for relevant message type
	for _, h := range p.hooksByMsgType(pullMsgType) {
		h(itemIDs, items, m)
	}
}

func (p *pullMediatorImpl) Stop() {
	p.engine.Stop()
}

// RegisterMsgHook registers a message hook to a specific type of pull message
func (p *pullMediatorImpl) RegisterMsgHook(pullMsgType MsgType, hook MessageHook) {
	p.Lock()
	defer p.Unlock()
	p.msgType2Hook[pullMsgType] = append(p.msgType2Hook[pullMsgType], hook)

}

// Add adds a GossipMessage to the store
func (p *pullMediatorImpl) Add(msg *proto.SignedGossipMessage) {
	p.Lock()
	defer p.Unlock()
	itemID := p.IdExtractor(msg)
	p.itemID2Msg[itemID] = msg
	p.engine.Add(itemID)
}

// Remove removes a GossipMessage from the Mediator with a matching digest,
// if such a message exits
func (p *pullMediatorImpl) Remove(digest string) {
	p.Lock()
	defer p.Unlock()
	delete(p.itemID2Msg, digest)
	p.engine.Remove(digest)
}

// SelectPeers returns a slice of peers which the engine will initiate the protocol with
func (p *pullMediatorImpl) SelectPeers() []string {
	remotePeers := SelectEndpoints(p.config.PeerCountToSelect, p.MemSvc.GetMembership())
	endpoints := make([]string, len(remotePeers))
	for i, peer := range remotePeers {
		endpoints[i] = peer.Endpoint
	}
	return endpoints
}

// Hello sends a hello message to initiate the protocol
// and returns an NONCE that is expected to be returned
// in the digest message.
func (p *pullMediatorImpl) Hello(dest string, nonce uint64) {
	helloMsg := &proto.GossipMessage{
		Channel: p.config.Channel,
		Tag:     p.config.Tag,
		Content: &proto.GossipMessage_Hello{
			Hello: &proto.GossipHello{
				Nonce:    nonce,
				Metadata: nil,
				MsgType:  p.config.MsgType,
			},
		},
	}

	p.logger.Debug("Sending", p.config.MsgType, "hello to", dest)
	sMsg, err := helloMsg.NoopSign()
	if err != nil {
		p.logger.Errorf("Failed creating SignedGossipMessage: %+v", errors.WithStack(err))
		return
	}
	p.Sndr.Send(sMsg, p.peersWithEndpoints(dest)...)
}

// SendDigest sends a digest to a remote PullEngine.
// The context parameter specifies the remote engine to send to.
func (p *pullMediatorImpl) SendDigest(digest []string, nonce uint64, context interface{}) {
	digMsg := &proto.GossipMessage{
		Channel: p.config.Channel,
		Tag:     p.config.Tag,
		Nonce:   0,
		Content: &proto.GossipMessage_DataDig{
			DataDig: &proto.DataDigest{
				MsgType: p.config.MsgType,
				Nonce:   nonce,
				Digests: util.StringsToBytes(digest),
			},
		},
	}
	remotePeer := context.(proto.ReceivedMessage).GetConnectionInfo()
	if p.logger.IsEnabledFor(zapcore.DebugLevel) {
		p.logger.Debug("Sending", p.config.MsgType, "digest:", digMsg.GetDataDig().FormattedDigests(), "to", remotePeer)
	}

	context.(proto.ReceivedMessage).Respond(digMsg)
}

// SendReq sends an array of items to a certain remote PullEngine identified
// by a string
func (p *pullMediatorImpl) SendReq(dest string, items []string, nonce uint64) {
	req := &proto.GossipMessage{
		Channel: p.config.Channel,
		Tag:     p.config.Tag,
		Nonce:   0,
		Content: &proto.GossipMessage_DataReq{
			DataReq: &proto.DataRequest{
				MsgType: p.config.MsgType,
				Nonce:   nonce,
				Digests: util.StringsToBytes(items),
			},
		},
	}
	if p.logger.IsEnabledFor(zapcore.DebugLevel) {
		p.logger.Debug("Sending", req.GetDataReq().FormattedDigests(), "to", dest)
	}
	sMsg, err := req.NoopSign()
	if err != nil {
		p.logger.Warningf("Failed creating SignedGossipMessage: %+v", errors.WithStack(err))
		return
	}
	p.Sndr.Send(sMsg, p.peersWithEndpoints(dest)...)
}

// SendRes sends an array of items to a remote PullEngine identified by a context.
func (p *pullMediatorImpl) SendRes(items []string, context interface{}, nonce uint64) {
	items2return := []*proto.Envelope{}
	p.RLock()
	defer p.RUnlock()
	for _, item := range items {
		if msg, exists := p.itemID2Msg[item]; exists {
			items2return = append(items2return, msg.Envelope)
		}
	}
	returnedUpdate := &proto.GossipMessage{
		Channel: p.config.Channel,
		Tag:     p.config.Tag,
		Nonce:   0,
		Content: &proto.GossipMessage_DataUpdate{
			DataUpdate: &proto.DataUpdate{
				MsgType: p.config.MsgType,
				Nonce:   nonce,
				Data:    items2return,
			},
		},
	}
	remotePeer := context.(proto.ReceivedMessage).GetConnectionInfo()
	p.logger.Debug("Sending", len(returnedUpdate.GetDataUpdate().Data), p.config.MsgType, "items to", remotePeer)
	context.(proto.ReceivedMessage).Respond(returnedUpdate)
}

func (p *pullMediatorImpl) peersWithEndpoints(endpoints ...string) []*comm.RemotePeer {
	peers := []*comm.RemotePeer{}
	for _, member := range p.MemSvc.GetMembership() {
		for _, endpoint := range endpoints {
			if member.PreferredEndpoint() == endpoint {
				peers = append(peers, &comm.RemotePeer{Endpoint: member.PreferredEndpoint(), PKIID: member.PKIid})
			}
		}
	}
	return peers
}

func (p *pullMediatorImpl) hooksByMsgType(msgType MsgType) []MessageHook {
	p.RLock()
	defer p.RUnlock()
	returnedHooks := []MessageHook{}
	for _, h := range p.msgType2Hook[msgType] {
		returnedHooks = append(returnedHooks, h)
	}
	return returnedHooks
}

// SelectEndpoints select k peers from peerPool and returns them.
func SelectEndpoints(k int, peerPool []discovery.NetworkMember) []*comm.RemotePeer {
	if len(peerPool) < k {
		k = len(peerPool)
	}

	indices := util.GetRandomIndices(k, len(peerPool)-1)
	endpoints := make([]*comm.RemotePeer, len(indices))
	for i, j := range indices {
		endpoints[i] = &comm.RemotePeer{Endpoint: peerPool[j].PreferredEndpoint(), PKIID: peerPool[j].PKIid}
	}
	return endpoints
}
<commit_msg>[FAB-13204] refactor gossip pull message class<commit_after>/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package pull

import (
	"sync"
	"time"

	"github.com/hyperledger/fabric/gossip/comm"
	"github.com/hyperledger/fabric/gossip/common"
	"github.com/hyperledger/fabric/gossip/discovery"
	"github.com/hyperledger/fabric/gossip/gossip/algo"
	"github.com/hyperledger/fabric/gossip/util"
	proto "github.com/hyperledger/fabric/protos/gossip"
	"github.com/pkg/errors"
	"go.uber.org/zap/zapcore"
)

// Constants go here.
const (
	HelloMsgType MsgType = iota
	DigestMsgType
	RequestMsgType
	ResponseMsgType
)

// MsgType defines the type of a message that is sent to the PullStore
type MsgType int

// MessageHook defines a function that will run after a certain pull message is received
type MessageHook func(itemIDs []string, items []*proto.SignedGossipMessage, msg proto.ReceivedMessage)

// Sender sends messages to remote peers
type Sender interface {
	// Send sends a message to a list of remote peers
	Send(msg *proto.SignedGossipMessage, peers ...*comm.RemotePeer)
}

// MembershipService obtains membership information of alive peers
type MembershipService interface {
	// GetMembership returns the membership of
	GetMembership() []discovery.NetworkMember
}

// Config defines the configuration of the pull mediator
type Config struct {
	ID                string
	PullInterval      time.Duration // Duration between pull invocations
	Channel           common.ChainID
	PeerCountToSelect int // Number of peers to initiate pull with
	Tag               proto.GossipMessage_Tag
	MsgType           proto.PullMsgType
}

// IngressDigestFilter filters out entities in digests that are received from remote peers
type IngressDigestFilter func(digestMsg *proto.DataDigest) *proto.DataDigest

// EgressDigestFilter filters digests to be sent to a remote peer, that
// sent a hello with the following message
type EgressDigestFilter func(helloMsg proto.ReceivedMessage) func(digestItem string) bool

// byContext converts this EgressDigFilter to an algo.DigestFilter
func (df EgressDigestFilter) byContext() algo.DigestFilter {
	return func(context interface{}) func(digestItem string) bool {
		return func(digestItem string) bool {
			return df(context.(proto.ReceivedMessage))(digestItem)
		}
	}
}

// PullAdapter defines methods of the pullStore to interact
// with various modules of gossip
type PullAdapter struct {
	Sndr             Sender
	MemSvc           MembershipService
	IdExtractor      proto.IdentifierExtractor
	MsgCons          proto.MsgConsumer
	EgressDigFilter  EgressDigestFilter
	IngressDigFilter IngressDigestFilter
}

// Mediator is a component wrap a PullEngine and provides the methods
// it needs to perform pull synchronization.
// The specialization of a pull mediator to a certain type of message is
// done by the configuration, a IdentifierExtractor, IdentifierExtractor
// given at construction, and also hooks that can be registered for each
// type of pullMsgType (hello, digest, req, res).
type Mediator interface {
	// Stop stop the Mediator
	Stop()

	// RegisterMsgHook registers a message hook to a specific type of pull message
	RegisterMsgHook(MsgType, MessageHook)

	// Add adds a GossipMessage to the Mediator
	Add(*proto.SignedGossipMessage)

	// Remove removes a GossipMessage from the Mediator with a matching digest,
	// if such a message exits
	Remove(digest string)

	// HandleMessage handles a message from some remote peer
	HandleMessage(msg proto.ReceivedMessage)
}

// pullMediatorImpl is an implementation of Mediator
type pullMediatorImpl struct {
	sync.RWMutex
	*PullAdapter
	msgType2Hook map[MsgType][]MessageHook
	config       Config
	logger       util.Logger
	itemID2Msg   map[string]*proto.SignedGossipMessage
	engine       *algo.PullEngine
}

// NewPullMediator returns a new Mediator
func NewPullMediator(config Config, adapter *PullAdapter) Mediator {
	egressDigFilter := adapter.EgressDigFilter

	acceptAllFilter := func(_ proto.ReceivedMessage) func(string) bool {
		return func(_ string) bool {
			return true
		}
	}

	if egressDigFilter == nil {
		egressDigFilter = acceptAllFilter
	}

	p := &pullMediatorImpl{
		PullAdapter:  adapter,
		msgType2Hook: make(map[MsgType][]MessageHook),
		config:       config,
		logger:       util.GetLogger(util.PullLogger, config.ID),
		itemID2Msg:   make(map[string]*proto.SignedGossipMessage),
	}

	p.engine = algo.NewPullEngineWithFilter(p, config.PullInterval, egressDigFilter.byContext())

	if adapter.IngressDigFilter == nil {
		// Create accept all filter
		adapter.IngressDigFilter = func(digestMsg *proto.DataDigest) *proto.DataDigest {
			return digestMsg
		}
	}
	return p

}

func (p *pullMediatorImpl) HandleMessage(m proto.ReceivedMessage) {
	if m.GetGossipMessage() == nil || !m.GetGossipMessage().IsPullMsg() {
		return
	}
	msg := m.GetGossipMessage()
	msgType := msg.GetPullMsgType()
	if msgType != p.config.MsgType {
		return
	}

	p.logger.Debug(msg)

	itemIDs := []string{}
	items := []*proto.SignedGossipMessage{}
	var pullMsgType MsgType

	if helloMsg := msg.GetHello(); helloMsg != nil {
		pullMsgType = HelloMsgType
		p.engine.OnHello(helloMsg.Nonce, m)
	} else if digest := msg.GetDataDig(); digest != nil {
		d := p.PullAdapter.IngressDigFilter(digest)
		itemIDs = util.BytesToStrings(d.Digests)
		pullMsgType = DigestMsgType
		p.engine.OnDigest(itemIDs, d.Nonce, m)
	} else if req := msg.GetDataReq(); req != nil {
		itemIDs = util.BytesToStrings(req.Digests)
		pullMsgType = RequestMsgType
		p.engine.OnReq(itemIDs, req.Nonce, m)
	} else if res := msg.GetDataUpdate(); res != nil {
		itemIDs = make([]string, len(res.Data))
		items = make([]*proto.SignedGossipMessage, len(res.Data))
		pullMsgType = ResponseMsgType
		for i, pulledMsg := range res.Data {
			msg, err := pulledMsg.ToGossipMessage()
			if err != nil {
				p.logger.Warningf("Data update contains an invalid message: %+v", errors.WithStack(err))
				return
			}
			p.MsgCons(msg)
			itemIDs[i] = p.IdExtractor(msg)
			items[i] = msg
			p.Lock()
			p.itemID2Msg[itemIDs[i]] = msg
			p.Unlock()
		}
		p.engine.OnRes(itemIDs, res.Nonce)
	}

	// Invoke hooks for relevant message type
	for _, h := range p.hooksByMsgType(pullMsgType) {
		h(itemIDs, items, m)
	}
}

func (p *pullMediatorImpl) Stop() {
	p.engine.Stop()
}

// RegisterMsgHook registers a message hook to a specific type of pull message
func (p *pullMediatorImpl) RegisterMsgHook(pullMsgType MsgType, hook MessageHook) {
	p.Lock()
	defer p.Unlock()
	p.msgType2Hook[pullMsgType] = append(p.msgType2Hook[pullMsgType], hook)

}

// Add adds a GossipMessage to the store
func (p *pullMediatorImpl) Add(msg *proto.SignedGossipMessage) {
	p.Lock()
	defer p.Unlock()
	itemID := p.IdExtractor(msg)
	p.itemID2Msg[itemID] = msg
	p.engine.Add(itemID)
}

// Remove removes a GossipMessage from the Mediator with a matching digest,
// if such a message exits
func (p *pullMediatorImpl) Remove(digest string) {
	p.Lock()
	defer p.Unlock()
	delete(p.itemID2Msg, digest)
	p.engine.Remove(digest)
}

// SelectPeers returns a slice of peers which the engine will initiate the protocol with
func (p *pullMediatorImpl) SelectPeers() []string {
	remotePeers := SelectEndpoints(p.config.PeerCountToSelect, p.MemSvc.GetMembership())
	endpoints := make([]string, len(remotePeers))
	for i, peer := range remotePeers {
		endpoints[i] = peer.Endpoint
	}
	return endpoints
}

// Hello sends a hello message to initiate the protocol
// and returns an NONCE that is expected to be returned
// in the digest message.
func (p *pullMediatorImpl) Hello(dest string, nonce uint64) {
	helloMsg := &proto.GossipMessage{
		Channel: p.config.Channel,
		Tag:     p.config.Tag,
		Content: &proto.GossipMessage_Hello{
			Hello: &proto.GossipHello{
				Nonce:    nonce,
				Metadata: nil,
				MsgType:  p.config.MsgType,
			},
		},
	}

	p.logger.Debug("Sending", p.config.MsgType, "hello to", dest)
	sMsg, err := helloMsg.NoopSign()
	if err != nil {
		p.logger.Errorf("Failed creating SignedGossipMessage: %+v", errors.WithStack(err))
		return
	}
	p.Sndr.Send(sMsg, p.peersWithEndpoints(dest)...)
}

// SendDigest sends a digest to a remote PullEngine.
// The context parameter specifies the remote engine to send to.
func (p *pullMediatorImpl) SendDigest(digest []string, nonce uint64, context interface{}) {
	digMsg := &proto.GossipMessage{
		Channel: p.config.Channel,
		Tag:     p.config.Tag,
		Nonce:   0,
		Content: &proto.GossipMessage_DataDig{
			DataDig: &proto.DataDigest{
				MsgType: p.config.MsgType,
				Nonce:   nonce,
				Digests: util.StringsToBytes(digest),
			},
		},
	}
	remotePeer := context.(proto.ReceivedMessage).GetConnectionInfo()
	if p.logger.IsEnabledFor(zapcore.DebugLevel) {
		p.logger.Debug("Sending", p.config.MsgType, "digest:", digMsg.GetDataDig().FormattedDigests(), "to", remotePeer)
	}

	context.(proto.ReceivedMessage).Respond(digMsg)
}

// SendReq sends an array of items to a certain remote PullEngine identified
// by a string
func (p *pullMediatorImpl) SendReq(dest string, items []string, nonce uint64) {
	req := &proto.GossipMessage{
		Channel: p.config.Channel,
		Tag:     p.config.Tag,
		Nonce:   0,
		Content: &proto.GossipMessage_DataReq{
			DataReq: &proto.DataRequest{
				MsgType: p.config.MsgType,
				Nonce:   nonce,
				Digests: util.StringsToBytes(items),
			},
		},
	}
	if p.logger.IsEnabledFor(zapcore.DebugLevel) {
		p.logger.Debug("Sending", req.GetDataReq().FormattedDigests(), "to", dest)
	}
	sMsg, err := req.NoopSign()
	if err != nil {
		p.logger.Warningf("Failed creating SignedGossipMessage: %+v", errors.WithStack(err))
		return
	}
	p.Sndr.Send(sMsg, p.peersWithEndpoints(dest)...)
}

// SendRes sends an array of items to a remote PullEngine identified by a context.
func (p *pullMediatorImpl) SendRes(items []string, context interface{}, nonce uint64) {
	items2return := []*proto.Envelope{}
	p.RLock()
	defer p.RUnlock()
	for _, item := range items {
		if msg, exists := p.itemID2Msg[item]; exists {
			items2return = append(items2return, msg.Envelope)
		}
	}
	returnedUpdate := &proto.GossipMessage{
		Channel: p.config.Channel,
		Tag:     p.config.Tag,
		Nonce:   0,
		Content: &proto.GossipMessage_DataUpdate{
			DataUpdate: &proto.DataUpdate{
				MsgType: p.config.MsgType,
				Nonce:   nonce,
				Data:    items2return,
			},
		},
	}
	remotePeer := context.(proto.ReceivedMessage).GetConnectionInfo()
	p.logger.Debug("Sending", len(returnedUpdate.GetDataUpdate().Data), p.config.MsgType, "items to", remotePeer)
	context.(proto.ReceivedMessage).Respond(returnedUpdate)
}

func (p *pullMediatorImpl) peersWithEndpoints(endpoints ...string) []*comm.RemotePeer {
	peers := []*comm.RemotePeer{}
	for _, member := range p.MemSvc.GetMembership() {
		for _, endpoint := range endpoints {
			if member.PreferredEndpoint() == endpoint {
				peers = append(peers, &comm.RemotePeer{Endpoint: member.PreferredEndpoint(), PKIID: member.PKIid})
			}
		}
	}
	return peers
}

func (p *pullMediatorImpl) hooksByMsgType(msgType MsgType) []MessageHook {
	p.RLock()
	defer p.RUnlock()
	returnedHooks := []MessageHook{}
	for _, h := range p.msgType2Hook[msgType] {
		returnedHooks = append(returnedHooks, h)
	}
	return returnedHooks
}

// SelectEndpoints select k peers from peerPool and returns them.
func SelectEndpoints(k int, peerPool []discovery.NetworkMember) []*comm.RemotePeer {
	if len(peerPool) < k {
		k = len(peerPool)
	}

	indices := util.GetRandomIndices(k, len(peerPool)-1)
	endpoints := make([]*comm.RemotePeer, len(indices))
	for i, j := range indices {
		endpoints[i] = &comm.RemotePeer{Endpoint: peerPool[j].PreferredEndpoint(), PKIID: peerPool[j].PKIid}
	}
	return endpoints
}
<|endoftext|>


<commit_before>// Copyright 2011, 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package ec2

import (
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs/imagemetadata"
	"launchpad.net/juju-core/environs/instances"
	"launchpad.net/juju-core/environs/simplestreams"
	"launchpad.net/juju-core/testing/testbase"
)

type imageSuite struct {
	testbase.LoggingSuite
}

var _ = gc.Suite(&imageSuite{})

func (s *imageSuite) SetUpSuite(c *gc.C) {
	s.LoggingSuite.SetUpSuite(c)
	UseTestImageData(TestImagesData)
}

func (s *imageSuite) TearDownSuite(c *gc.C) {
	UseTestImageData(nil)
	s.LoggingSuite.TearDownTest(c)
}

type specSuite struct {
	testbase.LoggingSuite
}

var _ = gc.Suite(&specSuite{})

func (s *specSuite) SetUpSuite(c *gc.C) {
	s.LoggingSuite.SetUpSuite(c)
	UseTestImageData(TestImagesData)
	UseTestInstanceTypeData(TestInstanceTypeCosts)
	UseTestRegionData(TestRegions)
}

func (s *specSuite) TearDownSuite(c *gc.C) {
	UseTestInstanceTypeData(nil)
	UseTestImageData(nil)
	UseTestRegionData(nil)
	s.LoggingSuite.TearDownSuite(c)
}

var findInstanceSpecTests = []struct {
	series string
	arches []string
	cons   string
	itype  string
	image  string
}{
	{
		series: "precise",
		arches: both,
		itype:  "m1.small",
		image:  "ami-00000033",
	}, {
		series: "quantal",
		arches: both,
		itype:  "m1.small",
		image:  "ami-01000034",
	}, {
		series: "precise",
		arches: both,
		cons:   "cpu-cores=4",
		itype:  "m1.xlarge",
		image:  "ami-00000033",
	}, {
		series: "precise",
		arches: both,
		cons:   "cpu-cores=2 arch=i386",
		itype:  "c1.medium",
		image:  "ami-00000034",
	}, {
		series: "precise",
		arches: both,
		cons:   "mem=10G",
		itype:  "m1.xlarge",
		image:  "ami-00000033",
	}, {
		series: "precise",
		arches: both,
		cons:   "mem=",
		itype:  "m1.small",
		image:  "ami-00000033",
	}, {
		series: "precise",
		arches: both,
		cons:   "cpu-power=",
		itype:  "t1.micro",
		image:  "ami-00000033",
	}, {
		series: "precise",
		arches: both,
		cons:   "cpu-power=800",
		itype:  "m1.xlarge",
		image:  "ami-00000033",
	}, {
		series: "precise",
		arches: both,
		cons:   "cpu-power=500 arch=i386",
		itype:  "c1.medium",
		image:  "ami-00000034",
	}, {
		series: "precise",
		arches: []string{"i386"},
		cons:   "cpu-power=400",
		itype:  "c1.medium",
		image:  "ami-00000034",
	}, {
		series: "quantal",
		arches: both,
		cons:   "arch=amd64",
		itype:  "cc1.4xlarge",
		image:  "ami-01000035",
	},
}

func (s *specSuite) TestFindInstanceSpec(c *gc.C) {
	for i, t := range findInstanceSpecTests {
		c.Logf("test %d", i)
		stor := ebsStorage
		spec, err := findInstanceSpec(
			[]simplestreams.DataSource{simplestreams.NewURLDataSource("test:", simplestreams.VerifySSLHostnames), simplestreams.VerifySSLHostnames},
			&instances.InstanceConstraint{
				Region:      "test",
				Series:      t.series,
				Arches:      t.arches,
				Constraints: constraints.MustParse(t.cons),
				Storage:     &stor,
			})
		c.Assert(err, gc.IsNil)
		c.Check(spec.InstanceType.Name, gc.Equals, t.itype)
		c.Check(spec.Image.Id, gc.Equals, t.image)
	}
}

var findInstanceSpecErrorTests = []struct {
	series string
	arches []string
	cons   string
	err    string
}{
	{
		series: "bad",
		arches: both,
		err:    `invalid series "bad"`,
	}, {
		series: "precise",
		arches: []string{"arm"},
		err:    `no "precise" images in test with arches \[arm\]`,
	}, {
		series: "raring",
		arches: both,
		cons:   "mem=4G",
		err:    `no "raring" images in test matching instance types \[m1.large m1.xlarge c1.xlarge cc1.4xlarge cc2.8xlarge\]`,
	},
}

func (s *specSuite) TestFindInstanceSpecErrors(c *gc.C) {
	for i, t := range findInstanceSpecErrorTests {
		c.Logf("test %d", i)
		_, err := findInstanceSpec(
			[]simplestreams.DataSource{simplestreams.NewURLDataSource("test:")}, &instances.InstanceConstraint{
				Region:      "test",
				Series:      t.series,
				Arches:      t.arches,
				Constraints: constraints.MustParse(t.cons),
			})
		c.Check(err, gc.ErrorMatches, t.err)
	}
}

func (*specSuite) TestFilterImagesAcceptsNil(c *gc.C) {
	c.Check(filterImages(nil), gc.HasLen, 0)
}

func (*specSuite) TestFilterImagesAcceptsImageWithEBSStorage(c *gc.C) {
	input := []*imagemetadata.ImageMetadata{{Id: "yay", Storage: "ebs"}}
	c.Check(filterImages(input), gc.DeepEquals, input)
}

func (*specSuite) TestFilterImagesRejectsImageWithoutEBSStorage(c *gc.C) {
	input := []*imagemetadata.ImageMetadata{{Id: "boo", Storage: "ftp"}}
	c.Check(filterImages(input), gc.HasLen, 0)
}

func (*specSuite) TestFilterImagesReturnsSelectively(c *gc.C) {
	good := imagemetadata.ImageMetadata{Id: "good", Storage: "ebs"}
	bad := imagemetadata.ImageMetadata{Id: "bad", Storage: "ftp"}
	input := []*imagemetadata.ImageMetadata{&good, &bad}
	expectation := []*imagemetadata.ImageMetadata{&good}
	c.Check(filterImages(input), gc.DeepEquals, expectation)
}

func (*specSuite) TestFilterImagesMaintainsOrdering(c *gc.C) {
	input := []*imagemetadata.ImageMetadata{
		{Id: "one", Storage: "ebs"},
		{Id: "two", Storage: "ebs"},
		{Id: "three", Storage: "ebs"},
	}
	c.Check(filterImages(input), gc.DeepEquals, input)
}
<commit_msg>Fix some typos<commit_after>// Copyright 2011, 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package ec2

import (
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs/imagemetadata"
	"launchpad.net/juju-core/environs/instances"
	"launchpad.net/juju-core/environs/simplestreams"
	"launchpad.net/juju-core/testing/testbase"
)

type imageSuite struct {
	testbase.LoggingSuite
}

var _ = gc.Suite(&imageSuite{})

func (s *imageSuite) SetUpSuite(c *gc.C) {
	s.LoggingSuite.SetUpSuite(c)
	UseTestImageData(TestImagesData)
}

func (s *imageSuite) TearDownSuite(c *gc.C) {
	UseTestImageData(nil)
	s.LoggingSuite.TearDownTest(c)
}

type specSuite struct {
	testbase.LoggingSuite
}

var _ = gc.Suite(&specSuite{})

func (s *specSuite) SetUpSuite(c *gc.C) {
	s.LoggingSuite.SetUpSuite(c)
	UseTestImageData(TestImagesData)
	UseTestInstanceTypeData(TestInstanceTypeCosts)
	UseTestRegionData(TestRegions)
}

func (s *specSuite) TearDownSuite(c *gc.C) {
	UseTestInstanceTypeData(nil)
	UseTestImageData(nil)
	UseTestRegionData(nil)
	s.LoggingSuite.TearDownSuite(c)
}

var findInstanceSpecTests = []struct {
	series string
	arches []string
	cons   string
	itype  string
	image  string
}{
	{
		series: "precise",
		arches: both,
		itype:  "m1.small",
		image:  "ami-00000033",
	}, {
		series: "quantal",
		arches: both,
		itype:  "m1.small",
		image:  "ami-01000034",
	}, {
		series: "precise",
		arches: both,
		cons:   "cpu-cores=4",
		itype:  "m1.xlarge",
		image:  "ami-00000033",
	}, {
		series: "precise",
		arches: both,
		cons:   "cpu-cores=2 arch=i386",
		itype:  "c1.medium",
		image:  "ami-00000034",
	}, {
		series: "precise",
		arches: both,
		cons:   "mem=10G",
		itype:  "m1.xlarge",
		image:  "ami-00000033",
	}, {
		series: "precise",
		arches: both,
		cons:   "mem=",
		itype:  "m1.small",
		image:  "ami-00000033",
	}, {
		series: "precise",
		arches: both,
		cons:   "cpu-power=",
		itype:  "t1.micro",
		image:  "ami-00000033",
	}, {
		series: "precise",
		arches: both,
		cons:   "cpu-power=800",
		itype:  "m1.xlarge",
		image:  "ami-00000033",
	}, {
		series: "precise",
		arches: both,
		cons:   "cpu-power=500 arch=i386",
		itype:  "c1.medium",
		image:  "ami-00000034",
	}, {
		series: "precise",
		arches: []string{"i386"},
		cons:   "cpu-power=400",
		itype:  "c1.medium",
		image:  "ami-00000034",
	}, {
		series: "quantal",
		arches: both,
		cons:   "arch=amd64",
		itype:  "cc1.4xlarge",
		image:  "ami-01000035",
	},
}

func (s *specSuite) TestFindInstanceSpec(c *gc.C) {
	for i, t := range findInstanceSpecTests {
		c.Logf("test %d", i)
		stor := ebsStorage
		spec, err := findInstanceSpec(
			[]simplestreams.DataSource{simplestreams.NewURLDataSource("test:", simplestreams.VerifySSLHostnames)},
			&instances.InstanceConstraint{
				Region:      "test",
				Series:      t.series,
				Arches:      t.arches,
				Constraints: constraints.MustParse(t.cons),
				Storage:     &stor,
			})
		c.Assert(err, gc.IsNil)
		c.Check(spec.InstanceType.Name, gc.Equals, t.itype)
		c.Check(spec.Image.Id, gc.Equals, t.image)
	}
}

var findInstanceSpecErrorTests = []struct {
	series string
	arches []string
	cons   string
	err    string
}{
	{
		series: "bad",
		arches: both,
		err:    `invalid series "bad"`,
	}, {
		series: "precise",
		arches: []string{"arm"},
		err:    `no "precise" images in test with arches \[arm\]`,
	}, {
		series: "raring",
		arches: both,
		cons:   "mem=4G",
		err:    `no "raring" images in test matching instance types \[m1.large m1.xlarge c1.xlarge cc1.4xlarge cc2.8xlarge\]`,
	},
}

func (s *specSuite) TestFindInstanceSpecErrors(c *gc.C) {
	for i, t := range findInstanceSpecErrorTests {
		c.Logf("test %d", i)
		_, err := findInstanceSpec(
			[]simplestreams.DataSource{simplestreams.NewURLDataSource("test:", simplestreams.VerifySSLHostnames)}, &instances.InstanceConstraint{
				Region:      "test",
				Series:      t.series,
				Arches:      t.arches,
				Constraints: constraints.MustParse(t.cons),
			})
		c.Check(err, gc.ErrorMatches, t.err)
	}
}

func (*specSuite) TestFilterImagesAcceptsNil(c *gc.C) {
	c.Check(filterImages(nil), gc.HasLen, 0)
}

func (*specSuite) TestFilterImagesAcceptsImageWithEBSStorage(c *gc.C) {
	input := []*imagemetadata.ImageMetadata{{Id: "yay", Storage: "ebs"}}
	c.Check(filterImages(input), gc.DeepEquals, input)
}

func (*specSuite) TestFilterImagesRejectsImageWithoutEBSStorage(c *gc.C) {
	input := []*imagemetadata.ImageMetadata{{Id: "boo", Storage: "ftp"}}
	c.Check(filterImages(input), gc.HasLen, 0)
}

func (*specSuite) TestFilterImagesReturnsSelectively(c *gc.C) {
	good := imagemetadata.ImageMetadata{Id: "good", Storage: "ebs"}
	bad := imagemetadata.ImageMetadata{Id: "bad", Storage: "ftp"}
	input := []*imagemetadata.ImageMetadata{&good, &bad}
	expectation := []*imagemetadata.ImageMetadata{&good}
	c.Check(filterImages(input), gc.DeepEquals, expectation)
}

func (*specSuite) TestFilterImagesMaintainsOrdering(c *gc.C) {
	input := []*imagemetadata.ImageMetadata{
		{Id: "one", Storage: "ebs"},
		{Id: "two", Storage: "ebs"},
		{Id: "three", Storage: "ebs"},
	}
	c.Check(filterImages(input), gc.DeepEquals, input)
}
<|endoftext|>


<commit_before>package network

import (
	"errors"
	"net"
	"reflect"
	"time"

	"github.com/NebulousLabs/Sia/encoding"
)

// reflection helpers
var (
	typeOfError = reflect.TypeOf((*error)(nil)).Elem()
	typeOfConn  = reflect.TypeOf((*net.Conn)(nil)).Elem()
)

// handlerName truncates a string to 8 bytes. If len(name) < 8, the remaining
// bytes are 0. A handlerName is specified at the beginning of each network
// call, indicating which function should handle the connection.
func handlerName(name string) []byte {
	b := make([]byte, 8)
	copy(b, name)
	return b
}

// Call establishes a TCP connection to the Address, calls the provided
// function on it, and closes the connection.
func (na Address) Call(name string, fn func(net.Conn) error) error {
	conn, err := net.DialTimeout("tcp", string(na), timeout)
	if err != nil {
		return err
	}
	defer conn.Close()
	// set default deadline
	// note: fn can extend this deadline as needed
	conn.SetDeadline(time.Now().Add(timeout))
	// write header
	if _, err := conn.Write(handlerName(name)); err != nil {
		return err
	}
	return fn(conn)
}

// RPC performs a Remote Procedure Call by sending the procedure name and
// encoded argument, and decoding the response into the supplied object.
// 'resp' must be a pointer. If arg is nil, no object is sent. If 'resp' is
// nil, no response is read.
func (na *Address) RPC(name string, arg, resp interface{}) error {
	return na.Call(name, func(conn net.Conn) error {
		// write arg
		if arg != nil {
			if _, err := encoding.WriteObject(conn, arg); err != nil {
				return err
			}
		}
		// read resp
		if resp != nil {
			if err := encoding.ReadObject(conn, resp, maxMsgLen); err != nil {
				return err
			}
		}
		// read err
		var errStr string
		if err := encoding.ReadObject(conn, &errStr, maxMsgLen); err != nil {
			return err
		} else if errStr != "" {
			return errors.New(errStr)
		}
		return nil
	})
}

// Broadcast calls the RPC on each peer in the address book.
func (tcps *TCPServer) Broadcast(name string, arg, resp interface{}) {
	for _, addr := range tcps.AddressBook() {
		err := addr.RPC(name, arg, resp)
		// remove unresponsive peers
		if err != nil {
			tcps.RemovePeer(addr)
		}
	}
}

// RegisterRPC registers a function as an RPC handler for a given identifier.
// The function must be one of four possible types:
//     func(net.Conn) error
//     func(Type) (Type, error)
//     func(Type) error
//     func() (Type, error)
// To call an RPC, use Address.RPC, supplying the same identifier given to
// RegisterRPC. Identifiers should always use PascalCase.
func (tcps *TCPServer) RegisterRPC(name string, fn interface{}) error {
	// all handlers are functions with 0 or 1 ins and 1 or 2 outs, the last of
	// which must be an error.
	val, typ := reflect.ValueOf(fn), reflect.TypeOf(fn)
	if typ.Kind() != reflect.Func || typ.NumIn() > 1 || typ.NumOut() > 2 ||
		typ.NumOut() < 1 || typ.Out(typ.NumOut()-1) != typeOfError {
		panic("registered function has wrong type signature")
	}

	var handler func(net.Conn) error
	switch {
	// func(net.Conn) error
	case typ.NumIn() == 1 && typ.NumOut() == 1 && typ.In(0) == typeOfConn:
		handler = fn.(func(net.Conn) error)
	// func(Type) (Type, error)
	case typ.NumIn() == 1 && typ.NumOut() == 2:
		handler = registerRPC(val, typ)
	// func(Type) error
	case typ.NumIn() == 1 && typ.NumOut() == 1:
		handler = registerArg(val, typ)
	// func() (Type, error)
	case typ.NumIn() == 0 && typ.NumOut() == 2:
		handler = registerResp(val, typ)

	default:
		panic("registered function has wrong type signature")
	}

	ident := string(handlerName(name))
	tcps.Lock()
	tcps.handlerMap[ident] = handler
	tcps.Unlock()

	return nil
}

// registerRPC is for handlers that take an argument return a value. The input
// is decoded and passed to fn, whose return value is written back to the
// caller. fn must have the type signature:
//   func(Type, *Type) error
func registerRPC(fn reflect.Value, typ reflect.Type) func(net.Conn) error {
	return func(conn net.Conn) error {
		// read arg
		arg := reflect.New(typ.In(0))
		if err := encoding.ReadObject(conn, arg.Interface(), maxMsgLen); err != nil {
			return err
		}
		// call fn
		retvals := fn.Call([]reflect.Value{arg.Elem()})
		resp, errInter := retvals[0].Interface(), retvals[1].Interface()
		// write resp
		if _, err := encoding.WriteObject(conn, resp); err != nil {
			return err
		}
		// write err
		var errStr string
		if errInter != nil {
			errStr = errInter.(error).Error()
		}
		_, err := encoding.WriteObject(conn, errStr)
		return err
	}
}

// registerArg is for RPCs that do not return a value.
func registerArg(fn reflect.Value, typ reflect.Type) func(net.Conn) error {
	return func(conn net.Conn) error {
		// read arg
		arg := reflect.New(typ.In(0))
		if err := encoding.ReadObject(conn, arg.Interface(), maxMsgLen); err != nil {
			return err
		}
		// call fn on object
		errInter := fn.Call([]reflect.Value{arg.Elem()})[0].Interface()
		// write err
		var errStr string
		if errInter != nil {
			errStr = errInter.(error).Error()
		}
		_, err := encoding.WriteObject(conn, errStr)
		return err
	}
}

// registerResp is for RPCs that do not take a value.
func registerResp(fn reflect.Value, typ reflect.Type) func(net.Conn) error {
	return func(conn net.Conn) error {
		// call fn
		retvals := fn.Call(nil)
		resp, errInter := retvals[0].Interface(), retvals[1].Interface()
		// write resp
		if _, err := encoding.WriteObject(conn, resp); err != nil {
			return err
		}
		// write err
		var errStr string
		if errInter != nil {
			errStr = errInter.(error).Error()
		}
		_, err := encoding.WriteObject(conn, errStr)
		return err
	}
}
<commit_msg>HOTFIX: prevent bootstrap panic<commit_after>package network

import (
	"errors"
	"net"
	"reflect"
	"time"

	"github.com/NebulousLabs/Sia/encoding"
)

// reflection helpers
var (
	typeOfError = reflect.TypeOf((*error)(nil)).Elem()
	typeOfConn  = reflect.TypeOf((*net.Conn)(nil)).Elem()
)

// handlerName truncates a string to 8 bytes. If len(name) < 8, the remaining
// bytes are 0. A handlerName is specified at the beginning of each network
// call, indicating which function should handle the connection.
func handlerName(name string) []byte {
	b := make([]byte, 8)
	copy(b, name)
	return b
}

// Call establishes a TCP connection to the Address, calls the provided
// function on it, and closes the connection.
func (na Address) Call(name string, fn func(net.Conn) error) error {
	conn, err := net.DialTimeout("tcp", string(na), timeout)
	if err != nil {
		return err
	}
	defer conn.Close()
	// set default deadline
	// note: fn can extend this deadline as needed
	conn.SetDeadline(time.Now().Add(timeout))
	// write header
	if _, err := conn.Write(handlerName(name)); err != nil {
		return err
	}
	return fn(conn)
}

// RPC performs a Remote Procedure Call by sending the procedure name and
// encoded argument, and decoding the response into the supplied object.
// 'resp' must be a pointer. If arg is nil, no object is sent. If 'resp' is
// nil, no response is read.
func (na *Address) RPC(name string, arg, resp interface{}) error {
	return na.Call(name, func(conn net.Conn) error {
		// write arg
		if arg != nil {
			if _, err := encoding.WriteObject(conn, arg); err != nil {
				return err
			}
		}
		// read resp
		if resp != nil {
			if err := encoding.ReadObject(conn, resp, maxMsgLen); err != nil {
				return err
			}
		}
		// read err
		var errStr string
		if err := encoding.ReadObject(conn, &errStr, maxMsgLen); err != nil {
			return err
		} else if errStr != "" {
			return errors.New(errStr)
		}
		return nil
	})
}

// Broadcast calls the RPC on each peer in the address book.
func (tcps *TCPServer) Broadcast(name string, arg, resp interface{}) {
	for _, addr := range tcps.AddressBook() {
		// TODO: remove unresponsive peers
		_ = addr.RPC(name, arg, resp)
	}
}

// RegisterRPC registers a function as an RPC handler for a given identifier.
// The function must be one of four possible types:
//     func(net.Conn) error
//     func(Type) (Type, error)
//     func(Type) error
//     func() (Type, error)
// To call an RPC, use Address.RPC, supplying the same identifier given to
// RegisterRPC. Identifiers should always use PascalCase.
func (tcps *TCPServer) RegisterRPC(name string, fn interface{}) error {
	// all handlers are functions with 0 or 1 ins and 1 or 2 outs, the last of
	// which must be an error.
	val, typ := reflect.ValueOf(fn), reflect.TypeOf(fn)
	if typ.Kind() != reflect.Func || typ.NumIn() > 1 || typ.NumOut() > 2 ||
		typ.NumOut() < 1 || typ.Out(typ.NumOut()-1) != typeOfError {
		panic("registered function has wrong type signature")
	}

	var handler func(net.Conn) error
	switch {
	// func(net.Conn) error
	case typ.NumIn() == 1 && typ.NumOut() == 1 && typ.In(0) == typeOfConn:
		handler = fn.(func(net.Conn) error)
	// func(Type) (Type, error)
	case typ.NumIn() == 1 && typ.NumOut() == 2:
		handler = registerRPC(val, typ)
	// func(Type) error
	case typ.NumIn() == 1 && typ.NumOut() == 1:
		handler = registerArg(val, typ)
	// func() (Type, error)
	case typ.NumIn() == 0 && typ.NumOut() == 2:
		handler = registerResp(val, typ)

	default:
		panic("registered function has wrong type signature")
	}

	ident := string(handlerName(name))
	tcps.Lock()
	tcps.handlerMap[ident] = handler
	tcps.Unlock()

	return nil
}

// registerRPC is for handlers that take an argument return a value. The input
// is decoded and passed to fn, whose return value is written back to the
// caller. fn must have the type signature:
//   func(Type, *Type) error
func registerRPC(fn reflect.Value, typ reflect.Type) func(net.Conn) error {
	return func(conn net.Conn) error {
		// read arg
		arg := reflect.New(typ.In(0))
		if err := encoding.ReadObject(conn, arg.Interface(), maxMsgLen); err != nil {
			return err
		}
		// call fn
		retvals := fn.Call([]reflect.Value{arg.Elem()})
		resp, errInter := retvals[0].Interface(), retvals[1].Interface()
		// write resp
		if _, err := encoding.WriteObject(conn, resp); err != nil {
			return err
		}
		// write err
		var errStr string
		if errInter != nil {
			errStr = errInter.(error).Error()
		}
		_, err := encoding.WriteObject(conn, errStr)
		return err
	}
}

// registerArg is for RPCs that do not return a value.
func registerArg(fn reflect.Value, typ reflect.Type) func(net.Conn) error {
	return func(conn net.Conn) error {
		// read arg
		arg := reflect.New(typ.In(0))
		if err := encoding.ReadObject(conn, arg.Interface(), maxMsgLen); err != nil {
			return err
		}
		// call fn on object
		errInter := fn.Call([]reflect.Value{arg.Elem()})[0].Interface()
		// write err
		var errStr string
		if errInter != nil {
			errStr = errInter.(error).Error()
		}
		_, err := encoding.WriteObject(conn, errStr)
		return err
	}
}

// registerResp is for RPCs that do not take a value.
func registerResp(fn reflect.Value, typ reflect.Type) func(net.Conn) error {
	return func(conn net.Conn) error {
		// call fn
		retvals := fn.Call(nil)
		resp, errInter := retvals[0].Interface(), retvals[1].Interface()
		// write resp
		if _, err := encoding.WriteObject(conn, resp); err != nil {
			return err
		}
		// write err
		var errStr string
		if errInter != nil {
			errStr = errInter.(error).Error()
		}
		_, err := encoding.WriteObject(conn, errStr)
		return err
	}
}
<|endoftext|>


<commit_before>// Package gitlab implements the OAuth2 protocol for authenticating users through gitlab.
// This package can be used as a reference implementation of an OAuth2 provider for Goth.
package gitlab

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"

	"github.com/markbates/goth"
	"golang.org/x/oauth2"
	"fmt"
)

// These vars define the Authentication, Token, and Profile URLS for Gitlab. If
// using Gitlab CE or EE, you should change these values before calling New.
//
// Examples:
//	gitlab.AuthURL = "https://gitlab.acme.com/oauth/authorize
//	gitlab.TokenURL = "https://gitlab.acme.com/oauth/token
//	gitlab.ProfileURL = "https://gitlab.acme.com/api/v3/user
var (
	AuthURL    = "https://gitlab.com/oauth/authorize"
	TokenURL   = "https://gitlab.com/oauth/token"
	ProfileURL = "https://gitlab.com/api/v3/user"
)

// Provider is the implementation of `goth.Provider` for accessing Gitlab.
type Provider struct {
	ClientKey    string
	Secret       string
	CallbackURL  string
	HTTPClient   *http.Client
	config       *oauth2.Config
	providerName string
}

// New creates a new Gitlab provider and sets up important connection details.
// You should always call `gitlab.New` to get a new provider.  Never try to
// create one manually.
func New(clientKey, secret, callbackURL string, scopes ...string) *Provider {
	p := &Provider{
		ClientKey:           clientKey,
		Secret:              secret,
		CallbackURL:         callbackURL,
		providerName:        "gitlab",
	}
	p.config = newConfig(p, scopes)
	return p
}

// Name is the name used to retrieve this provider later.
func (p *Provider) Name() string {
	return p.providerName
}

// SetName is to update the name of the provider (needed in case of multiple providers of 1 type)
func (p *Provider) SetName(name string) {
	p.providerName = name
}

func (p *Provider) Client() *http.Client {
	return goth.HTTPClientWithFallBack(p.HTTPClient)
}

// Debug is a no-op for the gitlab package.
func (p *Provider) Debug(debug bool) {}

// BeginAuth asks Gitlab for an authentication end-point.
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
	return &Session{
		AuthURL: p.config.AuthCodeURL(state),
	}, nil
}

// FetchUser will go to Gitlab and access basic information about the user.
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
	sess := session.(*Session)
	user := goth.User{
		AccessToken:  sess.AccessToken,
		Provider:     p.Name(),
		RefreshToken: sess.RefreshToken,
		ExpiresAt:    sess.ExpiresAt,
	}

	if user.AccessToken == "" {
		// data is not yet retrieved since accessToken is still empty
		return user, fmt.Errorf("%s cannot get user information without accessToken", p.providerName)
	}

	response, err := p.Client().Get(ProfileURL + "?access_token=" + url.QueryEscape(sess.AccessToken))
	if err != nil {
		if response != nil {
			response.Body.Close()
		}
		return user, err
	}

	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return user, fmt.Errorf("%s responded with a %d trying to fetch user information", p.providerName, response.StatusCode)
	}

	bits, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return user, err
	}

	err = json.NewDecoder(bytes.NewReader(bits)).Decode(&user.RawData)
	if err != nil {
		return user, err
	}

	err = userFromReader(bytes.NewReader(bits), &user)

	return user, err
}

func newConfig(provider *Provider, scopes []string) *oauth2.Config {
	c := &oauth2.Config{
		ClientID:     provider.ClientKey,
		ClientSecret: provider.Secret,
		RedirectURL:  provider.CallbackURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  AuthURL,
			TokenURL: TokenURL,
		},
		Scopes: []string{},
	}

	if len(scopes) > 0 {
		for _, scope := range scopes {
			c.Scopes = append(c.Scopes, scope)
		}
	}
	return c
}

func userFromReader(r io.Reader, user *goth.User) error {
	u := struct {
		Name      string `json:"name"`
		Email     string `json:"email"`
		NickName  string `json:"username"`
		ID        int    `json:"id"`
		AvatarURL string `json:"avatar_url"`
	}{}
	err := json.NewDecoder(r).Decode(&u)
	if err != nil {
		return err
	}
	user.Email = u.Email
	user.Name = u.Name
	user.NickName = u.NickName
	user.UserID = strconv.Itoa(u.ID)
	user.AvatarURL = u.AvatarURL
	return nil
}

//RefreshTokenAvailable refresh token is provided by auth provider or not
func (p *Provider) RefreshTokenAvailable() bool {
	return true
}

//RefreshToken get new access token based on the refresh token
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
	token := &oauth2.Token{RefreshToken: refreshToken}
	ts := p.config.TokenSource(goth.ContextForClient(p.Client()), token)
	newToken, err := ts.Token()
	if err != nil {
		return nil, err
	}
	return newToken, err
}
<commit_msg>gitlab add new customised url function<commit_after>// Package gitlab implements the OAuth2 protocol for authenticating users through gitlab.
// This package can be used as a reference implementation of an OAuth2 provider for Goth.
package gitlab

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"

	"github.com/markbates/goth"
	"golang.org/x/oauth2"
	"fmt"
)

// These vars define the Authentication, Token, and Profile URLS for Gitlab. If
// using Gitlab CE or EE, you should change these values before calling New.
//
// Examples:
//	gitlab.AuthURL = "https://gitlab.acme.com/oauth/authorize
//	gitlab.TokenURL = "https://gitlab.acme.com/oauth/token
//	gitlab.ProfileURL = "https://gitlab.acme.com/api/v3/user
var (
	AuthURL    = "https://gitlab.com/oauth/authorize"
	TokenURL   = "https://gitlab.com/oauth/token"
	ProfileURL = "https://gitlab.com/api/v3/user"
)

// Provider is the implementation of `goth.Provider` for accessing Gitlab.
type Provider struct {
	ClientKey    string
	Secret       string
	CallbackURL  string
	HTTPClient   *http.Client
	config       *oauth2.Config
	providerName string
	authURL      string
	tokenURL     string
	profileURL   string
}

// New creates a new Gitlab provider and sets up important connection details.
// You should always call `gitlab.New` to get a new provider.  Never try to
// create one manually.
func New(clientKey, secret, callbackURL string, scopes ...string) *Provider {
	return NewCustomisedURL(clientKey, secret, callbackURL, AuthURL, TokenURL, ProfileURL, scopes...)
}

// NewCustomisedURL is similar to New(...) but can be used to set custom URLs to connect to
func NewCustomisedURL(clientKey, secret, callbackURL, authURL, tokenURL, profileURL string, scopes ...string) *Provider {
	p := &Provider{
		ClientKey:    clientKey,
		Secret:       secret,
		CallbackURL:  callbackURL,
		providerName: "gitlab",
		authURL:      authURL,
		tokenURL:     tokenURL,
		profileURL:   profileURL,
	}
	p.config = newConfig(p, scopes)
	return p
}

// Name is the name used to retrieve this provider later.
func (p *Provider) Name() string {
	return p.providerName
}

// SetName is to update the name of the provider (needed in case of multiple providers of 1 type)
func (p *Provider) SetName(name string) {
	p.providerName = name
}

func (p *Provider) Client() *http.Client {
	return goth.HTTPClientWithFallBack(p.HTTPClient)
}

// Debug is a no-op for the gitlab package.
func (p *Provider) Debug(debug bool) {}

// BeginAuth asks Gitlab for an authentication end-point.
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
	return &Session{
		AuthURL: p.config.AuthCodeURL(state),
	}, nil
}

// FetchUser will go to Gitlab and access basic information about the user.
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
	sess := session.(*Session)
	user := goth.User{
		AccessToken:  sess.AccessToken,
		Provider:     p.Name(),
		RefreshToken: sess.RefreshToken,
		ExpiresAt:    sess.ExpiresAt,
	}

	if user.AccessToken == "" {
		// data is not yet retrieved since accessToken is still empty
		return user, fmt.Errorf("%s cannot get user information without accessToken", p.providerName)
	}

	response, err := p.Client().Get(p.profileURL + "?access_token=" + url.QueryEscape(sess.AccessToken))
	if err != nil {
		if response != nil {
			response.Body.Close()
		}
		return user, err
	}

	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return user, fmt.Errorf("%s responded with a %d trying to fetch user information", p.providerName, response.StatusCode)
	}

	bits, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return user, err
	}

	err = json.NewDecoder(bytes.NewReader(bits)).Decode(&user.RawData)
	if err != nil {
		return user, err
	}

	err = userFromReader(bytes.NewReader(bits), &user)

	return user, err
}

func newConfig(provider *Provider, scopes []string) *oauth2.Config {
	c := &oauth2.Config{
		ClientID:     provider.ClientKey,
		ClientSecret: provider.Secret,
		RedirectURL:  provider.CallbackURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  provider.authURL,
			TokenURL: provider.tokenURL,
		},
		Scopes: []string{},
	}

	if len(scopes) > 0 {
		for _, scope := range scopes {
			c.Scopes = append(c.Scopes, scope)
		}
	}
	return c
}

func userFromReader(r io.Reader, user *goth.User) error {
	u := struct {
		Name      string `json:"name"`
		Email     string `json:"email"`
		NickName  string `json:"username"`
		ID        int    `json:"id"`
		AvatarURL string `json:"avatar_url"`
	}{}
	err := json.NewDecoder(r).Decode(&u)
	if err != nil {
		return err
	}
	user.Email = u.Email
	user.Name = u.Name
	user.NickName = u.NickName
	user.UserID = strconv.Itoa(u.ID)
	user.AvatarURL = u.AvatarURL
	return nil
}

//RefreshTokenAvailable refresh token is provided by auth provider or not
func (p *Provider) RefreshTokenAvailable() bool {
	return true
}

//RefreshToken get new access token based on the refresh token
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
	token := &oauth2.Token{RefreshToken: refreshToken}
	ts := p.config.TokenSource(goth.ContextForClient(p.Client()), token)
	newToken, err := ts.Token()
	if err != nil {
		return nil, err
	}
	return newToken, err
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"regexp"
	"strconv"
	"time"

	"github.com/crowdmob/goamz/aws"
	"github.com/crowdmob/goamz/s3"
	"github.com/goji/httpauth"
	"github.com/zenazn/goji"
	"github.com/zenazn/goji/web"
)

var (
	bucketName string
	baseURL    string
)

func init() {
	flag.StringVar(&bucketName, "b", "", "Bucket Name")
	flag.StringVar(&baseURL, "u", "", "Base URL")
}

func rootHandler(c web.C, w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Nothing to see here!")
}

func tweetbot(c web.C, w http.ResponseWriter, r *http.Request) {
	multiReader, err := r.MultipartReader()
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Open Bucket
	auth, err := aws.EnvAuth()
	if err != nil {
		panic(err.Error())
	}
	s := s3.New(auth, aws.EUWest)
	bucket := s.Bucket(bucketName)

	// We only want the first part, the media
	part, err := multiReader.NextPart()
	if err == io.EOF {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	originalFilename := part.FileName()
	timeStamp := time.Now().Unix()
	randomValue := rand.Intn(999999)
	filename := fmt.Sprintf("%x-%x-%s", timeStamp, randomValue, originalFilename)
	path := fmt.Sprintf("tweetbot/%s", filename)

	contentType := part.Header.Get("Content-Type")
	contentLength, err := strconv.ParseInt(part.Header.Get("Content-Length"), 10, 64)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	err = bucket.PutReader(path, part, contentLength, contentType, s3.PublicRead, s3.Options{CacheControl: "public, max-age=315360000"})
	if err != nil {
		panic(err.Error())
	}

	fmt.Printf("\nFile %s (%s) uploaded successfully.\n", originalFilename, path)

	url := fmt.Sprintf("%s/%s", baseURL, path)

	responseMap := map[string]string{"url": url}
	jsonResponse, _ := json.Marshal(responseMap)

	w.Header().Set("Content-Type", "application/json")
	fmt.Fprintf(w, string(jsonResponse))
}

func main() {
	goji.Use(httpauth.SimpleBasicAuth("x", "password"))

	goji.Get("/", rootHandler)

	re := regexp.MustCompile("/tweetbot")
	goji.Post(re, tweetbot)

	goji.Serve()
}
<commit_msg>Better filenames<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"math/big"
	"mime"
	"net/http"
	"path/filepath"
	"regexp"
	"strconv"
	"time"

	"github.com/crowdmob/goamz/aws"
	"github.com/crowdmob/goamz/s3"
	"github.com/goji/httpauth"
	"github.com/tv42/base58"
	"github.com/zenazn/goji"
	"github.com/zenazn/goji/web"
)

var (
	bucketName string
	baseURL    string
)

func init() {
	flag.StringVar(&bucketName, "b", "", "Bucket Name")
	flag.StringVar(&baseURL, "u", "", "Base URL")
}

func rootHandler(c web.C, w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Nothing to see here!")
}

func tweetbot(c web.C, w http.ResponseWriter, r *http.Request) {
	multiReader, err := r.MultipartReader()
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Open Bucket
	auth, err := aws.EnvAuth()
	if err != nil {
		panic(err.Error())
	}
	s := s3.New(auth, aws.EUWest)
	bucket := s.Bucket(bucketName)

	// We only want the first part, the media
	part, err := multiReader.NextPart()
	if err == io.EOF {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	originalFilename := part.FileName()
	fileExt := filepath.Ext(originalFilename)
	unixTime := time.Now().UTC().Unix()
	b58buf := base58.EncodeBig(nil, big.NewInt(unixTime))

	filename := fmt.Sprintf("%s%s", b58buf, fileExt)
	path := "tweetbot/" + filename

	contentType := part.Header.Get("Content-Type")
	if contentType == "" {
		contentType = mime.TypeByExtension(fileExt)

		if contentType == "" {
			contentType = "application/octet-stream"
		}
	}

	contentLength, err := strconv.ParseInt(part.Header.Get("Content-Length"), 10, 64)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	err = bucket.PutReader(path, part, contentLength, contentType, s3.PublicRead, s3.Options{CacheControl: "public, max-age=315360000"})
	if err != nil {
		panic(err.Error())
	}

	// fmt.Printf("\nFile %s (%s) uploaded successfully.\n", originalFilename, path)

	url := fmt.Sprintf("%s/%s", baseURL, path)

	responseMap := map[string]string{"url": url}
	jsonResponse, _ := json.Marshal(responseMap)

	w.Header().Set("Content-Type", "application/json")
	fmt.Fprintf(w, string(jsonResponse))
}

func main() {
	goji.Use(httpauth.SimpleBasicAuth("x", "password"))

	goji.Get("/", rootHandler)

	re := regexp.MustCompile("/tweetbot")
	goji.Post(re, tweetbot)

	goji.Serve()
}
<|endoftext|>


<commit_before>// Package goutils provides common utilities
//
// See also:
// https://godoc.org/github.com/palsivertsen/goutils/converters
// https://godoc.org/github.com/palsivertsen/goutils/validators
package goutils
<commit_msg>Removed links to sub packages as they are listed implicitly<commit_after>// Package goutils provides common utilities
package goutils
<|endoftext|>


<commit_before>// Copyright (c) 2014 Datacratic. All rights reserved.

package nfork

import (
	"encoding/json"
	"math/rand"
	"strconv"
	"sync"
	"time"
)

// Stats contains the stats of an outbound at a given point in time.
type Stats struct {

	// Requests counts the number of requests made.
	Requests uint64

	// Errors counts the number of errors encountered.
	Errors uint64

	// Timeouts counts the number of timeouts encountered.
	Timeouts uint64

	// Latency is the latency distribution of all requests.
	Latency Distribution

	// Responses counts the number of responses received for an HTTP status
	// code.
	Responses map[int]uint64
}

// MarshalJSON defines a custom JSON format for encoding/json.
func (stats *Stats) MarshalJSON() ([]byte, error) {
	var statsJSON struct {
		Requests  uint64            `json:"requests"`
		Errors    uint64            `json:"errors"`
		Timeouts  uint64            `json:"timeouts"`
		Latency   map[string]string `json:"latency"`
		Responses map[string]uint64 `json:"responses"`
	}

	statsJSON.Requests = stats.Requests
	statsJSON.Errors = stats.Errors
	statsJSON.Timeouts = stats.Timeouts
	statsJSON.Latency = make(map[string]string)
	statsJSON.Responses = make(map[string]uint64)

	p50, p90, p99, max := stats.Latency.Percentiles()
	statsJSON.Latency["p50"] = time.Duration(p50).String()
	statsJSON.Latency["p90"] = time.Duration(p90).String()
	statsJSON.Latency["p99"] = time.Duration(p99).String()
	statsJSON.Latency["pmx"] = time.Duration(max).String()

	for code, count := range stats.Responses {
		statsJSON.Responses[strconv.Itoa(code)] = count
	}

	return json.Marshal(&statsJSON)
}

// Event contains the outcome of an HTTP request.
type Event struct {

	// Error indicates that an error occured.
	Error bool

	// Timeout indicates that the request timed out.
	Timeout bool

	// Response is the HTTP response code received.
	Response int

	// Latency mesures the latency of the request.
	Latency time.Duration
}

// DefaultSampleRate is used if Rate is not set set in StatsRecorder.
const DefaultSampleRate = 10 * time.Second

// StatsRecorder records stats for a given outbound and updates them at a
// given rate.
type StatsRecorder struct {

	// Rate at which stats are updated.
	Rate time.Duration

	// Rand is the RNG used for stats sampling.
	Rand *rand.Rand

	initialize sync.Once

	mutex         sync.Mutex
	current, prev *Stats

	shutdownC chan int
}

// Init initializes the object.
func (recorder *StatsRecorder) Init() {
	recorder.initialize.Do(recorder.init)
}

func (recorder *StatsRecorder) init() {
	if recorder.Rate == 0 {
		recorder.Rate = DefaultSampleRate
	}

	if recorder.Rand == nil {
		recorder.Rand = rand.New(rand.NewSource(0))
	}

	recorder.prev = new(Stats)
	recorder.current = new(Stats)

	recorder.shutdownC = make(chan int)
	go recorder.run()
}

// Close terminates the stats recorder.
func (recorder *StatsRecorder) Close() {
	recorder.Init()
	recorder.shutdownC <- 1
}

// Record records the given outcome.
func (recorder *StatsRecorder) Record(event Event) {
	recorder.Init()
	recorder.mutex.Lock()

	stats := recorder.current

	stats.Requests++
	stats.Latency.Sample(uint64(event.Latency))

	if event.Error {
		stats.Errors++

	} else if event.Timeout {
		stats.Timeouts++

	} else {
		if stats.Responses == nil {
			stats.Responses = make(map[int]uint64)
		}
		stats.Responses[event.Response]++
	}

	recorder.mutex.Unlock()
}

// Read returns the last updated stats.
func (recorder *StatsRecorder) Read() (stats *Stats) {
	recorder.Init()
	recorder.mutex.Lock()

	stats = recorder.prev

	recorder.mutex.Unlock()
	return
}

func (recorder *StatsRecorder) run() {
	tick := time.NewTicker(recorder.Rate)
	for {
		select {
		case <-tick.C:
			recorder.mutex.Lock()

			recorder.prev = recorder.current
			recorder.current = new(Stats)

			recorder.mutex.Unlock()

		case <-recorder.shutdownC:
			tick.Stop()
			return
		}
	}
}
<commit_msg>undo unnecessary white spaces<commit_after>// Copyright (c) 2014 Datacratic. All rights reserved.

package nfork

import (
	"encoding/json"
	"math/rand"
	"strconv"
	"sync"
	"time"
)

// Stats contains the stats of an outbound at a given point in time.
type Stats struct {

	// Requests counts the number of requests made.
	Requests uint64

	// Errors counts the number of errors encountered.
	Errors uint64

	// Timeouts counts the number of timeouts encountered.
	Timeouts uint64

	// Latency is the latency distribution of all requests.
	Latency Distribution

	// Responses counts the number of responses received for an HTTP status
	// code.
	Responses map[int]uint64
}

// MarshalJSON defines a custom JSON format for encoding/json.
func (stats *Stats) MarshalJSON() ([]byte, error) {
	var statsJSON struct {
		Requests  uint64            `json:"requests"`
		Errors    uint64            `json:"errors"`
		Timeouts  uint64            `json:"timeouts"`
		Latency   map[string]string `json:"latency"`
		Responses map[string]uint64 `json:"responses"`
	}

	statsJSON.Requests = stats.Requests
	statsJSON.Errors = stats.Errors
	statsJSON.Timeouts = stats.Timeouts
	statsJSON.Latency = make(map[string]string)
	statsJSON.Responses = make(map[string]uint64)

	p50, p90, p99, max := stats.Latency.Percentiles()
	statsJSON.Latency["p50"] = time.Duration(p50).String()
	statsJSON.Latency["p90"] = time.Duration(p90).String()
	statsJSON.Latency["p99"] = time.Duration(p99).String()
	statsJSON.Latency["pmx"] = time.Duration(max).String()

	for code, count := range stats.Responses {
		statsJSON.Responses[strconv.Itoa(code)] = count
	}

	return json.Marshal(&statsJSON)
}

// Event contains the outcome of an HTTP request.
type Event struct {

	// Error indicates that an error occured.
	Error bool

	// Timeout indicates that the request timed out.
	Timeout bool

	// Response is the HTTP response code received.
	Response int

	// Latency mesures the latency of the request.
	Latency time.Duration
}

// DefaultSampleRate is used if Rate is not set set in StatsRecorder.
const DefaultSampleRate = 10 * time.Second

// StatsRecorder records stats for a given outbound and updates them at a
// given rate.
type StatsRecorder struct {

	// Rate at which stats are updated.
	Rate time.Duration

	// Rand is the RNG used for stats sampling.
	Rand *rand.Rand

	initialize sync.Once

	mutex         sync.Mutex
	current, prev *Stats

	shutdownC chan int
}

// Init initializes the object.
func (recorder *StatsRecorder) Init() {
	recorder.initialize.Do(recorder.init)
}

func (recorder *StatsRecorder) init() {
	if recorder.Rate == 0 {
		recorder.Rate = DefaultSampleRate
	}

	if recorder.Rand == nil {
		recorder.Rand = rand.New(rand.NewSource(0))
	}

	recorder.prev = new(Stats)
	recorder.current = new(Stats)

	recorder.shutdownC = make(chan int)
	go recorder.run()
}

// Close terminates the stats recorder.
func (recorder *StatsRecorder) Close() {
	recorder.Init()
	recorder.shutdownC <- 1
}

// Record records the given outcome.
func (recorder *StatsRecorder) Record(event Event) {
	recorder.Init()
	recorder.mutex.Lock()

	stats := recorder.current

	stats.Requests++
	stats.Latency.Sample(uint64(event.Latency))

	if event.Error {
		stats.Errors++

	} else if event.Timeout {
		stats.Timeouts++

	} else {
		if stats.Responses == nil {
			stats.Responses = make(map[int]uint64)
		}
		stats.Responses[event.Response]++
	}

	recorder.mutex.Unlock()
}

// Read returns the last updated stats.
func (recorder *StatsRecorder) Read() (stats *Stats) {
	recorder.Init()
	recorder.mutex.Lock()

	stats = recorder.prev

	recorder.mutex.Unlock()
	return
}

func (recorder *StatsRecorder) run() {
	tick := time.NewTicker(recorder.Rate)
	for {
		select {
		case <-tick.C:
			recorder.mutex.Lock()

			recorder.prev = recorder.current
			recorder.current = new(Stats)

			recorder.mutex.Unlock()

		case <-recorder.shutdownC:
			tick.Stop()
			return
		}
	}
}
<|endoftext|>