<commit_before>package topom

import (
	"math"

	"github.com/wandoulabs/codis/pkg/models"
	"github.com/wandoulabs/codis/pkg/proxy"
	"github.com/wandoulabs/codis/pkg/utils/errors"
	"github.com/wandoulabs/codis/pkg/utils/log"
)

func (s *Topom) ListGroup() []*models.Group {
	s.mu.RLock()
	defer s.mu.RUnlock()
	glist := make([]*models.Group, 0, len(s.groups))
	for _, g := range s.groups {
		glist = append(glist, g)
	}
	return glist
}

func (s *Topom) getGroup(groupId int) (*models.Group, error) {
	if g := s.groups[groupId]; g != nil {
		return g, nil
	}
	return nil, errors.New("group does not exist")
}

func (s *Topom) getGroupMaster(groupId int) string {
	if g := s.groups[groupId]; g != nil && len(g.Servers) != 0 {
		return g.Servers[0]
	}
	return ""
}

func (s *Topom) CreateGroup(groupId int) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedTopom
	}

	if groupId <= 0 || groupId > math.MaxInt32 {
		return errors.New("invalid group id")
	}
	if s.groups[groupId] != nil {
		return errors.New("group already exists")
	}

	g := &models.Group{
		Id: groupId,
	}
	if err := s.store.CreateGroup(groupId, g); err != nil {
		log.WarnErrorf(err, "group-[%d] create failed", groupId)
		return errors.New("group create failed")
	}

	log.Infof("[%p] create group-[%d]", s, groupId)

	s.groups[groupId] = g
	return nil
}

func (s *Topom) RemoveGroup(groupId int) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedTopom
	}

	_, err := s.getGroup(groupId)
	if err != nil {
		return err
	}
	if len(s.getSlotsByGroup(groupId)) != 0 {
		return errors.New("group is still busy")
	}

	if err := s.store.RemoveGroup(groupId); err != nil {
		log.WarnErrorf(err, "group-[%d] remove failed", groupId)
		return errors.New("group remove failed")
	}

	log.Infof("[%p] remove group-[%d]", s, groupId)

	delete(s.groups, groupId)
	return nil
}

func (s *Topom) repairGroup(groupId int) error {
	g, err := s.getGroup(groupId)
	if err != nil {
		return err
	}
	for i, addr := range g.Servers {
		c, err := s.redisp.GetClient(addr)
		if err != nil {
			log.WarnErrorf(err, "group-[%d] repair failed, server[%d] = %s", groupId, i, addr)
			return errors.New("create redis client failed")
		}
		defer s.redisp.PutClient(c)

		var master = ""
		if i != 0 {
			master = g.Servers[0]
		}
		if err := c.SlaveOf(master); err != nil {
			log.WarnErrorf(err, "group-[%d] repair failed, server[%d] = %s", groupId, i, addr)
			return errors.New("server set slaveof failed")
		}
	}
	return nil
}

func (s *Topom) resyncGroup(groupId int, locked bool) map[string]error {
	slots := s.getSlotsByGroup(groupId)
	if locked {
		for _, slot := range slots {
			slot.Locked = true
		}
	}
	return s.broadcast(func(p *models.Proxy, c *proxy.ApiClient) error {
		if len(slots) == 0 {
			return nil
		}
		if err := c.FillSlots(slots...); err != nil {
			log.WarnErrorf(err, "proxy-[%s] resync group-[%d] failed", p.Token, groupId)
			return errors.New("proxy resync group failed")
		}
		return nil
	})
}

func (s *Topom) RepairGroup(groupId int) (map[string]error, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed {
		return nil, ErrClosedTopom
	}
	if err := s.repairGroup(groupId); err != nil {
		return nil, err
	}
	return s.resyncGroup(groupId, false), nil
}

func (s *Topom) ResyncGroup(groupId int) (map[string]error, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed {
		return nil, ErrClosedTopom
	}
	_, err := s.getGroup(groupId)
	if err != nil {
		return nil, err
	}
	return s.resyncGroup(groupId, false), nil
}

func (s *Topom) GroupAddNewServer(groupId int, addr string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedTopom
	}

	for _, g := range s.groups {
		for _, x := range g.Servers {
			if x == addr {
				return errors.New("server already exists")
			}
		}
	}
	g, err := s.getGroup(groupId)
	if err != nil {
		return err
	}

	c, err := s.redisp.GetClient(addr)
	if err != nil {
		log.WarnErrorf(err, "group-[%d] add server failed, server = %s", groupId, addr)
		return errors.New("create redis client failed")
	}
	defer s.redisp.PutClient(c)

	if _, err := c.SlotsInfo(); err != nil {
		log.WarnErrorf(err, "group-[%d] add server failed, server = %s", groupId, addr)
		return errors.New("server check slots failed")
	}
	if err := c.SlaveOf(s.getGroupMaster(groupId)); err != nil {
		log.WarnErrorf(err, "group-[%d] add server failed, server = %s", groupId, addr)
		return errors.New("server set slaveof failed")
	}

	log.Infof("[%p] group-[%d] add server = %s", s, groupId, addr)

	n := &models.Group{
		Id:      groupId,
		Servers: append(g.Servers, addr),
	}
	if err := s.store.UpdateGroup(groupId, n); err != nil {
		log.WarnErrorf(err, "group-[%d] update failed", groupId)
		return errors.New("group update failed")
	}

	log.Infof("[%p] update group: %+v", s, n)

	s.groups[groupId] = n
	return nil
}

func (s *Topom) GroupRemoveServer(groupId int, addr string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedTopom
	}

	g, err := s.getGroup(groupId)
	if err != nil {
		return err
	}
	slist := []string{}
	for _, x := range g.Servers {
		if x != addr {
			slist = append(slist, x)
		}
	}
	if len(slist) == len(g.Servers) {
		return errors.New("server does not exist")
	}
	if addr == g.Servers[0] {
		if len(g.Servers) != 1 || len(s.getSlotsByGroup(groupId)) != 0 {
			return errors.New("server is still busy")
		}
	}

	log.Infof("[%p] group-[%d] remove server = %s", s, groupId, addr)

	n := &models.Group{
		Id:      groupId,
		Servers: slist,
	}
	if err := s.store.UpdateGroup(groupId, n); err != nil {
		log.WarnErrorf(err, "group-[%d] update failed", groupId)
		return errors.New("group update failed")
	}

	log.Infof("[%p] update group: %+v", s, n)

	s.groups[groupId] = n
	return nil
}

func (s *Topom) GroupPromoteServer(groupId int, addr string, force bool) (map[string]error, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, ErrClosedTopom
	}

	g, err := s.getGroup(groupId)
	if err != nil {
		return nil, err
	}
	slist := []string{}
	for _, x := range g.Servers {
		if x != addr {
			slist = append(slist, x)
		}
	}
	if len(slist) == len(g.Servers) {
		return nil, errors.New("server does not exist")
	}
	if addr == g.Servers[0] {
		return nil, errors.New("server promote again")
	} else {
		slist = append([]string{addr}, slist[1:]...)
	}

	log.Infof("group-[%d] will be locked for server promotion", groupId)

	if errs := s.resyncGroup(groupId, true); len(errs) != 0 {
		if !force {
			return errs, nil
		}
		log.Warnf("group-[%d] force promote with resync errors", groupId)
	}

	log.Infof("[%p] group-[%d] promote server = %s", s, groupId, addr)

	n := &models.Group{
		Id:      groupId,
		Servers: slist,
	}
	if err := s.store.UpdateGroup(groupId, n); err != nil {
		log.WarnErrorf(err, "group-[%d] update failed", groupId)
		return nil, errors.New("group update failed")
	}

	log.Infof("[%p] update group: %+v", s, n)

	s.groups[groupId] = n

	if err := s.repairGroup(groupId); err != nil {
		log.WarnErrorf(err, "group-[%d] repair failed", groupId)
		return nil, errors.New("group repair failed")
	}
	return s.resyncGroup(groupId, false), nil
}
<commit_msg>Update, update topom_group, more logs<commit_after>package topom

import (
	"math"
	"time"

	"github.com/wandoulabs/codis/pkg/models"
	"github.com/wandoulabs/codis/pkg/proxy"
	"github.com/wandoulabs/codis/pkg/utils/errors"
	"github.com/wandoulabs/codis/pkg/utils/log"
)

func (s *Topom) daemonRedisPool() {
	var ticker = time.NewTicker(time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-s.exit.C:
			return
		case <-ticker.C:
			s.redisp.Cleanup()
		}
	}
}

func (s *Topom) ListGroup() []*models.Group {
	s.mu.RLock()
	defer s.mu.RUnlock()
	glist := make([]*models.Group, 0, len(s.groups))
	for _, g := range s.groups {
		glist = append(glist, g)
	}
	return glist
}

func (s *Topom) getGroup(groupId int) (*models.Group, error) {
	if g := s.groups[groupId]; g != nil {
		return g, nil
	}
	return nil, errors.New("group does not exist")
}

func (s *Topom) getGroupMaster(groupId int) string {
	if g := s.groups[groupId]; g != nil && len(g.Servers) != 0 {
		return g.Servers[0]
	}
	return ""
}

func (s *Topom) CreateGroup(groupId int) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedTopom
	}

	if groupId <= 0 || groupId > math.MaxInt32 {
		return errors.New("invalid group id")
	}
	if s.groups[groupId] != nil {
		return errors.New("group already exists")
	}

	g := &models.Group{
		Id: groupId,
	}
	if err := s.store.CreateGroup(groupId, g); err != nil {
		log.WarnErrorf(err, "group-[%d] create failed", groupId)
		return errors.New("group create failed")
	}

	log.Infof("[%p] create group-[%d]", s, groupId)

	s.groups[groupId] = g
	return nil
}

func (s *Topom) RemoveGroup(groupId int) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedTopom
	}

	_, err := s.getGroup(groupId)
	if err != nil {
		return err
	}
	if len(s.getSlotsByGroup(groupId)) != 0 {
		return errors.New("group is still busy")
	}

	if err := s.store.RemoveGroup(groupId); err != nil {
		log.WarnErrorf(err, "group-[%d] remove failed", groupId)
		return errors.New("group remove failed")
	}

	log.Infof("[%p] remove group-[%d]", s, groupId)

	delete(s.groups, groupId)
	return nil
}

func (s *Topom) repairGroup(groupId int) error {
	g, err := s.getGroup(groupId)
	if err != nil {
		return err
	}
	for i, addr := range g.Servers {
		c, err := s.redisp.GetClient(addr)
		if err != nil {
			log.WarnErrorf(err, "group-[%d] repair failed, server[%d] = %s", groupId, i, addr)
			return errors.New("create redis client failed")
		}
		defer s.redisp.PutClient(c)

		var master = ""
		if i != 0 {
			master = g.Servers[0]
		}
		if err := c.SlaveOf(master); err != nil {
			log.WarnErrorf(err, "group-[%d] repair failed, server[%d] = %s", groupId, i, addr)
			return errors.New("server set slaveof failed")
		}
	}
	return nil
}

func (s *Topom) resyncGroup(groupId int, locked bool) map[string]error {
	slots := s.getSlotsByGroup(groupId)
	if locked {
		for _, slot := range slots {
			slot.Locked = true
		}
	}
	return s.broadcast(func(p *models.Proxy, c *proxy.ApiClient) error {
		if len(slots) == 0 {
			return nil
		}
		if err := c.FillSlots(slots...); err != nil {
			log.WarnErrorf(err, "proxy-[%s] resync group-[%d] failed", p.Token, groupId)
			return errors.New("proxy resync group failed")
		}
		return nil
	})
}

func (s *Topom) RepairGroup(groupId int) (map[string]error, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed {
		return nil, ErrClosedTopom
	}
	if err := s.repairGroup(groupId); err != nil {
		return nil, err
	}
	return s.resyncGroup(groupId, false), nil
}

func (s *Topom) ResyncGroup(groupId int) (map[string]error, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed {
		return nil, ErrClosedTopom
	}
	_, err := s.getGroup(groupId)
	if err != nil {
		return nil, err
	}
	return s.resyncGroup(groupId, false), nil
}

func (s *Topom) GroupAddNewServer(groupId int, addr string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedTopom
	}

	for _, g := range s.groups {
		for _, x := range g.Servers {
			if x == addr {
				return errors.New("server already exists")
			}
		}
	}
	g, err := s.getGroup(groupId)
	if err != nil {
		return err
	}

	c, err := s.redisp.GetClient(addr)
	if err != nil {
		log.WarnErrorf(err, "group-[%d] add server failed, server = %s", groupId, addr)
		return errors.New("create redis client failed")
	}
	defer s.redisp.PutClient(c)

	if _, err := c.SlotsInfo(); err != nil {
		log.WarnErrorf(err, "group-[%d] add server failed, server = %s", groupId, addr)
		return errors.New("server check slots failed")
	}
	if err := c.SlaveOf(s.getGroupMaster(groupId)); err != nil {
		log.WarnErrorf(err, "group-[%d] add server failed, server = %s", groupId, addr)
		return errors.New("server set slaveof failed")
	}

	log.Infof("[%p] group-[%d] add server = %s", s, groupId, addr)

	n := &models.Group{
		Id:      groupId,
		Servers: append(g.Servers, addr),
	}
	if err := s.store.UpdateGroup(groupId, n); err != nil {
		log.WarnErrorf(err, "group-[%d] update failed", groupId)
		return errors.New("group update failed")
	}

	log.Infof("[%p] update group: %+v", s, n)

	s.groups[groupId] = n
	return nil
}

func (s *Topom) GroupRemoveServer(groupId int, addr string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedTopom
	}

	g, err := s.getGroup(groupId)
	if err != nil {
		return err
	}
	slist := []string{}
	for _, x := range g.Servers {
		if x != addr {
			slist = append(slist, x)
		}
	}
	if len(slist) == len(g.Servers) {
		return errors.New("server does not exist")
	}
	if addr == g.Servers[0] {
		if len(g.Servers) != 1 || len(s.getSlotsByGroup(groupId)) != 0 {
			return errors.New("server is still busy")
		}
	}

	log.Infof("[%p] group-[%d] remove server = %s", s, groupId, addr)

	n := &models.Group{
		Id:      groupId,
		Servers: slist,
	}
	if err := s.store.UpdateGroup(groupId, n); err != nil {
		log.WarnErrorf(err, "group-[%d] update failed", groupId)
		return errors.New("group update failed")
	}

	log.Infof("[%p] update group: %+v", s, n)

	s.groups[groupId] = n
	return nil
}

func (s *Topom) GroupPromoteServer(groupId int, addr string, force bool) (map[string]error, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, ErrClosedTopom
	}

	g, err := s.getGroup(groupId)
	if err != nil {
		return nil, err
	}
	slist := []string{}
	for _, x := range g.Servers {
		if x != addr {
			slist = append(slist, x)
		}
	}
	if len(slist) == len(g.Servers) {
		return nil, errors.New("server does not exist")
	}
	if addr == g.Servers[0] {
		return nil, errors.New("server promote again")
	} else {
		slist = append([]string{addr}, slist[1:]...)
	}

	log.Infof("group-[%d] will be locked by force", groupId)

	if errs := s.resyncGroup(groupId, true); len(errs) != 0 {
		if !force {
			return errs, nil
		}
		log.Warnf("group-[%d] force promote with resync errors", groupId)
	}

	log.Infof("[%p] group-[%d] promote server = %s", s, groupId, addr)

	n := &models.Group{
		Id:      groupId,
		Servers: slist,
	}
	if err := s.store.UpdateGroup(groupId, n); err != nil {
		log.WarnErrorf(err, "group-[%d] update failed", groupId)
		return nil, errors.New("group update failed")
	}

	log.Infof("[%p] update group: %+v", s, n)

	s.groups[groupId] = n

	if err := s.repairGroup(groupId); err != nil {
		log.WarnErrorf(err, "group-[%d] repair failed", groupId)
		return nil, errors.New("group repair failed")
	}

	log.Infof("group-[%d] will be recovered", groupId)

	return s.resyncGroup(groupId, false), nil
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package broker

import (
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/stats"
	"github.com/apex/log"
	"github.com/brocaar/lorawan"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// component implements the core.Component interface
type component struct {
	NetworkController
	ctx log.Interface
}

// New construct a new Broker component
func New(controller NetworkController, ctx log.Interface) core.BrokerServer {
	return component{NetworkController: controller, ctx: ctx}
}

// HandleData implements the core.RouterClient interface
func (b component) HandleData(bctx context.Context, req *core.DataBrokerReq) (*core.DataBrokerRes, error) {
	// Get some logs / analytics
	stats.MarkMeter("broker.uplink.in")
	b.ctx.Debug("Handling uplink packet")

	// Validate incoming data
	uplinkPayload, err := core.NewLoRaWANData(req.Payload, true)
	if err != nil {
		return nil, errors.New(errors.Structural, err)
	}
	devAddr := req.Payload.MACPayload.FHDR.DevAddr // No nil ref, ensured by NewLoRaWANData()
	ctx := b.ctx.WithField("DevAddr", devAddr)

	// Check whether we should handle it
	entries, err := b.LookupDevices(devAddr)
	if err != nil {
		switch err.(errors.Failure).Nature {
		case errors.NotFound:
			stats.MarkMeter("broker.uplink.handler_lookup.device_not_found")
			ctx.Debug("Uplink device not found")
		default:
			b.ctx.Warn("Database lookup failed")
		}
		return nil, err
	}
	stats.UpdateHistogram("broker.uplink.handler_lookup.entries", int64(len(entries)))

	// Several handlers might be associated to the same device, we distinguish them using
	// MIC check. Only one should verify the MIC check
	// The device only stores a 16-bits counter but could reflect a 32-bits one.
	// The counter is used for the MIC computation, thus, we're gonna try both 16-bits and
	// 32-bits counters.
	// We keep track of the real counter in the network controller.
	fhdr := &uplinkPayload.MACPayload.(*lorawan.MACPayload).FHDR // No nil ref, ensured by NewLoRaWANData()
	fcnt16 := fhdr.FCnt                                          // Keep a reference to the original counter

	var mEntry *devEntry
	for _, entry := range entries {
		// retrieve the network session key
		key := lorawan.AES128Key(entry.NwkSKey)

		// Check with 16-bits counters
		fhdr.FCnt = fcnt16
		ok, err := uplinkPayload.ValidateMIC(key)
		if err != nil {
			continue
		}

		if !ok && entry.FCntUp > 65535 { // Check with 32-bits counter
			fcnt32, err := b.WholeCounter(fcnt16, entry.FCntUp)
			if err != nil {
				continue
			}
			fhdr.FCnt = fcnt32
			ok, err = uplinkPayload.ValidateMIC(key)
		}

		if ok {
			mEntry = &entry
			stats.MarkMeter("broker.uplink.handler_lookup.mic_match")
			ctx.WithField("handler", entry.HandlerNet).Debug("MIC check succeeded")
			break // We stop at the first valid check ...
		}
	}

	if mEntry == nil {
		stats.MarkMeter("broker.uplink.handler_lookup.no_mic_match")
		err := errors.New(errors.NotFound, "MIC check returned no matches")
		ctx.Debug(err.Error())
		return nil, err
	}

	// It does matter here to use the DevEUI from the entry and not from the packet.
	// The packet actually holds a DevAddr and the real DevEUI has been determined thanks
	// to the MIC check + persistence
	if err := b.UpdateFCnt(mEntry.AppEUI, mEntry.DevEUI, fhdr.FCnt); err != nil {
		return nil, err
	}

	// Then we forward the packet to the handler and wait for the response
	conn, err := grpc.Dial(mEntry.HandlerNet, grpc.WithInsecure(), grpc.WithTimeout(time.Second*2))
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	handler := core.NewHandlerClient(conn)
	resp, err := handler.HandleData(context.Background(), &core.DataHandlerReq{
		Payload:  req.Payload.MACPayload.FRMPayload,
		DevEUI:   mEntry.DevEUI,
		AppEUI:   mEntry.AppEUI,
		FCnt:     fhdr.FCnt,
		MType:    req.Payload.MHDR.MType,
		Metadata: req.Metadata,
	})

	if err != nil && !strings.Contains(err.Error(), string(errors.NotFound)) { // FIXME Find better way to analyze error
		stats.MarkMeter("broker.uplink.bad_handler_response")
		return nil, errors.New(errors.Operational, err)
	}
	stats.MarkMeter("broker.uplink.ok")

	// No response, we stop here and propagate the "no answer".
	// In case of confirmed data, the handler is in charge of creating the confirmation
	if resp == nil {
		return nil, nil
	}

	// If a response was sent, i.e. a downlink data, we need to compute the right MIC
	downlinkPayload, err := core.NewLoRaWANData(resp.Payload, false)
	if err != nil {
		return nil, errors.New(errors.Structural, err)
	}
	stats.MarkMeter("broker.downlink.in")
	if err := downlinkPayload.SetMIC(lorawan.AES128Key(mEntry.NwkSKey)); err != nil {
		return nil, errors.New(errors.Structural, "Unable to set response MIC")
	}
	resp.Payload.MIC = downlinkPayload.MIC[:]

	// And finally, we acknowledge the answer
	return &core.DataBrokerRes{
		Payload:  resp.Payload,
		Metadata: resp.Metadata,
	}, nil
}

// Register implements the core.BrokerServer interface
func (b component) SubscribePersonalized(bctx context.Context, req *core.SubBrokerReq) (*core.SubBrokerRes, error) {
	b.ctx.Debug("Handling personalized subscription")

	// Ensure the entry is valid
	if len(req.AppEUI) != 8 {
		return nil, errors.New(errors.Structural, "Invalid Application EUI")
	}

	if len(req.DevEUI) != 8 {
		return nil, errors.New(errors.Structural, "Invalid Device EUI")
	}

	var nwkSKey [16]byte
	if len(req.NwkSKey) != 16 {
		return nil, errors.New(errors.Structural, "Invalid Network Session Key")
	}
	copy(nwkSKey[:], req.NwkSKey)

	re := regexp.MustCompile("^[-\\w]+\\.([-\\w]+\\.?)+:\\d+$")
	if !re.MatchString(req.HandlerNet) {
		return nil, errors.New(errors.Structural, fmt.Sprintf("Invalid Handler Net Address. Should match: %s", re))
	}

	return nil, b.StoreDevice(devEntry{
		HandlerNet: req.HandlerNet,
		AppEUI:     req.AppEUI,
		DevEUI:     req.DevEUI,
		NwkSKey:    nwkSKey,
		FCntUp:     0,
	})
}
<commit_msg>[refactor/grpc] Make broker consistent with new handler<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package broker

import (
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/stats"
	"github.com/apex/log"
	"github.com/brocaar/lorawan"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// component implements the core.Component interface
type component struct {
	NetworkController
	ctx log.Interface
}

// New construct a new Broker component
func New(controller NetworkController, ctx log.Interface) core.BrokerServer {
	return component{NetworkController: controller, ctx: ctx}
}

// HandleData implements the core.RouterClient interface
func (b component) HandleData(bctx context.Context, req *core.DataBrokerReq) (*core.DataBrokerRes, error) {
	// Get some logs / analytics
	stats.MarkMeter("broker.uplink.in")
	b.ctx.Debug("Handling uplink packet")

	// Validate incoming data
	uplinkPayload, err := core.NewLoRaWANData(req.Payload, true)
	if err != nil {
		return nil, errors.New(errors.Structural, err)
	}
	devAddr := req.Payload.MACPayload.FHDR.DevAddr // No nil ref, ensured by NewLoRaWANData()
	ctx := b.ctx.WithField("DevAddr", devAddr)

	// Check whether we should handle it
	entries, err := b.LookupDevices(devAddr)
	if err != nil {
		switch err.(errors.Failure).Nature {
		case errors.NotFound:
			stats.MarkMeter("broker.uplink.handler_lookup.device_not_found")
			ctx.Debug("Uplink device not found")
		default:
			b.ctx.Warn("Database lookup failed")
		}
		return nil, err
	}
	stats.UpdateHistogram("broker.uplink.handler_lookup.entries", int64(len(entries)))

	// Several handlers might be associated to the same device, we distinguish them using
	// MIC check. Only one should verify the MIC check
	// The device only stores a 16-bits counter but could reflect a 32-bits one.
	// The counter is used for the MIC computation, thus, we're gonna try both 16-bits and
	// 32-bits counters.
	// We keep track of the real counter in the network controller.
	fhdr := &uplinkPayload.MACPayload.(*lorawan.MACPayload).FHDR // No nil ref, ensured by NewLoRaWANData()
	fcnt16 := fhdr.FCnt                                          // Keep a reference to the original counter

	var mEntry *devEntry
	for _, entry := range entries {
		// retrieve the network session key
		key := lorawan.AES128Key(entry.NwkSKey)

		// Check with 16-bits counters
		fhdr.FCnt = fcnt16
		ok, err := uplinkPayload.ValidateMIC(key)
		if err != nil {
			continue
		}

		if !ok && entry.FCntUp > 65535 { // Check with 32-bits counter
			fcnt32, err := b.WholeCounter(fcnt16, entry.FCntUp)
			if err != nil {
				continue
			}
			fhdr.FCnt = fcnt32
			ok, err = uplinkPayload.ValidateMIC(key)
		}

		if ok {
			mEntry = &entry
			stats.MarkMeter("broker.uplink.handler_lookup.mic_match")
			ctx.WithField("handler", entry.HandlerNet).Debug("MIC check succeeded")
			break // We stop at the first valid check ...
		}
	}

	if mEntry == nil {
		stats.MarkMeter("broker.uplink.handler_lookup.no_mic_match")
		err := errors.New(errors.NotFound, "MIC check returned no matches")
		ctx.Debug(err.Error())
		return nil, err
	}

	// It does matter here to use the DevEUI from the entry and not from the packet.
	// The packet actually holds a DevAddr and the real DevEUI has been determined thanks
	// to the MIC check + persistence
	if err := b.UpdateFCnt(mEntry.AppEUI, mEntry.DevEUI, fhdr.FCnt); err != nil {
		return nil, err
	}

	// Then we forward the packet to the handler and wait for the response
	conn, err := grpc.Dial(mEntry.HandlerNet, grpc.WithInsecure(), grpc.WithTimeout(time.Second*2))
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	handler := core.NewHandlerClient(conn)
	resp, err := handler.HandleDataUp(context.Background(), &core.DataUpHandlerReq{
		Payload:  req.Payload.MACPayload.FRMPayload,
		DevEUI:   mEntry.DevEUI,
		AppEUI:   mEntry.AppEUI,
		FCnt:     fhdr.FCnt,
		MType:    req.Payload.MHDR.MType,
		Metadata: req.Metadata,
	})

	if err != nil && !strings.Contains(err.Error(), string(errors.NotFound)) { // FIXME Find better way to analyze error
		stats.MarkMeter("broker.uplink.bad_handler_response")
		return nil, errors.New(errors.Operational, err)
	}
	stats.MarkMeter("broker.uplink.ok")

	// No response, we stop here and propagate the "no answer".
	// In case of confirmed data, the handler is in charge of creating the confirmation
	if resp == nil {
		return nil, nil
	}

	// If a response was sent, i.e. a downlink data, we need to compute the right MIC
	downlinkPayload, err := core.NewLoRaWANData(resp.Payload, false)
	if err != nil {
		return nil, errors.New(errors.Structural, err)
	}
	stats.MarkMeter("broker.downlink.in")
	if err := downlinkPayload.SetMIC(lorawan.AES128Key(mEntry.NwkSKey)); err != nil {
		return nil, errors.New(errors.Structural, "Unable to set response MIC")
	}
	resp.Payload.MIC = downlinkPayload.MIC[:]

	// And finally, we acknowledge the answer
	return &core.DataBrokerRes{
		Payload:  resp.Payload,
		Metadata: resp.Metadata,
	}, nil
}

// Register implements the core.BrokerServer interface
func (b component) SubscribePersonalized(bctx context.Context, req *core.SubPersoBrokerReq) (*core.SubPersoBrokerRes, error) {
	b.ctx.Debug("Handling personalized subscription")

	// Ensure the entry is valid
	if len(req.AppEUI) != 8 {
		return nil, errors.New(errors.Structural, "Invalid Application EUI")
	}

	if len(req.DevAddr) != 4 {
		return nil, errors.New(errors.Structural, "Invalid Device Address")
	}
	devEUI := make([]byte, 8, 8)
	copy(devEUI[4:], req.DevAddr)

	var nwkSKey [16]byte
	if len(req.NwkSKey) != 16 {
		return nil, errors.New(errors.Structural, "Invalid Network Session Key")
	}
	copy(nwkSKey[:], req.NwkSKey)

	re := regexp.MustCompile("^[-\\w]+\\.([-\\w]+\\.?)+:\\d+$")
	if !re.MatchString(req.HandlerNet) {
		return nil, errors.New(errors.Structural, fmt.Sprintf("Invalid Handler Net Address. Should match: %s", re))
	}

	return nil, b.StoreDevice(devEntry{
		HandlerNet: req.HandlerNet,
		AppEUI:     req.AppEUI,
		DevEUI:     devEUI,
		NwkSKey:    nwkSKey,
		FCntUp:     0,
	})
}
<|endoftext|>


<commit_before>package api

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"

	"github.com/FederationOfFathers/dashboard/bot"
)

func init() {
	Router.Path("/api/v0/slack/login").Methods("POST").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.ParseForm()
		user := r.Form.Get("user_id")
		if home := os.Getenv("SERVICE_DIR"); home == "" {
			bot.SendLogin(user)
		} else {
			bot.SendDevLogin(user)
		}
	})
	Router.Path("/api/v0/gh/ui-rebuild").Methods("POST").HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			var payload struct {
				Ref string `json:"ref"`
			}
			if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
				Logger.Warn(err.Error())
			}
			r.Body.Close()
			Logger.Warn(fmt.Sprintf("%#v", payload))
			fp, _ := os.OpenFile("queue-rebuild", os.O_RDONLY|os.O_CREATE, 0666)
			if fp != nil {
				fp.Close()
			}
		},
	)
}
<commit_msg>different file for dev and production<commit_after>package api

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"

	"github.com/FederationOfFathers/dashboard/bot"
)

func init() {
	Router.Path("/api/v0/slack/login").Methods("POST").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.ParseForm()
		user := r.Form.Get("user_id")
		if home := os.Getenv("SERVICE_DIR"); home == "" {
			bot.SendLogin(user)
		} else {
			bot.SendDevLogin(user)
		}
	})
	Router.Path("/api/v0/gh/ui-rebuild").Methods("POST").HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			var payload struct {
				Ref string `json:"ref"`
			}
			if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
				Logger.Warn(err.Error())
			}
			r.Body.Close()
			Logger.Warn(fmt.Sprintf("%#v", payload))
			touchFile := ""
			if payload.Ref == "refs/heads/dev" {
				touchFile = "queue-rebuild-dev"
			} else if payload.Ref == "refs/heads/master" {
				touchFile = "queue-rebuild-prod"
			}
			fp, _ := os.OpenFile(touchFile, os.O_RDONLY|os.O_CREATE, 0666)
			if fp != nil {
				fp.Close()
			}
		},
	)
}
<|endoftext|>


<commit_before>package locking

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/filepathfilter"
	"github.com/git-lfs/git-lfs/git"
	"github.com/git-lfs/git-lfs/tools"
)

// GetLockablePatterns returns a list of patterns in .gitattributes which are
// marked as lockable
func (c *Client) GetLockablePatterns() []string {
	c.ensureLockablesLoaded()
	return c.lockablePatterns
}

// getLockableFilter returns the internal filter used to check if a file is lockable
func (c *Client) getLockableFilter() *filepathfilter.Filter {
	c.ensureLockablesLoaded()
	return c.lockableFilter
}

func (c *Client) ensureLockablesLoaded() {
	c.lockableMutex.Lock()
	defer c.lockableMutex.Unlock()

	// Only load once
	if c.lockablePatterns == nil {
		c.refreshLockablePatterns()
	}
}

// Internal function to repopulate lockable patterns
// You must have locked the c.lockableMutex in the caller
func (c *Client) refreshLockablePatterns() {
	// Always make non-nil even if empty
	c.lockablePatterns = make([]string, 0, 10)

	paths := git.GetAttributePaths(c.LocalWorkingDir, c.LocalGitDir)
	for _, p := range paths {
		if p.Lockable {
			c.lockablePatterns = append(c.lockablePatterns, p.Path)
		}
	}
	c.lockableFilter = filepathfilter.New(c.lockablePatterns, nil)
}

// IsFileLockable returns whether a specific file path is marked as Lockable,
// ie has the 'lockable' attribute in .gitattributes
// Lockable patterns are cached once for performance, unless you call RefreshLockablePatterns
// path should be relative to repository root
func (c *Client) IsFileLockable(path string) bool {
	return c.getLockableFilter().Allows(path)
}

// FixAllLockableFileWriteFlags recursively scans the repo looking for files which
// are lockable, and makes sure their write flags are set correctly based on
// whether they are currently locked or unlocked.
// Files which are unlocked are made read-only, files which are locked are made
// writeable.
// This function can be used after a clone or checkout to ensure that file
// state correctly reflects the locking state
func (c *Client) FixAllLockableFileWriteFlags() error {
	return c.fixFileWriteFlags(c.LocalWorkingDir, c.LocalWorkingDir, c.getLockableFilter(), nil)
}

// FixFileWriteFlagsInDir scans dir (which can either be a relative dir
// from the root of the repo, or an absolute dir within the repo) looking for
// files to change permissions for.
// If lockablePatterns is non-nil, then any file matching those patterns will be
// checked to see if it is currently locked by the current committer, and if so
// it will be writeable, and if not locked it will be read-only.
// If unlockablePatterns is non-nil, then any file matching those patterns will
// be made writeable if it is not already. This can be used to reset files to
// writeable when their 'lockable' attribute is turned off.
func (c *Client) FixFileWriteFlagsInDir(dir string, lockablePatterns, unlockablePatterns []string) error {

	// early-out if no patterns
	if len(lockablePatterns) == 0 && len(unlockablePatterns) == 0 {
		return nil
	}

	absPath := dir
	if !filepath.IsAbs(dir) {
		absPath = filepath.Join(c.LocalWorkingDir, dir)
	}
	stat, err := os.Stat(absPath)
	if err != nil {
		return err
	}
	if !stat.IsDir() {
		return fmt.Errorf("%q is not a valid directory", dir)
	}

	var lockableFilter *filepathfilter.Filter
	var unlockableFilter *filepathfilter.Filter
	if lockablePatterns != nil {
		lockableFilter = filepathfilter.New(lockablePatterns, nil)
	}
	if unlockablePatterns != nil {
		unlockableFilter = filepathfilter.New(unlockablePatterns, nil)
	}

	return c.fixFileWriteFlags(absPath, c.LocalWorkingDir, lockableFilter, unlockableFilter)
}

// Internal implementation of fixing file write flags with precompiled filters
func (c *Client) fixFileWriteFlags(absPath, workingDir string, lockable, unlockable *filepathfilter.Filter) error {

	var errs []error
	var errMux sync.Mutex

	addErr := func(err error) {
		errMux.Lock()
		defer errMux.Unlock()

		errs = append(errs, err)
	}

	tools.FastWalkGitRepo(absPath, func(parentDir string, fi os.FileInfo, err error) {
		if err != nil {
			addErr(err)
			return
		}
		// Skip dirs, we only need to check files
		if fi.IsDir() {
			return
		}
		abschild := filepath.Join(parentDir, fi.Name())

		// This is a file, get relative to repo root
		relpath, err := filepath.Rel(workingDir, abschild)
		if err != nil {
			addErr(err)
			return
		}

		err = c.fixSingleFileWriteFlags(relpath, lockable, unlockable)
		if err != nil {
			addErr(err)
		}

	})
	return errors.Combine(errs)
}

// FixLockableFileWriteFlags checks each file in the provided list, and for
// those which are lockable, makes sure their write flags are set correctly
// based on whether they are currently locked or unlocked. Files which are
// unlocked are made read-only, files which are locked are made writeable.
// Files which are not lockable are ignored.
// This function can be used after a clone or checkout to ensure that file
// state correctly reflects the locking state, and is more efficient than
// FixAllLockableFileWriteFlags when you know which files changed
func (c *Client) FixLockableFileWriteFlags(files []string) error {
	// early-out if no lockable patterns
	if len(c.GetLockablePatterns()) == 0 {
		return nil
	}

	var errs []error
	for _, f := range files {
		err := c.fixSingleFileWriteFlags(f, c.getLockableFilter(), nil)
		if err != nil {
			errs = append(errs, err)
		}
	}

	return errors.Combine(errs)
}

// fixSingleFileWriteFlags fixes write flags on a single file
// If lockablePatterns is non-nil, then any file matching those patterns will be
// checked to see if it is currently locked by the current committer, and if so
// it will be writeable, and if not locked it will be read-only.
// If unlockablePatterns is non-nil, then any file matching those patterns will
// be made writeable if it is not already. This can be used to reset files to
// writeable when their 'lockable' attribute is turned off.
func (c *Client) fixSingleFileWriteFlags(file string, lockable, unlockable *filepathfilter.Filter) error {
	// Convert to git-style forward slash separators if necessary
	// Necessary to match attributes
	if filepath.Separator == '\\' {
		file = strings.Replace(file, "\\", "/", -1)
	}
	if lockable != nil && lockable.Allows(file) {
		// Lockable files are writeable only if they're currently locked
		err := tools.SetFileWriteFlag(file, c.IsFileLockedByCurrentCommitter(file))
		// Ignore not exist errors
		if err != nil && !os.IsNotExist(err) {
			return err
		}
	} else if unlockable != nil && unlockable.Allows(file) {
		// Unlockable files are always writeable
		// We only check files which match the incoming patterns to avoid
		// checking every file in the system all the time, and only do it
		// when a file has had its lockable attribute removed
		err := tools.SetFileWriteFlag(file, true)
		if err != nil && !os.IsNotExist(err) {
			return err
		}
	}
	return nil
}
<commit_msg>Better guide size for lockable patterns<commit_after>package locking

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/filepathfilter"
	"github.com/git-lfs/git-lfs/git"
	"github.com/git-lfs/git-lfs/tools"
)

// GetLockablePatterns returns a list of patterns in .gitattributes which are
// marked as lockable
func (c *Client) GetLockablePatterns() []string {
	c.ensureLockablesLoaded()
	return c.lockablePatterns
}

// getLockableFilter returns the internal filter used to check if a file is lockable
func (c *Client) getLockableFilter() *filepathfilter.Filter {
	c.ensureLockablesLoaded()
	return c.lockableFilter
}

func (c *Client) ensureLockablesLoaded() {
	c.lockableMutex.Lock()
	defer c.lockableMutex.Unlock()

	// Only load once
	if c.lockablePatterns == nil {
		c.refreshLockablePatterns()
	}
}

// Internal function to repopulate lockable patterns
// You must have locked the c.lockableMutex in the caller
func (c *Client) refreshLockablePatterns() {

	paths := git.GetAttributePaths(c.LocalWorkingDir, c.LocalGitDir)
	// Always make non-nil even if empty
	c.lockablePatterns = make([]string, 0, len(paths))
	for _, p := range paths {
		if p.Lockable {
			c.lockablePatterns = append(c.lockablePatterns, p.Path)
		}
	}
	c.lockableFilter = filepathfilter.New(c.lockablePatterns, nil)
}

// IsFileLockable returns whether a specific file path is marked as Lockable,
// ie has the 'lockable' attribute in .gitattributes
// Lockable patterns are cached once for performance, unless you call RefreshLockablePatterns
// path should be relative to repository root
func (c *Client) IsFileLockable(path string) bool {
	return c.getLockableFilter().Allows(path)
}

// FixAllLockableFileWriteFlags recursively scans the repo looking for files which
// are lockable, and makes sure their write flags are set correctly based on
// whether they are currently locked or unlocked.
// Files which are unlocked are made read-only, files which are locked are made
// writeable.
// This function can be used after a clone or checkout to ensure that file
// state correctly reflects the locking state
func (c *Client) FixAllLockableFileWriteFlags() error {
	return c.fixFileWriteFlags(c.LocalWorkingDir, c.LocalWorkingDir, c.getLockableFilter(), nil)
}

// FixFileWriteFlagsInDir scans dir (which can either be a relative dir
// from the root of the repo, or an absolute dir within the repo) looking for
// files to change permissions for.
// If lockablePatterns is non-nil, then any file matching those patterns will be
// checked to see if it is currently locked by the current committer, and if so
// it will be writeable, and if not locked it will be read-only.
// If unlockablePatterns is non-nil, then any file matching those patterns will
// be made writeable if it is not already. This can be used to reset files to
// writeable when their 'lockable' attribute is turned off.
func (c *Client) FixFileWriteFlagsInDir(dir string, lockablePatterns, unlockablePatterns []string) error {

	// early-out if no patterns
	if len(lockablePatterns) == 0 && len(unlockablePatterns) == 0 {
		return nil
	}

	absPath := dir
	if !filepath.IsAbs(dir) {
		absPath = filepath.Join(c.LocalWorkingDir, dir)
	}
	stat, err := os.Stat(absPath)
	if err != nil {
		return err
	}
	if !stat.IsDir() {
		return fmt.Errorf("%q is not a valid directory", dir)
	}

	var lockableFilter *filepathfilter.Filter
	var unlockableFilter *filepathfilter.Filter
	if lockablePatterns != nil {
		lockableFilter = filepathfilter.New(lockablePatterns, nil)
	}
	if unlockablePatterns != nil {
		unlockableFilter = filepathfilter.New(unlockablePatterns, nil)
	}

	return c.fixFileWriteFlags(absPath, c.LocalWorkingDir, lockableFilter, unlockableFilter)
}

// Internal implementation of fixing file write flags with precompiled filters
func (c *Client) fixFileWriteFlags(absPath, workingDir string, lockable, unlockable *filepathfilter.Filter) error {

	var errs []error
	var errMux sync.Mutex

	addErr := func(err error) {
		errMux.Lock()
		defer errMux.Unlock()

		errs = append(errs, err)
	}

	tools.FastWalkGitRepo(absPath, func(parentDir string, fi os.FileInfo, err error) {
		if err != nil {
			addErr(err)
			return
		}
		// Skip dirs, we only need to check files
		if fi.IsDir() {
			return
		}
		abschild := filepath.Join(parentDir, fi.Name())

		// This is a file, get relative to repo root
		relpath, err := filepath.Rel(workingDir, abschild)
		if err != nil {
			addErr(err)
			return
		}

		err = c.fixSingleFileWriteFlags(relpath, lockable, unlockable)
		if err != nil {
			addErr(err)
		}

	})
	return errors.Combine(errs)
}

// FixLockableFileWriteFlags checks each file in the provided list, and for
// those which are lockable, makes sure their write flags are set correctly
// based on whether they are currently locked or unlocked. Files which are
// unlocked are made read-only, files which are locked are made writeable.
// Files which are not lockable are ignored.
// This function can be used after a clone or checkout to ensure that file
// state correctly reflects the locking state, and is more efficient than
// FixAllLockableFileWriteFlags when you know which files changed
func (c *Client) FixLockableFileWriteFlags(files []string) error {
	// early-out if no lockable patterns
	if len(c.GetLockablePatterns()) == 0 {
		return nil
	}

	var errs []error
	for _, f := range files {
		err := c.fixSingleFileWriteFlags(f, c.getLockableFilter(), nil)
		if err != nil {
			errs = append(errs, err)
		}
	}

	return errors.Combine(errs)
}

// fixSingleFileWriteFlags fixes write flags on a single file
// If lockablePatterns is non-nil, then any file matching those patterns will be
// checked to see if it is currently locked by the current committer, and if so
// it will be writeable, and if not locked it will be read-only.
// If unlockablePatterns is non-nil, then any file matching those patterns will
// be made writeable if it is not already. This can be used to reset files to
// writeable when their 'lockable' attribute is turned off.
func (c *Client) fixSingleFileWriteFlags(file string, lockable, unlockable *filepathfilter.Filter) error {
	// Convert to git-style forward slash separators if necessary
	// Necessary to match attributes
	if filepath.Separator == '\\' {
		file = strings.Replace(file, "\\", "/", -1)
	}
	if lockable != nil && lockable.Allows(file) {
		// Lockable files are writeable only if they're currently locked
		err := tools.SetFileWriteFlag(file, c.IsFileLockedByCurrentCommitter(file))
		// Ignore not exist errors
		if err != nil && !os.IsNotExist(err) {
			return err
		}
	} else if unlockable != nil && unlockable.Allows(file) {
		// Unlockable files are always writeable
		// We only check files which match the incoming patterns to avoid
		// checking every file in the system all the time, and only do it
		// when a file has had its lockable attribute removed
		err := tools.SetFileWriteFlag(file, true)
		if err != nil && !os.IsNotExist(err) {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package logan_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
)

func TestLogan(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Logan Suite")
}
<commit_msg>Deleted : logan_suite_test.go file<commit_after><|endoftext|>


<commit_before>package twofa

import (
	"bytes"
	"crypto"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"strings"

	"github.com/Cloud-Foundations/golib/pkg/log"
	"github.com/Cloud-Foundations/keymaster/lib/client/twofa/pushtoken"
	"github.com/Cloud-Foundations/keymaster/lib/client/twofa/totp"
	"github.com/Cloud-Foundations/keymaster/lib/client/twofa/u2f"
	"github.com/Cloud-Foundations/keymaster/lib/webapi/v0/proto"
	"github.com/flynn/u2f/u2fhid" // client side (interface with hardware)
	"golang.org/x/crypto/ssh"
)

const clientDataAuthenticationTypeValue = "navigator.id.getAssertion"

// This is now copy-paste from the server test side... probably make public and reuse.
func createKeyBodyRequest(method, urlStr, filedata string) (*http.Request, error) {
	//create attachment....
	bodyBuf := &bytes.Buffer{}
	bodyWriter := multipart.NewWriter(bodyBuf)

	//
	fileWriter, err := bodyWriter.CreateFormFile("pubkeyfile", "somefilename.pub")
	if err != nil {
		fmt.Println("error writing to buffer")
		return nil, err
	}
	// When using a file this used to be: fh, err := os.Open(pubKeyFilename)
	fh := strings.NewReader(filedata)

	_, err = io.Copy(fileWriter, fh)
	if err != nil {
		return nil, err
	}

	err = bodyWriter.WriteField("duration", (*Duration).String())
	if err != nil {
		return nil, err
	}

	contentType := bodyWriter.FormDataContentType()
	bodyWriter.Close()

	req, err := http.NewRequest(method, urlStr, bodyBuf)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", contentType)

	return req, nil
}

func doCertRequest(client *http.Client,
	url, serializedPubkey string,
	certType string,
	addGroups bool,
	userAgentString string, logger log.Logger) ([]byte, error) {

	switch certType {
	case "x509-kubernetes", "x509":
		//logger.Debugf(1, "requesting x509 cert family, type=%s", certType)
	case "ssh":
		//logger.Debugf(1, "requesting ssh cert family, type=%s", certType)
	default:
		return nil, fmt.Errorf("invalid certType requested '%s'", certType)

	}

	return nil, fmt.Errorf("Not implemented")
}

func doCertRequestInternal(client *http.Client,
	url, filedata string,
	userAgentString string, logger log.Logger) ([]byte, error) {

	req, err := createKeyBodyRequest("POST", url, filedata)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", userAgentString)
	resp, err := client.Do(req) // Client.Get(targetUrl)
	if err != nil {
		logger.Printf("Failure to do cert request %s", err)
		return nil, err
	}

	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("got error from call %s, url='%s'\n", resp.Status, url)
	}
	return ioutil.ReadAll(resp.Body)

}

// This assumes the http client has a non-nul cookie jar
func authenticateUser(
	userName string,
	password []byte,
	baseUrl string,
	skip2fa bool,
	client *http.Client,
	userAgentString string,
	logger log.DebugLogger) (err error) {

	loginUrl := baseUrl + proto.LoginPath
	form := url.Values{}
	form.Add("username", userName)
	form.Add("password", string(password[:]))
	req, err := http.NewRequest("POST", loginUrl,
		strings.NewReader(form.Encode()))
	if err != nil {
		return err
	}
	req.Header.Add("Content-Length", strconv.Itoa(len(form.Encode())))
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Add("Accept", "application/json")
	req.Header.Set("User-Agent", userAgentString)

	logger.Debugf(1, "About to start login request\n")
	loginResp, err := client.Do(req) //client.Get(targetUrl)
	if err != nil {
		logger.Printf("got error from req")
		logger.Println(err)
		// TODO: differentiate between 400 and 500 errors
		// is OK to fail.. try next
		return err
	}
	defer loginResp.Body.Close()
	if loginResp.StatusCode != 200 {
		if loginResp.StatusCode == http.StatusUnauthorized {
			return fmt.Errorf("Unauthorized reponse from server. Check username and/or password")
		}
		logger.Debugf(1, "got error from login call %s", loginResp.Status)
		return fmt.Errorf("got error from login call %s", loginResp.Status)
	}
	//Enusre we have at least one cookie
	if len(loginResp.Cookies()) < 1 {
		err = errors.New("No cookies from login")
		return err
	}

	loginJSONResponse := proto.LoginResponse{}
	//body := jsonrr.Result().Body
	err = json.NewDecoder(loginResp.Body).Decode(&loginJSONResponse)
	if err != nil {
		return err
	}
	io.Copy(ioutil.Discard, loginResp.Body) // We also need to read ALL of the body
	loginResp.Body.Close()                  //so that we can reuse the channel
	logger.Debugf(1, "This the login response=%v\n", loginJSONResponse)

	allowVIP := false
	allowU2F := false
	allowTOTP := false
	allowOkta2FA := false
	for _, backend := range loginJSONResponse.CertAuthBackend {
		if backend == proto.AuthTypePassword {
			skip2fa = true
		}
		if backend == proto.AuthTypeSymantecVIP {
			allowVIP = true
			//remote next statemente later
			//skipu2f = true
		}
		if backend == proto.AuthTypeU2F {
			allowU2F = true
		}
		if backend == proto.AuthTypeTOTP {
			allowTOTP = true
		}
		if backend == proto.AuthTypeOkta2FA {
			allowOkta2FA = true
		}
	}

	// Dont try U2F if chosen by user
	if *noU2F {
		allowU2F = false
	}
	if *noTOTP {
		allowTOTP = false
	}
	if *noVIPAccess {
		allowVIP = false
	}

	// on linux disable U2F is the /sys/class/hidraw is missing
	if runtime.GOOS == "linux" && allowU2F {
		if _, err := os.Stat("/sys/class/hidraw"); os.IsNotExist(err) {
			allowU2F = false
		}

	}
	// upgrade to u2f
	successful2fa := false
	if !skip2fa {
		if allowU2F {
			devices, err := u2fhid.Devices()
			if err != nil {
				logger.Fatal(err)
				return err
			}
			if len(devices) > 0 {

				err = u2f.DoU2FAuthenticate(
					client, baseUrl, userAgentString, logger)
				if err != nil {

					return err
				}
				successful2fa = true
			}
		}

		if allowTOTP && !successful2fa {
			err = totp.DoTOTPAuthenticate(
				client, baseUrl, userAgentString, logger)
			if err != nil {

				return err
			}
			successful2fa = true
		}
		if allowVIP && !successful2fa {
			err = pushtoken.DoVIPAuthenticate(
				client, baseUrl, userAgentString, logger)
			if err != nil {

				return err
			}
			successful2fa = true
		}
		// TODO: do better logic when both VIP and OKTA are configured
		if allowOkta2FA && !successful2fa {
			err = pushtoken.DoOktaAuthenticate(
				client, baseUrl, userAgentString, logger)
			if err != nil {
				return err
			}
			successful2fa = true
		}

		if !successful2fa {
			err = errors.New("Failed to Pefrom 2FA (as requested from server)")
			return err
		}

	}

	logger.Debugf(1, "Authentication Phase complete")

	return nil

}

func getCertsFromServer(
	signer crypto.Signer,
	userName string,
	password []byte,
	baseUrl string,
	skip2fa bool,
	addGroups bool,
	client *http.Client,
	userAgentString string,
	logger log.DebugLogger) (sshCert []byte, x509Cert []byte, kubernetesCert []byte, err error) {

	//This requires an http client with valid cookie jar

	err = authenticateUser(
		userName,
		password,
		baseUrl,
		skip2fa,
		client,
		userAgentString,
		logger)
	if err != nil {
		return nil, nil, nil, err
	}

	//now get x509 cert
	pubKey := signer.Public()
	derKey, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		return nil, nil, nil, err
	}
	pemKey := string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: derKey}))

	var urlPostfix string
	if addGroups {
		urlPostfix = "&addGroups=true"
		logger.Debugln(0, "adding \"addGroups\" to request")
	}
	// TODO: urlencode the userName
	x509Cert, err = doCertRequestInternal(
		client,
		baseUrl+"/certgen/"+userName+"?type=x509"+urlPostfix,
		pemKey,
		userAgentString,
		logger)
	if err != nil {
		return nil, nil, nil, err
	}

	kubernetesCert, err = doCertRequestInternal(
		client,
		baseUrl+"/certgen/"+userName+"?type=x509-kubernetes",
		pemKey,
		userAgentString,
		logger)
	if err != nil {
		//logger.Printf("Warning: could not get the kubernets cert (old server?) err=%s \n", err)
		kubernetesCert = nil
		//return nil, nil, nil, err
	}

	//// Now we do sshCert!
	// generate and write public key
	sshPub, err := ssh.NewPublicKey(pubKey)
	if err != nil {
		return nil, nil, nil, err
	}
	sshAuthFile := string(ssh.MarshalAuthorizedKey(sshPub))
	sshCert, err = doCertRequestInternal(
		client,
		baseUrl+"/certgen/"+userName+"?type=ssh",
		sshAuthFile,
		userAgentString,
		logger)
	if err != nil {
		return nil, nil, nil, err
	}

	return sshCert, x509Cert, kubernetesCert, nil
}

func getCertFromTargetUrls(
	signer crypto.Signer,
	userName string,
	password []byte,
	targetUrls []string,
	skipu2f bool,
	addGroups bool,
	client *http.Client,
	userAgentString string,
	logger log.DebugLogger) (sshCert []byte, x509Cert []byte, kubernetesCert []byte, err error) {
	success := false

	for _, baseUrl := range targetUrls {
		logger.Printf("attempting to target '%s' for '%s'\n", baseUrl, userName)
		sshCert, x509Cert, kubernetesCert, err = getCertsFromServer(
			signer, userName, password, baseUrl, skipu2f, addGroups,
			client, userAgentString, logger)
		if err != nil {
			logger.Println(err)
			continue
		}
		success = true
		break

	}
	if !success {
		err := errors.New("Failed to get creds")
		return nil, nil, nil, err
	}

	return sshCert, x509Cert, kubernetesCert, nil
}

func authenticateToTargetUrls(
	userName string,
	password []byte,
	targetUrls []string,
	skip2fa bool,
	client *http.Client,
	userAgentString string,
	logger log.DebugLogger) (baseUrl string, err error) {

	for _, baseUrl = range targetUrls {
		logger.Printf("attempting to target '%s' for '%s'\n", baseUrl, userName)
		err = authenticateUser(
			userName,
			password,
			baseUrl,
			skip2fa,
			client,
			userAgentString,
			logger)
		if err != nil {
			continue
		}
		return baseUrl, nil

	}
	return "", fmt.Errorf("Failed to Authenticate to any URL")
}
<commit_msg>cleaner separation<commit_after>package twofa

import (
	"bytes"
	"crypto"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"strings"

	"github.com/Cloud-Foundations/golib/pkg/log"
	"github.com/Cloud-Foundations/keymaster/lib/client/twofa/pushtoken"
	"github.com/Cloud-Foundations/keymaster/lib/client/twofa/totp"
	"github.com/Cloud-Foundations/keymaster/lib/client/twofa/u2f"
	"github.com/Cloud-Foundations/keymaster/lib/webapi/v0/proto"
	"github.com/flynn/u2f/u2fhid" // client side (interface with hardware)
	"golang.org/x/crypto/ssh"
)

const clientDataAuthenticationTypeValue = "navigator.id.getAssertion"

// This is now copy-paste from the server test side... probably make public and reuse.
func createKeyBodyRequest(method, urlStr, filedata string) (*http.Request, error) {
	//create attachment....
	bodyBuf := &bytes.Buffer{}
	bodyWriter := multipart.NewWriter(bodyBuf)

	//
	fileWriter, err := bodyWriter.CreateFormFile("pubkeyfile", "somefilename.pub")
	if err != nil {
		fmt.Println("error writing to buffer")
		return nil, err
	}
	// When using a file this used to be: fh, err := os.Open(pubKeyFilename)
	fh := strings.NewReader(filedata)

	_, err = io.Copy(fileWriter, fh)
	if err != nil {
		return nil, err
	}

	err = bodyWriter.WriteField("duration", (*Duration).String())
	if err != nil {
		return nil, err
	}

	contentType := bodyWriter.FormDataContentType()
	bodyWriter.Close()

	req, err := http.NewRequest(method, urlStr, bodyBuf)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", contentType)

	return req, nil
}

func doCertRequest(signer crypto.Signer, client *http.Client, userName string,
	baseUrl,
	certType string,
	addGroups bool,
	userAgentString string, logger log.DebugLogger) ([]byte, error) {

	pubKey := signer.Public()
	var serializedPubkey string
	switch certType {
	case "x509", "x509-kubernetes":
		derKey, err := x509.MarshalPKIXPublicKey(pubKey)
		if err != nil {
			return nil, err
		}
		serializedPubkey = string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: derKey}))
	case "ssh":
		sshPub, err := ssh.NewPublicKey(pubKey)
		if err != nil {
			return nil, err
		}
		serializedPubkey = string(ssh.MarshalAuthorizedKey(sshPub))
	default:
		return nil, fmt.Errorf("invalid certType requested '%s'", certType)

	}

	var urlPostfix string
	// addgroups only makes sense for x509 plain .. maybe set as a check insetad of dropping?
	if certType == "x509" && addGroups {
		urlPostfix = "&addGroups=true"
		logger.Debugln(0, "adding \"addGroups\" to request")
	}

	requestURL := baseUrl + "/certgen/" + userName + "?type=" + certType + urlPostfix

	return doCertRequestInternal(client, requestURL, serializedPubkey, userAgentString, logger)
}

func doCertRequestInternal(client *http.Client,
	url, filedata string,
	userAgentString string, logger log.Logger) ([]byte, error) {

	req, err := createKeyBodyRequest("POST", url, filedata)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", userAgentString)
	resp, err := client.Do(req) // Client.Get(targetUrl)
	if err != nil {
		logger.Printf("Failure to do cert request %s", err)
		return nil, err
	}

	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("got error from call %s, url='%s'\n", resp.Status, url)
	}
	return ioutil.ReadAll(resp.Body)

}

// This assumes the http client has a non-nul cookie jar
func authenticateUser(
	userName string,
	password []byte,
	baseUrl string,
	skip2fa bool,
	client *http.Client,
	userAgentString string,
	logger log.DebugLogger) (err error) {

	loginUrl := baseUrl + proto.LoginPath
	form := url.Values{}
	form.Add("username", userName)
	form.Add("password", string(password[:]))
	req, err := http.NewRequest("POST", loginUrl,
		strings.NewReader(form.Encode()))
	if err != nil {
		return err
	}
	req.Header.Add("Content-Length", strconv.Itoa(len(form.Encode())))
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Add("Accept", "application/json")
	req.Header.Set("User-Agent", userAgentString)

	logger.Debugf(1, "About to start login request\n")
	loginResp, err := client.Do(req) //client.Get(targetUrl)
	if err != nil {
		logger.Printf("got error from req")
		logger.Println(err)
		// TODO: differentiate between 400 and 500 errors
		// is OK to fail.. try next
		return err
	}
	defer loginResp.Body.Close()
	if loginResp.StatusCode != 200 {
		if loginResp.StatusCode == http.StatusUnauthorized {
			return fmt.Errorf("Unauthorized reponse from server. Check username and/or password")
		}
		logger.Debugf(1, "got error from login call %s", loginResp.Status)
		return fmt.Errorf("got error from login call %s", loginResp.Status)
	}
	//Enusre we have at least one cookie
	if len(loginResp.Cookies()) < 1 {
		err = errors.New("No cookies from login")
		return err
	}

	loginJSONResponse := proto.LoginResponse{}
	//body := jsonrr.Result().Body
	err = json.NewDecoder(loginResp.Body).Decode(&loginJSONResponse)
	if err != nil {
		return err
	}
	io.Copy(ioutil.Discard, loginResp.Body) // We also need to read ALL of the body
	loginResp.Body.Close()                  //so that we can reuse the channel
	logger.Debugf(1, "This the login response=%v\n", loginJSONResponse)

	allowVIP := false
	allowU2F := false
	allowTOTP := false
	allowOkta2FA := false
	for _, backend := range loginJSONResponse.CertAuthBackend {
		if backend == proto.AuthTypePassword {
			skip2fa = true
		}
		if backend == proto.AuthTypeSymantecVIP {
			allowVIP = true
			//remote next statemente later
			//skipu2f = true
		}
		if backend == proto.AuthTypeU2F {
			allowU2F = true
		}
		if backend == proto.AuthTypeTOTP {
			allowTOTP = true
		}
		if backend == proto.AuthTypeOkta2FA {
			allowOkta2FA = true
		}
	}

	// Dont try U2F if chosen by user
	if *noU2F {
		allowU2F = false
	}
	if *noTOTP {
		allowTOTP = false
	}
	if *noVIPAccess {
		allowVIP = false
	}

	// on linux disable U2F is the /sys/class/hidraw is missing
	if runtime.GOOS == "linux" && allowU2F {
		if _, err := os.Stat("/sys/class/hidraw"); os.IsNotExist(err) {
			allowU2F = false
		}

	}
	// upgrade to u2f
	successful2fa := false
	if !skip2fa {
		if allowU2F {
			devices, err := u2fhid.Devices()
			if err != nil {
				logger.Fatal(err)
				return err
			}
			if len(devices) > 0 {

				err = u2f.DoU2FAuthenticate(
					client, baseUrl, userAgentString, logger)
				if err != nil {

					return err
				}
				successful2fa = true
			}
		}

		if allowTOTP && !successful2fa {
			err = totp.DoTOTPAuthenticate(
				client, baseUrl, userAgentString, logger)
			if err != nil {

				return err
			}
			successful2fa = true
		}
		if allowVIP && !successful2fa {
			err = pushtoken.DoVIPAuthenticate(
				client, baseUrl, userAgentString, logger)
			if err != nil {

				return err
			}
			successful2fa = true
		}
		// TODO: do better logic when both VIP and OKTA are configured
		if allowOkta2FA && !successful2fa {
			err = pushtoken.DoOktaAuthenticate(
				client, baseUrl, userAgentString, logger)
			if err != nil {
				return err
			}
			successful2fa = true
		}

		if !successful2fa {
			err = errors.New("Failed to Pefrom 2FA (as requested from server)")
			return err
		}

	}

	logger.Debugf(1, "Authentication Phase complete")

	return nil

}

func getCertsFromServer(
	signer crypto.Signer,
	userName string,
	password []byte,
	baseUrl string,
	skip2fa bool,
	addGroups bool,
	client *http.Client,
	userAgentString string,
	logger log.DebugLogger) (sshCert []byte, x509Cert []byte, kubernetesCert []byte, err error) {

	//This requires an http client with valid cookie jar

	err = authenticateUser(
		userName,
		password,
		baseUrl,
		skip2fa,
		client,
		userAgentString,
		logger)
	if err != nil {
		return nil, nil, nil, err
	}

	//now get x509 cert
	pubKey := signer.Public()
	derKey, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		return nil, nil, nil, err
	}
	pemKey := string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: derKey}))

	var urlPostfix string
	if addGroups {
		urlPostfix = "&addGroups=true"
		logger.Debugln(0, "adding \"addGroups\" to request")
	}
	// TODO: urlencode the userName
	x509Cert, err = doCertRequestInternal(
		client,
		baseUrl+"/certgen/"+userName+"?type=x509"+urlPostfix,
		pemKey,
		userAgentString,
		logger)
	if err != nil {
		return nil, nil, nil, err
	}

	kubernetesCert, err = doCertRequestInternal(
		client,
		baseUrl+"/certgen/"+userName+"?type=x509-kubernetes",
		pemKey,
		userAgentString,
		logger)
	if err != nil {
		//logger.Printf("Warning: could not get the kubernets cert (old server?) err=%s \n", err)
		kubernetesCert = nil
		//return nil, nil, nil, err
	}

	//// Now we do sshCert!
	// generate and write public key
	sshPub, err := ssh.NewPublicKey(pubKey)
	if err != nil {
		return nil, nil, nil, err
	}
	sshAuthFile := string(ssh.MarshalAuthorizedKey(sshPub))
	sshCert, err = doCertRequestInternal(
		client,
		baseUrl+"/certgen/"+userName+"?type=ssh",
		sshAuthFile,
		userAgentString,
		logger)
	if err != nil {
		return nil, nil, nil, err
	}

	return sshCert, x509Cert, kubernetesCert, nil
}

func getCertFromTargetUrls(
	signer crypto.Signer,
	userName string,
	password []byte,
	targetUrls []string,
	skipu2f bool,
	addGroups bool,
	client *http.Client,
	userAgentString string,
	logger log.DebugLogger) (sshCert []byte, x509Cert []byte, kubernetesCert []byte, err error) {
	success := false

	for _, baseUrl := range targetUrls {
		logger.Printf("attempting to target '%s' for '%s'\n", baseUrl, userName)
		sshCert, x509Cert, kubernetesCert, err = getCertsFromServer(
			signer, userName, password, baseUrl, skipu2f, addGroups,
			client, userAgentString, logger)
		if err != nil {
			logger.Println(err)
			continue
		}
		success = true
		break

	}
	if !success {
		err := errors.New("Failed to get creds")
		return nil, nil, nil, err
	}

	return sshCert, x509Cert, kubernetesCert, nil
}

func authenticateToTargetUrls(
	userName string,
	password []byte,
	targetUrls []string,
	skip2fa bool,
	client *http.Client,
	userAgentString string,
	logger log.DebugLogger) (baseUrl string, err error) {

	for _, baseUrl = range targetUrls {
		logger.Printf("attempting to target '%s' for '%s'\n", baseUrl, userName)
		err = authenticateUser(
			userName,
			password,
			baseUrl,
			skip2fa,
			client,
			userAgentString,
			logger)
		if err != nil {
			continue
		}
		return baseUrl, nil

	}
	return "", fmt.Errorf("Failed to Authenticate to any URL")
}
<|endoftext|>


<commit_before>package plugins

import (
	"errors"
	"fmt"
	"sync"

	schematypes "github.com/taskcluster/go-schematypes"
	"github.com/taskcluster/taskcluster-worker/engines"
	"github.com/taskcluster/taskcluster-worker/runtime"
	"github.com/taskcluster/taskcluster-worker/runtime/atomics"
)

type pluginManager struct {
	payloadSchema schematypes.Object
	plugins       []Plugin
	pluginNames   []string
}

type taskPluginManager struct {
	taskPlugins []TaskPlugin
	working     atomics.Bool
}

// mergeErrors merges a list of errors into one error, ignoring nil entries,
// and ensuring that if there is only MalformedPayloadErrors these will be
// merged into one MalformedPayloadError
func mergeErrors(errs ...error) error {
	msg := ""
	for _, err := range errs {
		if err != nil {
			// TODO MalformedPayloadError needs special care, if we only have these
			//      it is very different from having a list of other errors, Hence,
			//      we need to maintain the type rather than mindlessly wrap them.
			msg += err.Error() + "\n"
		}
	}
	if len(msg) > 0 {
		return errors.New(msg)
	}
	return nil
}

// waitForErrors returns when it has received count results from the errors
// channel.
//
// Note, that errors might be nil, if all are nil it'll return nil otherwise
// it'll merge the errors.
func waitForErrors(errors <-chan error, count int) error {
	if count == 0 {
		return nil
	}
	errs := []error{}
	for err := range errors {
		if err != nil {
			errs = append(errs, err)
		}
		count--
		if count == 0 {
			break
		}
	}
	return mergeErrors(errs...)
}

// PluginManagerConfigSchema returns configuration for PluginOptions.Config for
// NewPluginManager.
func PluginManagerConfigSchema() schematypes.Object {
	plugins := Plugins()

	pluginNames := []string{}
	for name := range plugins {
		pluginNames = append(pluginNames, name)
	}

	s := schematypes.Object{
		MetaData: schematypes.MetaData{
			Title: "Plugin Configuration",
			Description: `Mapping from plugin name to plugin configuration.
                    The 'disabled' key is special and lists plugins that are
                    disabled. Plugins that are disabled do not require
                    configuration.`,
		},
		Properties: schematypes.Properties{
			"disabled": schematypes.Array{
				MetaData: schematypes.MetaData{
					Title: "Disabled Plugins",
					Description: `List of disabled plugins. If a plugin is not listed
												as disabled here, then its configuration key must be
												specified, unless it doesn't take any configuration.`,
				},
				Items: schematypes.StringEnum{
					Options: pluginNames,
				},
			},
		},
		Required: []string{"disabled"},
	}
	for name, provider := range plugins {
		cs := provider.ConfigSchema()
		if cs != nil {
			s.Properties[name] = cs
		}
	}
	return s
}

// stringContains returns true if list contains element
func stringContains(list []string, element string) bool {
	for _, s := range list {
		if s == element {
			return true
		}
	}
	return false
}

// NewPluginManager loads all plugins not disabled in configuration and
// returns a Plugin implementation that wraps all of the plugins.
//
// This expects options.Config satisfying schema from
// PluginManagerConfigSchema().
func NewPluginManager(options PluginOptions) (Plugin, error) {
	pluginProviders := Plugins()

	// Construct config schema
	configSchema := PluginManagerConfigSchema()

	// Ensure the config is valid
	if err := configSchema.Validate(options.Config); err != nil {
		return nil, fmt.Errorf("Invalid config, error: %s", err)
	}
	config := options.Config.(map[string]interface{})

	// Find plugins to load
	var enabled []string
	var disabled []string
	if configSchema.Properties["disabled"].Map(config["disabled"], &disabled) != nil {
		panic("internal error -- shouldn't be possible")
	}

	// Find list of enabled plugins and ensure that config is present if required
	for name, plugin := range pluginProviders {
		// Skip disabled plugins
		if stringContains(disabled, name) {
			continue
		}
		// Check that configuration is given if required
		if plugin.ConfigSchema() != nil {
			if _, ok := config[name]; !ok {
				return nil, fmt.Errorf("Missing configuration for plugin: '%s'", name)
			}
		}
		// List plugin as enabled
		enabled = append(enabled, name)
	}

	// Initialize all the plugins
	wg := sync.WaitGroup{}
	plugins := make([]Plugin, len(enabled))
	errors := make([]error, len(enabled))
	wg.Add(len(enabled))
	for index, name := range enabled {
		go func(index int, name string) {
			plugins[index], errors[index] = pluginProviders[name].NewPlugin(PluginOptions{
				Environment: options.Environment,
				Engine:      options.Engine,
				Log:         options.Log.WithField("plugin", name),
				Config:      config[name],
			})
			wg.Done()
		}(index, name) // needed to capture values not variables
	}
	wg.Wait()

	// Return the first error, if any
	for _, e := range errors {
		if e != nil {
			return nil, e
		}
	}

	// Construct payload schema
	schemas := []schematypes.Object{}
	for _, plugin := range plugins {
		schemas = append(schemas, plugin.PayloadSchema())
	}
	schema, err := schematypes.Merge(schemas...)
	if err != nil {
		return nil, fmt.Errorf("Conflicting payload schema types, error: %s", err)
	}

	return &pluginManager{
		plugins:       plugins,
		pluginNames:   enabled,
		payloadSchema: schema,
	}, nil
}

func (m *pluginManager) PayloadSchema() schematypes.Object {
	return m.payloadSchema
}

func (m *pluginManager) NewTaskPlugin(options TaskPluginOptions) (TaskPlugin, error) {
	// Input must be valid
	if m.payloadSchema.Validate(options.Payload) != nil {
		return nil, engines.ErrContractViolation
	}

	// Create a list of TaskPlugins and a mutex to guard access
	taskPlugins := []TaskPlugin{}
	mu := sync.Mutex{}

	errors := make(chan error)
	for i, p := range m.plugins {
		go func(name string, p Plugin) {
			taskPlugin, err := p.NewTaskPlugin(TaskPluginOptions{
				TaskInfo: options.TaskInfo,
				Payload:  p.PayloadSchema().Filter(options.Payload),
				Log:      options.Log.WithField("plugin", name),
			})
			if taskPlugin != nil {
				mu.Lock()
				taskPlugins = append(taskPlugins, taskPlugin)
				mu.Unlock()
			}
			errors <- err
		}(m.pluginNames[i], p)
	}

	// Wait for errors, if any we dispose and return a merged error
	err := waitForErrors(errors, len(m.plugins))
	manager := &taskPluginManager{taskPlugins: taskPlugins}
	if err != nil {
		err2 := manager.Dispose()
		return nil, mergeErrors(err, err2)
	}
	return manager, nil
}

func (m *taskPluginManager) Prepare(c *runtime.TaskContext) error {
	// Sanity check that no two methods on plugin is running in parallel, this way
	// plugins don't have to be thread-safe, and we ensure nothing is called after
	// Dispose() has been called.
	if m.working.Swap(true) {
		panic("Another plugin method is currently running, or Dispose() has been called!")
	}
	defer m.working.Set(false)

	// Run method on plugins in parallel
	errors := make(chan error)
	for _, p := range m.taskPlugins {
		go func(p TaskPlugin) { errors <- p.Prepare(c) }(p)
	}
	return waitForErrors(errors, len(m.taskPlugins))
}

func (m *taskPluginManager) BuildSandbox(b engines.SandboxBuilder) error {
	// Sanity check that no two methods on plugin is running in parallel, this way
	// plugins don't have to be thread-safe, and we ensure nothing is called after
	// Dispose() has been called.
	if m.working.Swap(true) {
		panic("Another plugin method is currently running, or Dispose() has been called!")
	}
	defer m.working.Set(false)

	// Run method on plugins in parallel
	errors := make(chan error)
	for _, p := range m.taskPlugins {
		go func(p TaskPlugin) { errors <- p.BuildSandbox(b) }(p)
	}
	return waitForErrors(errors, len(m.taskPlugins))
}

func (m *taskPluginManager) Started(s engines.Sandbox) error {
	// Sanity check that no two methods on plugin is running in parallel, this way
	// plugins don't have to be thread-safe, and we ensure nothing is called after
	// Dispose() has been called.
	if m.working.Swap(true) {
		panic("Another plugin method is currently running, or Dispose() has been called!")
	}
	defer m.working.Set(false)

	// Run method on plugins in parallel
	errors := make(chan error)
	for _, p := range m.taskPlugins {
		go func(p TaskPlugin) { errors <- p.Started(s) }(p)
	}
	return waitForErrors(errors, len(m.taskPlugins))
}

func (m *taskPluginManager) Stopped(r engines.ResultSet) (bool, error) {
	// Sanity check that no two methods on plugin is running in parallel, this way
	// plugins don't have to be thread-safe, and we ensure nothing is called after
	// Dispose() has been called.
	if m.working.Swap(true) {
		panic("Another plugin method is currently running, or Dispose() has been called!")
	}
	defer m.working.Set(false)

	// Use atomic bool to return true, if no plugin returns false
	result := atomics.NewBool(true)

	// Run method on plugins in parallel
	errors := make(chan error)
	for _, p := range m.taskPlugins {
		go func(p TaskPlugin) {
			success, err := p.Stopped(r)
			if !success {
				result.Set(false)
			}
			errors <- err
		}(p)
	}
	return result.Get(), waitForErrors(errors, len(m.taskPlugins))
}

func (m *taskPluginManager) Finished(s bool) error {
	// Sanity check that no two methods on plugin is running in parallel, this way
	// plugins don't have to be thread-safe, and we ensure nothing is called after
	// Dispose() has been called.
	if m.working.Swap(true) {
		panic("Another plugin method is currently running, or Dispose() has been called!")
	}
	defer m.working.Set(false)

	// Run method on plugins in parallel
	errors := make(chan error)
	for _, p := range m.taskPlugins {
		go func(p TaskPlugin) { errors <- p.Finished(s) }(p)
	}
	return waitForErrors(errors, len(m.taskPlugins))
}

func (m *taskPluginManager) Exception(r runtime.ExceptionReason) error {
	// Sanity check that no two methods on plugin is running in parallel, this way
	// plugins don't have to be thread-safe, and we ensure nothing is called after
	// Dispose() has been called.
	if m.working.Swap(true) {
		panic("Another plugin method is currently running, or Dispose() has been called!")
	}
	defer m.working.Set(false)

	// Run method on plugins in parallel
	errors := make(chan error)
	for _, p := range m.taskPlugins {
		go func(p TaskPlugin) { errors <- p.Exception(r) }(p)
	}
	return waitForErrors(errors, len(m.taskPlugins))
}

func (m *taskPluginManager) Dispose() error {
	// Sanity check that no two methods on plugin is running in parallel
	if m.working.Swap(true) {
		panic("Another plugin method is currently running, or Dispose() has been called!")
	}
	// Notice that we don't call: defer w.working.Set(false), as we don't want to
	// allow any calls to plugins after Dispose()

	// Run method on plugins in parallel
	errors := make(chan error)
	for _, p := range m.taskPlugins {
		go func(p TaskPlugin) { errors <- p.Dispose() }(p)
	}
	return waitForErrors(errors, len(m.taskPlugins))
}
<commit_msg>Clean up pluginmanager (#178)<commit_after>package plugins

import (
	"errors"
	"fmt"
	"sync"

	schematypes "github.com/taskcluster/go-schematypes"
	"github.com/taskcluster/taskcluster-worker/engines"
	"github.com/taskcluster/taskcluster-worker/runtime"
	"github.com/taskcluster/taskcluster-worker/runtime/atomics"
)

type pluginManager struct {
	payloadSchema schematypes.Object
	plugins       []Plugin
	pluginNames   []string
}

type taskPluginManager struct {
	taskPlugins []TaskPlugin
	working     atomics.Bool
}

// mergeErrors merges a list of errors into one error, ignoring nil entries,
// and ensuring that if there is only MalformedPayloadErrors these will be
// merged into one MalformedPayloadError
func mergeErrors(errs ...error) error {
	msg := ""
	for _, err := range errs {
		if err != nil {
			// TODO MalformedPayloadError needs special care, if we only have these
			//      it is very different from having a list of other errors, Hence,
			//      we need to maintain the type rather than mindlessly wrap them.
			msg += err.Error() + "\n"
		}
	}
	if len(msg) > 0 {
		return errors.New(msg)
	}
	return nil
}

// PluginManagerConfigSchema returns configuration for PluginOptions.Config for
// NewPluginManager.
func PluginManagerConfigSchema() schematypes.Object {
	plugins := Plugins()

	pluginNames := []string{}
	for name := range plugins {
		pluginNames = append(pluginNames, name)
	}

	s := schematypes.Object{
		MetaData: schematypes.MetaData{
			Title: "Plugin Configuration",
			Description: `Mapping from plugin name to plugin configuration.
                    The 'disabled' key is special and lists plugins that are
                    disabled. Plugins that are disabled do not require
                    configuration.`,
		},
		Properties: schematypes.Properties{
			"disabled": schematypes.Array{
				MetaData: schematypes.MetaData{
					Title: "Disabled Plugins",
					Description: `List of disabled plugins. If a plugin is not listed
												as disabled here, then its configuration key must be
												specified, unless it doesn't take any configuration.`,
				},
				Items: schematypes.StringEnum{
					Options: pluginNames,
				},
			},
		},
		Required: []string{"disabled"},
	}
	for name, provider := range plugins {
		cs := provider.ConfigSchema()
		if cs != nil {
			s.Properties[name] = cs
		}
	}
	return s
}

// stringContains returns true if list contains element
func stringContains(list []string, element string) bool {
	for _, s := range list {
		if s == element {
			return true
		}
	}
	return false
}

// NewPluginManager loads all plugins not disabled in configuration and
// returns a Plugin implementation that wraps all of the plugins.
//
// This expects options.Config satisfying schema from
// PluginManagerConfigSchema().
func NewPluginManager(options PluginOptions) (Plugin, error) {
	pluginProviders := Plugins()

	// Construct config schema
	configSchema := PluginManagerConfigSchema()

	// Ensure the config is valid
	if err := configSchema.Validate(options.Config); err != nil {
		return nil, fmt.Errorf("Invalid config, error: %s", err)
	}
	config := options.Config.(map[string]interface{})

	// Find plugins to load
	var enabled []string
	var disabled []string
	if configSchema.Properties["disabled"].Map(config["disabled"], &disabled) != nil {
		panic("internal error -- shouldn't be possible")
	}

	// Find list of enabled plugins and ensure that config is present if required
	for name, plugin := range pluginProviders {
		// Skip disabled plugins
		if stringContains(disabled, name) {
			continue
		}
		// Check that configuration is given if required
		if plugin.ConfigSchema() != nil {
			if _, ok := config[name]; !ok {
				return nil, fmt.Errorf("Missing configuration for plugin: '%s'", name)
			}
		}
		// List plugin as enabled
		enabled = append(enabled, name)
	}

	// Initialize all the plugins
	wg := sync.WaitGroup{}
	plugins := make([]Plugin, len(enabled))
	errors := make([]error, len(enabled))
	wg.Add(len(enabled))
	for i, j := range enabled {
		go func(index int, name string) {
			plugins[index], errors[index] = pluginProviders[name].NewPlugin(PluginOptions{
				Environment: options.Environment,
				Engine:      options.Engine,
				Log:         options.Log.WithField("plugin", name),
				Config:      config[name],
			})
			wg.Done()
		}(i, j) // needed to capture values not variables
	}
	wg.Wait()

	// Return the first error, if any
	for _, e := range errors {
		if e != nil {
			return nil, e
		}
	}

	// Construct payload schema
	schemas := []schematypes.Object{}
	for _, plugin := range plugins {
		schemas = append(schemas, plugin.PayloadSchema())
	}
	schema, err := schematypes.Merge(schemas...)
	if err != nil {
		return nil, fmt.Errorf("Conflicting payload schema types, error: %s", err)
	}

	return &pluginManager{
		plugins:       plugins,
		pluginNames:   enabled,
		payloadSchema: schema,
	}, nil
}

func (m *pluginManager) PayloadSchema() schematypes.Object {
	return m.payloadSchema
}

func (m *pluginManager) NewTaskPlugin(options TaskPluginOptions) (manager TaskPlugin, err error) {
	// Input must be valid
	if m.payloadSchema.Validate(options.Payload) != nil {
		return nil, engines.ErrContractViolation
	}

	taskPlugins := make([]TaskPlugin, len(m.plugins))
	errors := make([]error, len(m.plugins))

	var wg sync.WaitGroup
	for i, j := range m.plugins {
		wg.Add(1)
		go func(index int, p Plugin) {
			defer wg.Done()
			taskPlugins[index], errors[index] = p.NewTaskPlugin(TaskPluginOptions{
				TaskInfo: options.TaskInfo,
				Payload:  p.PayloadSchema().Filter(options.Payload),
				Log:      options.Log.WithField("plugin", m.pluginNames[index]),
			})
		}(i, j)
	}
	wg.Wait()
	err = mergeErrors(errors...)

	manager = &taskPluginManager{
		taskPlugins: taskPlugins,
	}

	if err != nil {
		err2 := manager.Dispose()
		return nil, mergeErrors(err, err2)
	}
	return
}

type TaskPluginPhase func(TaskPlugin) error

func (m *taskPluginManager) executePhase(f TaskPluginPhase) error {
	// Sanity check that no two methods on plugin is running in parallel, this way
	// plugins don't have to be thread-safe, and we ensure nothing is called after
	// Dispose() has been called.
	if m.working.Swap(true) {
		panic("Another plugin method is currently running, or Dispose() has been called!")
	}
	defer m.working.Set(false)

	// Execute phase on plugins in parallel
	errors := make([]error, len(m.taskPlugins))
	var wg sync.WaitGroup
	for i, j := range m.taskPlugins {
		wg.Add(1)
		go func(index int, tp TaskPlugin) {
			defer wg.Done()
			errors[index] = f(tp)
		}(i, j)
	}
	wg.Wait()
	// Returned error represents the merge of all errors which occurred against
	// any plugin, or nil if no error occurred.
	return mergeErrors(errors...)
}

func (m *taskPluginManager) Prepare(c *runtime.TaskContext) error {
	return m.executePhase(func(p TaskPlugin) error { return p.Prepare(c) })
}

func (m *taskPluginManager) BuildSandbox(b engines.SandboxBuilder) error {
	return m.executePhase(func(p TaskPlugin) error { return p.BuildSandbox(b) })
}

func (m *taskPluginManager) Started(s engines.Sandbox) error {
	return m.executePhase(func(p TaskPlugin) error { return p.Started(s) })
}

func (m *taskPluginManager) Stopped(r engines.ResultSet) (bool, error) {
	// Use atomic bool to return true, if no plugin returns false
	result := atomics.NewBool(true)

	// Run method on plugins in parallel
	err := m.executePhase(func(p TaskPlugin) error {
		success, err := p.Stopped(r)
		if !success {
			result.Set(false)
		}
		return err
	})
	return result.Get(), err
}

func (m *taskPluginManager) Finished(s bool) error {
	return m.executePhase(func(p TaskPlugin) error { return p.Finished(s) })
}

func (m *taskPluginManager) Exception(r runtime.ExceptionReason) error {
	return m.executePhase(func(p TaskPlugin) error { return p.Exception(r) })
}

func (m *taskPluginManager) Dispose() error {
	// we don't want to allow any calls to plugins after Dispose()
	defer m.working.Set(true)
	return m.executePhase(func(p TaskPlugin) error { return p.Dispose() })
}
<|endoftext|>