<commit_before>// Copyright (c) 2016 Timo Savola. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package runtime

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"

	"github.com/tsavola/gate/internal/error/badprogram"
	"github.com/tsavola/gate/packet"
	"github.com/tsavola/gate/snapshot"
)

const (
	packetReservedOffset = 7
)

func popServiceBuffers(frozen *snapshot.Buffers) (services []snapshot.Service) {
	if frozen == nil {
		return
	}

	services = frozen.Services
	frozen.Services = nil
	return
}

func popInputBuffer(frozen *snapshot.Buffers) (input []byte) {
	if frozen == nil {
		return
	}

	input = frozen.Input
	frozen.Input = nil
	return
}

func popOutputBuffer(frozen *snapshot.Buffers) (output []byte) {
	if frozen == nil {
		return
	}

	output = frozen.Output
	frozen.Output = nil
	return
}

type read struct {
	buf packet.Buf
	err error
}

// ioLoop mutates Process and IOState (if any).
func ioLoop(ctx context.Context, services ServiceRegistry, subject *Process, frozen *snapshot.Buffers,
) (err error) {
	if frozen == nil {
		subject.writerOut.Close()
		subject.writerOut = nil
	}

	var (
		suspended = subject.suspended
		done      = ctx.Done()
	)

	var (
		messageInput  = make(chan packet.Buf)
		messageOutput = make(chan packet.Buf)
	)
	discoverer, initialServiceState, err := services.StartServing(ctx, ServiceConfig{maxPacketSize}, popServiceBuffers(frozen), messageInput, messageOutput)
	if err != nil {
		return
	}
	defer func() {
		close(messageOutput)

		if frozen != nil && suspended == nil { // Suspended.
			frozen.Services = discoverer.ExtractState()
		}

		if e := discoverer.Close(); err == nil {
			err = e
		}
	}()

	pendingMsg, initialRead, err := splitBufferedPackets(popOutputBuffer(frozen), discoverer)
	if err != nil {
		return
	}

	var pendingEvs []packet.Buf
	if ev := popInputBuffer(frozen); len(ev) > 0 {
		pendingEvs = []packet.Buf{ev} // No need to split packets.
	}

	if len(initialServiceState) > 0 {
		pendingEvs = append(pendingEvs, makeServicesPacket(packet.DomainState, initialServiceState))
	}
	initialServiceState = nil

	subjectInput := subjectReadLoop(subject.reader, initialRead)
	defer func() {
		if subjectInput != nil {
			for range subjectInput {
			}
		}
	}()

	subjectOutput := subjectWriteLoop(subject.writer)
	subject.writer = nil
	defer func() {
		if subjectOutput != nil {
			close(subjectOutput)
		}
	}()

	doSuspend := func() {
		suspended = nil
		done = nil

		subject.killSuspend()

		close(subjectOutput)
		subjectOutput = nil
	}

	for subjectInput != nil || pendingMsg != nil {
		var (
			doEv            packet.Buf
			doMessageInput  <-chan packet.Buf
			doMessageOutput chan<- packet.Buf
			doSubjectInput  <-chan read
			doSubjectOutput chan<- packet.Buf
		)

		if len(pendingEvs) > 0 {
			doEv = pendingEvs[0]
		}

		if pendingMsg != nil {
			doMessageOutput = messageOutput
		}

		if doEv == nil {
			doMessageInput = messageInput
			if pendingMsg == nil {
				doSubjectInput = subjectInput
			}
		} else {
			doSubjectOutput = subjectOutput
		}

		select {
		case p := <-doMessageInput:
			pendingEvs = append(pendingEvs, initMessagePacket(p))

		case read, ok := <-doSubjectInput:
			if !ok {
				err = errors.New("read loop terminated unexpectedly")
				return
			}

			switch {
			case read.err == nil:
				msg, ev, opErr := handlePacket(ctx, read.buf, discoverer)
				if opErr != nil {
					err = opErr
					return
				}

				if msg != nil {
					pendingMsg = msg
				}
				if ev != nil {
					pendingEvs = append(pendingEvs, ev)
				}

			case read.err == io.EOF:
				subjectInput = nil

				if subjectOutput != nil {
					close(subjectOutput)
					subjectOutput = nil
				}

				if frozen != nil {
					if suspended == nil { // Suspended.
						frozen.Output = append(pendingMsg, read.buf...)
						pendingMsg = nil

						frozen.Input, err = ioutil.ReadAll(subject.writerOut)
						if err != nil {
							return
						}

						var pendingLen int
						for _, ev := range pendingEvs {
							pendingLen += len(ev)
						}

						if n := len(frozen.Input) + pendingLen; cap(frozen.Input) < n {
							frozen.Input = append(make([]byte, 0, n), frozen.Input...)
						}

						for _, ev := range pendingEvs {
							frozen.Input = append(frozen.Input, ev...)
						}
					}

					subject.writerOut.Close()
					subject.writerOut = nil
				}

			case read.err != nil:
				err = read.err
				return
			}

		case doMessageOutput <- pendingMsg:
			pendingMsg = nil

		case doSubjectOutput <- doEv:
			if len(pendingEvs) > 0 {
				pendingEvs = pendingEvs[1:]
			}

		case <-suspended:
			doSuspend()

		case <-done:
			doSuspend()
		}
	}

	return
}

func subjectReadLoop(r *os.File, partial []byte) <-chan read {
	reads := make(chan read)

	go func() {
		defer close(reads)

		header := make([]byte, packet.HeaderSize)

		for {
			offset := copy(header, partial)
			partial = partial[offset:]

			if n, err := io.ReadFull(r, header[offset:]); err != nil {
				if err != io.EOF {
					err = fmt.Errorf("subject read: %v", err)
				}
				reads <- read{buf: header[:offset+n], err: err}
				return
			}

			size := binary.LittleEndian.Uint32(header)
			if size < packet.HeaderSize || size > maxPacketSize {
				reads <- read{err: badprogram.Errorf("runtime: invalid op packet size: %d", size)}
				return
			}

			buf := make([]byte, size)
			offset = copy(buf, header)
			offset += copy(buf[offset:], partial)
			partial = nil

			if n, err := io.ReadFull(r, buf[offset:]); err != nil {
				if err != io.EOF {
					err = fmt.Errorf("subject read: %v", err)
				}
				reads <- read{buf: buf[:offset+n], err: err}
				return
			}

			reads <- read{buf: buf}
		}
	}()

	return reads
}

func subjectWriteLoop(w *os.File) chan<- packet.Buf {
	writes := make(chan packet.Buf)

	go func() {
		defer w.Close()

		for buf := range writes {
			if _, err := w.Write(buf); err != nil {
				return
			}
		}
	}()

	return writes
}

func initMessagePacket(p packet.Buf) packet.Buf {
	if len(p) < packet.HeaderSize || len(p) > maxPacketSize {
		panic(errors.New("invalid outgoing message packet buffer length"))
	}

	if p[packetReservedOffset] != 0 {
		panic(errors.New("reserved byte is nonzero in outgoing message packet header"))
	}

	if p.Code() < 0 {
		panic(errors.New("negative service code in outgoing message packet header"))
	}

	if p.Domain() > packet.DomainData {
		panic(errors.New("invalid domain in outgoing message packet header"))
	}

	// Service implementations only need to initialize code and domain fields.
	binary.LittleEndian.PutUint32(p[packet.OffsetSize:], uint32(len(p)))

	return p
}

func handlePacket(ctx context.Context, p packet.Buf, discoverer ServiceDiscoverer) (msg, reply packet.Buf, err error) {
	switch code := p.Code(); {
	case code >= 0:
		msg, err = checkServicePacket(p, discoverer)
		if err != nil {
			return
		}

	case code == packet.CodeServices:
		reply, err = handleServicesPacket(ctx, p, discoverer)
		if err != nil {
			return
		}

	default:
		err = badprogram.Errorf("invalid code in incoming packet: %d", code)
		return
	}

	return
}

func splitBufferedPackets(buf []byte, discoverer ServiceDiscoverer,
) (msg packet.Buf, tail []byte, err error) {
	if len(buf) < packet.HeaderSize {
		tail = buf
		return
	}

	size := binary.LittleEndian.Uint32(buf[packet.OffsetSize:])
	if size < packet.HeaderSize || size > maxPacketSize {
		err = badprogram.Errorf("buffered packet has invalid size: %d", size)
		return
	}

	if uint32(len(buf)) < size {
		tail = buf
		return
	}

	p := packet.Buf(buf[:size])

	switch code := p.Code(); {
	case code >= 0:
		msg, err = checkServicePacket(p, discoverer)
		if err != nil {
			return
		}

	default:
		err = badprogram.Errorf("invalid code in buffered packet: %d", code)
		return
	}

	tail = buf[size:]
	return
}

func checkServicePacket(p packet.Buf, discoverer ServiceDiscoverer) (msg packet.Buf, err error) {
	if x := p[packetReservedOffset]; x != 0 {
		err = badprogram.Errorf("reserved byte has value 0x%02x in buffered packet header", x)
		return
	}

	if int(p.Code()) >= discoverer.NumServices() {
		err = badprogram.Errorf("invalid service code in packet: %d", p.Code())
		return
	}

	switch p.Domain() {
	case packet.DomainCall, packet.DomainFlow:

	case packet.DomainData:
		if n := len(p); n < packet.DataHeaderSize {
			err = badprogram.Errorf("data packet is too short: %d bytes", n)
			return
		}

	default:
		err = badprogram.Errorf("invalid domain in packet: %d", p.Domain())
		return
	}

	msg = p
	return
}
<commit_msg>runtime: treat packet buffers as read-only if possible<commit_after>// Copyright (c) 2016 Timo Savola. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package runtime

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"

	"github.com/tsavola/gate/internal/error/badprogram"
	"github.com/tsavola/gate/packet"
	"github.com/tsavola/gate/snapshot"
)

const (
	packetReservedOffset = 7
)

func popServiceBuffers(frozen *snapshot.Buffers) (services []snapshot.Service) {
	if frozen == nil {
		return
	}

	services = frozen.Services
	frozen.Services = nil
	return
}

func popInputBuffer(frozen *snapshot.Buffers) (input []byte) {
	if frozen == nil {
		return
	}

	input = frozen.Input
	frozen.Input = nil
	return
}

func popOutputBuffer(frozen *snapshot.Buffers) (output []byte) {
	if frozen == nil {
		return
	}

	output = frozen.Output
	frozen.Output = nil
	return
}

type read struct {
	buf packet.Buf
	err error
}

// ioLoop mutates Process and IOState (if any).
func ioLoop(ctx context.Context, services ServiceRegistry, subject *Process, frozen *snapshot.Buffers,
) (err error) {
	if frozen == nil {
		subject.writerOut.Close()
		subject.writerOut = nil
	}

	var (
		suspended = subject.suspended
		done      = ctx.Done()
	)

	var (
		messageInput  = make(chan packet.Buf)
		messageOutput = make(chan packet.Buf)
	)
	discoverer, initialServiceState, err := services.StartServing(ctx, ServiceConfig{maxPacketSize}, popServiceBuffers(frozen), messageInput, messageOutput)
	if err != nil {
		return
	}
	defer func() {
		close(messageOutput)

		if frozen != nil && suspended == nil { // Suspended.
			frozen.Services = discoverer.ExtractState()
		}

		if e := discoverer.Close(); err == nil {
			err = e
		}
	}()

	pendingMsg, initialRead, err := splitBufferedPackets(popOutputBuffer(frozen), discoverer)
	if err != nil {
		return
	}

	var pendingEvs []packet.Buf
	if ev := popInputBuffer(frozen); len(ev) > 0 {
		pendingEvs = []packet.Buf{ev} // No need to split packets.
	}

	if len(initialServiceState) > 0 {
		pendingEvs = append(pendingEvs, makeServicesPacket(packet.DomainState, initialServiceState))
	}
	initialServiceState = nil

	subjectInput := subjectReadLoop(subject.reader, initialRead)
	defer func() {
		if subjectInput != nil {
			for range subjectInput {
			}
		}
	}()

	subjectOutput := subjectWriteLoop(subject.writer)
	subject.writer = nil
	defer func() {
		if subjectOutput != nil {
			close(subjectOutput)
		}
	}()

	doSuspend := func() {
		suspended = nil
		done = nil

		subject.killSuspend()

		close(subjectOutput)
		subjectOutput = nil
	}

	for subjectInput != nil || pendingMsg != nil {
		var (
			doEv            packet.Buf
			doMessageInput  <-chan packet.Buf
			doMessageOutput chan<- packet.Buf
			doSubjectInput  <-chan read
			doSubjectOutput chan<- packet.Buf
		)

		if len(pendingEvs) > 0 {
			doEv = pendingEvs[0]
		}

		if pendingMsg != nil {
			doMessageOutput = messageOutput
		}

		if doEv == nil {
			doMessageInput = messageInput
			if pendingMsg == nil {
				doSubjectInput = subjectInput
			}
		} else {
			doSubjectOutput = subjectOutput
		}

		select {
		case p := <-doMessageInput:
			pendingEvs = append(pendingEvs, initMessagePacket(p))

		case read, ok := <-doSubjectInput:
			if !ok {
				err = errors.New("read loop terminated unexpectedly")
				return
			}

			switch {
			case read.err == nil:
				msg, ev, opErr := handlePacket(ctx, read.buf, discoverer)
				if opErr != nil {
					err = opErr
					return
				}

				if msg != nil {
					pendingMsg = msg
				}
				if ev != nil {
					pendingEvs = append(pendingEvs, ev)
				}

			case read.err == io.EOF:
				subjectInput = nil

				if subjectOutput != nil {
					close(subjectOutput)
					subjectOutput = nil
				}

				if frozen != nil {
					if suspended == nil { // Suspended.
						frozen.Output = append(pendingMsg, read.buf...)
						pendingMsg = nil

						frozen.Input, err = ioutil.ReadAll(subject.writerOut)
						if err != nil {
							return
						}

						var pendingLen int
						for _, ev := range pendingEvs {
							pendingLen += len(ev)
						}

						if n := len(frozen.Input) + pendingLen; cap(frozen.Input) < n {
							frozen.Input = append(make([]byte, 0, n), frozen.Input...)
						}

						for _, ev := range pendingEvs {
							frozen.Input = append(frozen.Input, ev...)
						}
					}

					subject.writerOut.Close()
					subject.writerOut = nil
				}

			case read.err != nil:
				err = read.err
				return
			}

		case doMessageOutput <- pendingMsg:
			pendingMsg = nil

		case doSubjectOutput <- doEv:
			if len(pendingEvs) > 0 {
				pendingEvs = pendingEvs[1:]
			}

		case <-suspended:
			doSuspend()

		case <-done:
			doSuspend()
		}
	}

	return
}

func subjectReadLoop(r *os.File, partial []byte) <-chan read {
	reads := make(chan read)

	go func() {
		defer close(reads)

		header := make([]byte, packet.HeaderSize)

		for {
			offset := copy(header, partial)
			partial = partial[offset:]

			if n, err := io.ReadFull(r, header[offset:]); err != nil {
				if err != io.EOF {
					err = fmt.Errorf("subject read: %v", err)
				}
				reads <- read{buf: header[:offset+n], err: err}
				return
			}

			size := binary.LittleEndian.Uint32(header)
			if size < packet.HeaderSize || size > maxPacketSize {
				reads <- read{err: badprogram.Errorf("runtime: invalid op packet size: %d", size)}
				return
			}

			buf := make([]byte, size)
			offset = copy(buf, header)
			offset += copy(buf[offset:], partial)
			partial = nil

			if n, err := io.ReadFull(r, buf[offset:]); err != nil {
				if err != io.EOF {
					err = fmt.Errorf("subject read: %v", err)
				}
				reads <- read{buf: buf[:offset+n], err: err}
				return
			}

			reads <- read{buf: buf}
		}
	}()

	return reads
}

func subjectWriteLoop(w *os.File) chan<- packet.Buf {
	writes := make(chan packet.Buf)

	go func() {
		defer w.Close()

		for buf := range writes {
			if _, err := w.Write(buf); err != nil {
				return
			}
		}
	}()

	return writes
}

func initMessagePacket(p packet.Buf) packet.Buf {
	if len(p) < packet.HeaderSize || len(p) > maxPacketSize {
		panic(errors.New("invalid outgoing message packet buffer length"))
	}

	if p[packetReservedOffset] != 0 {
		panic(errors.New("reserved byte is nonzero in outgoing message packet header"))
	}

	if p.Code() < 0 {
		panic(errors.New("negative service code in outgoing message packet header"))
	}

	if p.Domain() > packet.DomainData {
		panic(errors.New("invalid domain in outgoing message packet header"))
	}

	// Service implementations only need to initialize code and domain fields.
	// But if the size field has been initialized correctly, treat the buffer
	// as read-only.
	if binary.LittleEndian.Uint32(p[packet.OffsetSize:]) != uint32(len(p)) {
		binary.LittleEndian.PutUint32(p[packet.OffsetSize:], uint32(len(p)))
	}

	return p
}

func handlePacket(ctx context.Context, p packet.Buf, discoverer ServiceDiscoverer) (msg, reply packet.Buf, err error) {
	switch code := p.Code(); {
	case code >= 0:
		msg, err = checkServicePacket(p, discoverer)
		if err != nil {
			return
		}

	case code == packet.CodeServices:
		reply, err = handleServicesPacket(ctx, p, discoverer)
		if err != nil {
			return
		}

	default:
		err = badprogram.Errorf("invalid code in incoming packet: %d", code)
		return
	}

	return
}

func splitBufferedPackets(buf []byte, discoverer ServiceDiscoverer,
) (msg packet.Buf, tail []byte, err error) {
	if len(buf) < packet.HeaderSize {
		tail = buf
		return
	}

	size := binary.LittleEndian.Uint32(buf[packet.OffsetSize:])
	if size < packet.HeaderSize || size > maxPacketSize {
		err = badprogram.Errorf("buffered packet has invalid size: %d", size)
		return
	}

	if uint32(len(buf)) < size {
		tail = buf
		return
	}

	p := packet.Buf(buf[:size])

	switch code := p.Code(); {
	case code >= 0:
		msg, err = checkServicePacket(p, discoverer)
		if err != nil {
			return
		}

	default:
		err = badprogram.Errorf("invalid code in buffered packet: %d", code)
		return
	}

	tail = buf[size:]
	return
}

func checkServicePacket(p packet.Buf, discoverer ServiceDiscoverer) (msg packet.Buf, err error) {
	if x := p[packetReservedOffset]; x != 0 {
		err = badprogram.Errorf("reserved byte has value 0x%02x in buffered packet header", x)
		return
	}

	if int(p.Code()) >= discoverer.NumServices() {
		err = badprogram.Errorf("invalid service code in packet: %d", p.Code())
		return
	}

	switch p.Domain() {
	case packet.DomainCall, packet.DomainFlow:

	case packet.DomainData:
		if n := len(p); n < packet.DataHeaderSize {
			err = badprogram.Errorf("data packet is too short: %d bytes", n)
			return
		}

	default:
		err = badprogram.Errorf("invalid domain in packet: %d", p.Domain())
		return
	}

	msg = p
	return
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package common

import (
	stdmath "math"
	"time"

	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/utils/math"
)

const (
	// MaxContainersPerMultiPut is the maximum number of containers that can be sent in a MultiPut
	MaxContainersPerMultiPut = 2000

	// StatusUpdateFrequency ... bootstrapper logs "processed X blocks/vertices" every [statusUpdateFrequency] blocks/vertices
	StatusUpdateFrequency = 2500

	// MaxOutstandingRequests is the maximum number of GetAncestors sent but not responsded to/failed
	MaxOutstandingRequests = 15
)

var (
	// MaxTimeFetchingAncestors is the maximum amount of time to spend fetching vertices during a call to GetAncestors
	MaxTimeFetchingAncestors = 100 * time.Millisecond
)

// Bootstrapper implements the Engine interface.
type Bootstrapper struct {
	Config

	// IDs of validators we have requested the accepted frontier from but haven't
	// received a reply from
	pendingAcceptedFrontier ids.ShortSet
	acceptedFrontier        ids.Set

	pendingAccepted ids.ShortSet
	acceptedVotes   map[[32]byte]uint64

	RequestID uint32
}

// Initialize implements the Engine interface.
func (b *Bootstrapper) Initialize(config Config) {
	b.Config = config

	for _, vdr := range b.Beacons.List() {
		vdrID := vdr.ID()
		b.pendingAcceptedFrontier.Add(vdrID)
		b.pendingAccepted.Add(vdrID)
	}

	b.acceptedVotes = make(map[[32]byte]uint64)
}

// Startup implements the Engine interface.
func (b *Bootstrapper) Startup() error {
	if b.pendingAcceptedFrontier.Len() == 0 {
		b.Context.Log.Info("Bootstrapping skipped due to no provided bootstraps")
		return b.Bootstrapable.ForceAccepted(ids.Set{})
	}

	// Ask each of the bootstrap validators to send their accepted frontier
	vdrs := ids.ShortSet{}
	vdrs.Union(b.pendingAcceptedFrontier)

	b.RequestID++
	b.Sender.GetAcceptedFrontier(vdrs, b.RequestID)
	return nil
}

// GetAcceptedFrontier implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFrontier(validatorID ids.ShortID, requestID uint32) error {
	b.Sender.AcceptedFrontier(validatorID, requestID, b.Bootstrapable.CurrentAcceptedFrontier())
	return nil
}

// GetAcceptedFrontierFailed implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFrontierFailed(validatorID ids.ShortID, requestID uint32) error {
	// If we can't get a response from [validatorID], act as though they said their accepted frontier is empty
	b.AcceptedFrontier(validatorID, requestID, ids.Set{})
	return nil
}

// AcceptedFrontier implements the Engine interface.
func (b *Bootstrapper) AcceptedFrontier(validatorID ids.ShortID, requestID uint32, containerIDs ids.Set) error {
	if !b.pendingAcceptedFrontier.Contains(validatorID) {
		b.Context.Log.Debug("Received an AcceptedFrontier message from %s unexpectedly", validatorID)
		return nil
	}
	// Mark that we received a response from [validatorID]
	b.pendingAcceptedFrontier.Remove(validatorID)

	// Union the reported accepted frontier from [validatorID] with the accepted frontier we got from others
	b.acceptedFrontier.Union(containerIDs)

	// We've received the accepted frontier from every bootstrap validator
	// Ask each bootstrap validator to filter the list of containers that we were
	// told are on the accepted frontier such that the list only contains containers
	// they think are accepted
	if b.pendingAcceptedFrontier.Len() == 0 {
		vdrs := ids.ShortSet{}
		vdrs.Union(b.pendingAccepted)

		b.RequestID++
		b.Sender.GetAccepted(vdrs, b.RequestID, b.acceptedFrontier)
	}
	return nil
}

// GetAccepted implements the Engine interface.
func (b *Bootstrapper) GetAccepted(validatorID ids.ShortID, requestID uint32, containerIDs ids.Set) error {
	b.Sender.Accepted(validatorID, requestID, b.Bootstrapable.FilterAccepted(containerIDs))
	return nil
}

// GetAcceptedFailed implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFailed(validatorID ids.ShortID, requestID uint32) error {
	// If we can't get a response from [validatorID], act as though they said
	// that they think none of the containers we sent them in GetAccepted are accepted
	return b.Accepted(validatorID, requestID, ids.Set{})
}

// Accepted implements the Engine interface.
func (b *Bootstrapper) Accepted(validatorID ids.ShortID, requestID uint32, containerIDs ids.Set) error {
	if !b.pendingAccepted.Contains(validatorID) {
		b.Context.Log.Debug("Received an Accepted message from %s unexpectedly", validatorID)
		return nil
	}
	// Mark that we received a response from [validatorID]
	b.pendingAccepted.Remove(validatorID)

	weight := uint64(0)
	if vdr, ok := b.Beacons.Get(validatorID); ok {
		weight = vdr.Weight()
	}

	for _, containerID := range containerIDs.List() {
		key := containerID.Key()
		previousWeight := b.acceptedVotes[key]
		newWeight, err := math.Add64(weight, previousWeight)
		if err != nil {
			newWeight = stdmath.MaxUint64
		}
		b.acceptedVotes[key] = newWeight
	}

	if b.pendingAccepted.Len() != 0 {
		return nil
	}

	// We've received the filtered accepted frontier from every bootstrap validator
	// Accept all containers that have a sufficient weight behind them
	accepted := ids.Set{}
	for key, weight := range b.acceptedVotes {
		if weight >= b.Config.Alpha {
			accepted.Add(ids.NewID(key))
		}
	}

	if size := accepted.Len(); size == 0 && b.Config.Beacons.Len() > 0 {
		b.Context.Log.Info("Bootstrapping finished with no accepted frontier. This is likely a result of failing to be able to connect to the specified bootstraps, or no transactions have been issued on this chain yet")
	} else {
		b.Context.Log.Info("Bootstrapping started syncing with %d vertices in the accepted frontier", size)
	}

	return b.Bootstrapable.ForceAccepted(accepted)
}
<commit_msg>change maximum # outstanding to 8 to reduce load on nodes<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package common

import (
	stdmath "math"
	"time"

	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/utils/math"
)

const (
	// MaxContainersPerMultiPut is the maximum number of containers that can be sent in a MultiPut
	MaxContainersPerMultiPut = 2000

	// StatusUpdateFrequency ... bootstrapper logs "processed X blocks/vertices" every [statusUpdateFrequency] blocks/vertices
	StatusUpdateFrequency = 2500

	// MaxOutstandingRequests is the maximum number of GetAncestors sent but not responsded to/failed
	MaxOutstandingRequests = 8
)

var (
	// MaxTimeFetchingAncestors is the maximum amount of time to spend fetching vertices during a call to GetAncestors
	MaxTimeFetchingAncestors = 100 * time.Millisecond
)

// Bootstrapper implements the Engine interface.
type Bootstrapper struct {
	Config

	// IDs of validators we have requested the accepted frontier from but haven't
	// received a reply from
	pendingAcceptedFrontier ids.ShortSet
	acceptedFrontier        ids.Set

	pendingAccepted ids.ShortSet
	acceptedVotes   map[[32]byte]uint64

	RequestID uint32
}

// Initialize implements the Engine interface.
func (b *Bootstrapper) Initialize(config Config) {
	b.Config = config

	for _, vdr := range b.Beacons.List() {
		vdrID := vdr.ID()
		b.pendingAcceptedFrontier.Add(vdrID)
		b.pendingAccepted.Add(vdrID)
	}

	b.acceptedVotes = make(map[[32]byte]uint64)
}

// Startup implements the Engine interface.
func (b *Bootstrapper) Startup() error {
	if b.pendingAcceptedFrontier.Len() == 0 {
		b.Context.Log.Info("Bootstrapping skipped due to no provided bootstraps")
		return b.Bootstrapable.ForceAccepted(ids.Set{})
	}

	// Ask each of the bootstrap validators to send their accepted frontier
	vdrs := ids.ShortSet{}
	vdrs.Union(b.pendingAcceptedFrontier)

	b.RequestID++
	b.Sender.GetAcceptedFrontier(vdrs, b.RequestID)
	return nil
}

// GetAcceptedFrontier implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFrontier(validatorID ids.ShortID, requestID uint32) error {
	b.Sender.AcceptedFrontier(validatorID, requestID, b.Bootstrapable.CurrentAcceptedFrontier())
	return nil
}

// GetAcceptedFrontierFailed implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFrontierFailed(validatorID ids.ShortID, requestID uint32) error {
	// If we can't get a response from [validatorID], act as though they said their accepted frontier is empty
	b.AcceptedFrontier(validatorID, requestID, ids.Set{})
	return nil
}

// AcceptedFrontier implements the Engine interface.
func (b *Bootstrapper) AcceptedFrontier(validatorID ids.ShortID, requestID uint32, containerIDs ids.Set) error {
	if !b.pendingAcceptedFrontier.Contains(validatorID) {
		b.Context.Log.Debug("Received an AcceptedFrontier message from %s unexpectedly", validatorID)
		return nil
	}
	// Mark that we received a response from [validatorID]
	b.pendingAcceptedFrontier.Remove(validatorID)

	// Union the reported accepted frontier from [validatorID] with the accepted frontier we got from others
	b.acceptedFrontier.Union(containerIDs)

	// We've received the accepted frontier from every bootstrap validator
	// Ask each bootstrap validator to filter the list of containers that we were
	// told are on the accepted frontier such that the list only contains containers
	// they think are accepted
	if b.pendingAcceptedFrontier.Len() == 0 {
		vdrs := ids.ShortSet{}
		vdrs.Union(b.pendingAccepted)

		b.RequestID++
		b.Sender.GetAccepted(vdrs, b.RequestID, b.acceptedFrontier)
	}
	return nil
}

// GetAccepted implements the Engine interface.
func (b *Bootstrapper) GetAccepted(validatorID ids.ShortID, requestID uint32, containerIDs ids.Set) error {
	b.Sender.Accepted(validatorID, requestID, b.Bootstrapable.FilterAccepted(containerIDs))
	return nil
}

// GetAcceptedFailed implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFailed(validatorID ids.ShortID, requestID uint32) error {
	// If we can't get a response from [validatorID], act as though they said
	// that they think none of the containers we sent them in GetAccepted are accepted
	return b.Accepted(validatorID, requestID, ids.Set{})
}

// Accepted implements the Engine interface.
func (b *Bootstrapper) Accepted(validatorID ids.ShortID, requestID uint32, containerIDs ids.Set) error {
	if !b.pendingAccepted.Contains(validatorID) {
		b.Context.Log.Debug("Received an Accepted message from %s unexpectedly", validatorID)
		return nil
	}
	// Mark that we received a response from [validatorID]
	b.pendingAccepted.Remove(validatorID)

	weight := uint64(0)
	if vdr, ok := b.Beacons.Get(validatorID); ok {
		weight = vdr.Weight()
	}

	for _, containerID := range containerIDs.List() {
		key := containerID.Key()
		previousWeight := b.acceptedVotes[key]
		newWeight, err := math.Add64(weight, previousWeight)
		if err != nil {
			newWeight = stdmath.MaxUint64
		}
		b.acceptedVotes[key] = newWeight
	}

	if b.pendingAccepted.Len() != 0 {
		return nil
	}

	// We've received the filtered accepted frontier from every bootstrap validator
	// Accept all containers that have a sufficient weight behind them
	accepted := ids.Set{}
	for key, weight := range b.acceptedVotes {
		if weight >= b.Config.Alpha {
			accepted.Add(ids.NewID(key))
		}
	}

	if size := accepted.Len(); size == 0 && b.Config.Beacons.Len() > 0 {
		b.Context.Log.Info("Bootstrapping finished with no accepted frontier. This is likely a result of failing to be able to connect to the specified bootstraps, or no transactions have been issued on this chain yet")
	} else {
		b.Context.Log.Info("Bootstrapping started syncing with %d vertices in the accepted frontier", size)
	}

	return b.Bootstrapable.ForceAccepted(accepted)
}
<|endoftext|>


<commit_before>package space

import (
	"fmt"
	"strings"

	. "github.com/cloudfoundry/cli/cf/i18n"
	"github.com/cloudfoundry/cli/cf/models"

	"github.com/cloudfoundry/cli/cf/api/space_quotas"
	"github.com/cloudfoundry/cli/cf/command_metadata"
	"github.com/cloudfoundry/cli/cf/configuration/core_config"
	"github.com/cloudfoundry/cli/cf/formatters"
	"github.com/cloudfoundry/cli/cf/requirements"
	"github.com/cloudfoundry/cli/cf/terminal"
	"github.com/codegangsta/cli"
)

type ShowSpace struct {
	ui        terminal.UI
	config    core_config.Reader
	spaceReq  requirements.SpaceRequirement
	quotaRepo space_quotas.SpaceQuotaRepository
}

func NewShowSpace(ui terminal.UI, config core_config.Reader, quotaRepo space_quotas.SpaceQuotaRepository) (cmd *ShowSpace) {
	cmd = new(ShowSpace)
	cmd.ui = ui
	cmd.config = config
	cmd.quotaRepo = quotaRepo
	return
}

func (cmd *ShowSpace) Metadata() command_metadata.CommandMetadata {
	return command_metadata.CommandMetadata{
		Name:        "space",
		Description: T("Show space info"),
		Usage:       T("CF_NAME space SPACE"),
	}
}

func (cmd *ShowSpace) GetRequirements(requirementsFactory requirements.Factory, c *cli.Context) (reqs []requirements.Requirement, err error) {
	if len(c.Args()) != 1 {
		cmd.ui.FailWithUsage(c)
	}

	cmd.spaceReq = requirementsFactory.NewSpaceRequirement(c.Args()[0])
	reqs = []requirements.Requirement{
		requirementsFactory.NewLoginRequirement(),
		requirementsFactory.NewTargetedOrgRequirement(),
		cmd.spaceReq,
	}
	return
}

func (cmd *ShowSpace) Run(c *cli.Context) {
	space := cmd.spaceReq.GetSpace()
	cmd.ui.Say(T("Getting info for space {{.TargetSpace}} in org {{.OrgName}} as {{.CurrentUser}}...",
		map[string]interface{}{
			"TargetSpace": terminal.EntityNameColor(space.Name),
			"OrgName":     terminal.EntityNameColor(space.Organization.Name),
			"CurrentUser": terminal.EntityNameColor(cmd.config.Username()),
		}))

	quotaString := cmd.quotaString(space)

	cmd.ui.Ok()
	cmd.ui.Say("")
	cmd.ui.Say(terminal.EntityNameColor(space.Name) + ":")

	table := terminal.NewTable(cmd.ui, []string{"", "", ""})
	table.Add("", T("Org:"), terminal.EntityNameColor(space.Organization.Name))

	apps := []string{}
	for _, app := range space.Applications {
		apps = append(apps, terminal.EntityNameColor(app.Name))
	}
	table.Add("", T("Apps:"), strings.Join(apps, ", "))

	domains := []string{}
	for _, domain := range space.Domains {
		domains = append(domains, terminal.EntityNameColor(domain.Name))
	}
	table.Add("", T("Domains:"), strings.Join(domains, ", "))

	services := []string{}
	for _, service := range space.ServiceInstances {
		services = append(services, terminal.EntityNameColor(service.Name))
	}
	table.Add("", T("Services:"), strings.Join(services, ", "))

	securityGroups := []string{}
	for _, group := range space.SecurityGroups {
		securityGroups = append(securityGroups, terminal.EntityNameColor(group.Name))
	}
	table.Add("", T("Security Groups:"), strings.Join(securityGroups, ", "))

	table.Add("", T("Space Quota:"), quotaString)

	table.Print()
}

func (cmd *ShowSpace) quotaString(space models.Space) string {
	var instance_memory string

	if space.SpaceQuotaGuid == "" {
		return ""
	}

	quota, err := cmd.quotaRepo.FindByGuid(space.SpaceQuotaGuid)
	if err != nil {
		cmd.ui.Failed(err.Error())
		return ""
	}

	if quota.InstanceMemoryLimit == -1 {
		instance_memory = "-1"
	} else {
		instance_memory = formatters.ByteSize(quota.InstanceMemoryLimit * formatters.MEGABYTE)
	}
	memory := formatters.ByteSize(quota.MemoryLimit * formatters.MEGABYTE)

	spaceQuota := fmt.Sprintf("%s (%s memory limit, %s instance memory limit, %d routes, %d services, paid services %s)", quota.Name, memory, instance_memory, quota.RoutesLimit, quota.ServicesLimit, formatters.Allowed(quota.NonBasicServicesAllowed))
	//	spaceQuota := fmt.Sprintf(T("{{.QuotaName}} ({{.MemoryLimit}} memory limit, {{.InstanceMemoryLimit}} instance memory limit, {{.RoutesLimit}} routes, {{.ServicesLimit}} services, paid services {{.NonBasicServicesAllowed}})",
	//		map[string]interface{}{
	//			"QuotaName":               quota.Name,
	//			"MemoryLimit":             memory,
	//			"InstanceMemoryLimit":     instance_memory,
	//			"RoutesLimit":             quota.RoutesLimit,
	//			"ServicesLimit":           quota.ServicesLimit,
	//			"NonBasicServicesAllowed": formatters.Allowed(quota.NonBasicServicesAllowed)}))

	return spaceQuota
}
<commit_msg>Remove space from command `cf space ...` between space and table.<commit_after>package space

import (
	"fmt"
	"strings"

	. "github.com/cloudfoundry/cli/cf/i18n"
	"github.com/cloudfoundry/cli/cf/models"

	"github.com/cloudfoundry/cli/cf/api/space_quotas"
	"github.com/cloudfoundry/cli/cf/command_metadata"
	"github.com/cloudfoundry/cli/cf/configuration/core_config"
	"github.com/cloudfoundry/cli/cf/formatters"
	"github.com/cloudfoundry/cli/cf/requirements"
	"github.com/cloudfoundry/cli/cf/terminal"
	"github.com/codegangsta/cli"
)

type ShowSpace struct {
	ui        terminal.UI
	config    core_config.Reader
	spaceReq  requirements.SpaceRequirement
	quotaRepo space_quotas.SpaceQuotaRepository
}

func NewShowSpace(ui terminal.UI, config core_config.Reader, quotaRepo space_quotas.SpaceQuotaRepository) (cmd *ShowSpace) {
	cmd = new(ShowSpace)
	cmd.ui = ui
	cmd.config = config
	cmd.quotaRepo = quotaRepo
	return
}

func (cmd *ShowSpace) Metadata() command_metadata.CommandMetadata {
	return command_metadata.CommandMetadata{
		Name:        "space",
		Description: T("Show space info"),
		Usage:       T("CF_NAME space SPACE"),
	}
}

func (cmd *ShowSpace) GetRequirements(requirementsFactory requirements.Factory, c *cli.Context) (reqs []requirements.Requirement, err error) {
	if len(c.Args()) != 1 {
		cmd.ui.FailWithUsage(c)
	}

	cmd.spaceReq = requirementsFactory.NewSpaceRequirement(c.Args()[0])
	reqs = []requirements.Requirement{
		requirementsFactory.NewLoginRequirement(),
		requirementsFactory.NewTargetedOrgRequirement(),
		cmd.spaceReq,
	}
	return
}

func (cmd *ShowSpace) Run(c *cli.Context) {
	space := cmd.spaceReq.GetSpace()
	cmd.ui.Say(T("Getting info for space {{.TargetSpace}} in org {{.OrgName}} as {{.CurrentUser}}...",
		map[string]interface{}{
			"TargetSpace": terminal.EntityNameColor(space.Name),
			"OrgName":     terminal.EntityNameColor(space.Organization.Name),
			"CurrentUser": terminal.EntityNameColor(cmd.config.Username()),
		}))

	quotaString := cmd.quotaString(space)

	cmd.ui.Ok()
	cmd.ui.Say("")
	table := terminal.NewTable(cmd.ui, []string{terminal.EntityNameColor(space.Name), "", ""})
	table.Add("", T("Org:"), terminal.EntityNameColor(space.Organization.Name))

	apps := []string{}
	for _, app := range space.Applications {
		apps = append(apps, terminal.EntityNameColor(app.Name))
	}
	table.Add("", T("Apps:"), strings.Join(apps, ", "))

	domains := []string{}
	for _, domain := range space.Domains {
		domains = append(domains, terminal.EntityNameColor(domain.Name))
	}
	table.Add("", T("Domains:"), strings.Join(domains, ", "))

	services := []string{}
	for _, service := range space.ServiceInstances {
		services = append(services, terminal.EntityNameColor(service.Name))
	}
	table.Add("", T("Services:"), strings.Join(services, ", "))

	securityGroups := []string{}
	for _, group := range space.SecurityGroups {
		securityGroups = append(securityGroups, terminal.EntityNameColor(group.Name))
	}
	table.Add("", T("Security Groups:"), strings.Join(securityGroups, ", "))

	table.Add("", T("Space Quota:"), quotaString)

	table.Print()
}

func (cmd *ShowSpace) quotaString(space models.Space) string {
	var instance_memory string

	if space.SpaceQuotaGuid == "" {
		return ""
	}

	quota, err := cmd.quotaRepo.FindByGuid(space.SpaceQuotaGuid)
	if err != nil {
		cmd.ui.Failed(err.Error())
		return ""
	}

	if quota.InstanceMemoryLimit == -1 {
		instance_memory = "-1"
	} else {
		instance_memory = formatters.ByteSize(quota.InstanceMemoryLimit * formatters.MEGABYTE)
	}
	memory := formatters.ByteSize(quota.MemoryLimit * formatters.MEGABYTE)

	spaceQuota := fmt.Sprintf("%s (%s memory limit, %s instance memory limit, %d routes, %d services, paid services %s)", quota.Name, memory, instance_memory, quota.RoutesLimit, quota.ServicesLimit, formatters.Allowed(quota.NonBasicServicesAllowed))
	//	spaceQuota := fmt.Sprintf(T("{{.QuotaName}} ({{.MemoryLimit}} memory limit, {{.InstanceMemoryLimit}} instance memory limit, {{.RoutesLimit}} routes, {{.ServicesLimit}} services, paid services {{.NonBasicServicesAllowed}})",
	//		map[string]interface{}{
	//			"QuotaName":               quota.Name,
	//			"MemoryLimit":             memory,
	//			"InstanceMemoryLimit":     instance_memory,
	//			"RoutesLimit":             quota.RoutesLimit,
	//			"ServicesLimit":           quota.ServicesLimit,
	//			"NonBasicServicesAllowed": formatters.Allowed(quota.NonBasicServicesAllowed)}))

	return spaceQuota
}
<|endoftext|>


<commit_before>package itembase

import (
	"crypto/rand"
	"errors"
	"fmt"
	"log"
	"net/http"

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

func (c *client) newConf() *oauth2.Config {
	var endpointURL string

	if c.production {
		endpointURL = "https://accounts.itembase.com/oauth/v2"
		c.me = "https://users.itembase.com/v1/me"
		c.root = "https://api.itembase.io/v1"
		c.activation = "https://solutionservice.itembase.com"
	} else {
		endpointURL = "http://sandbox.accounts.itembase.io/oauth/v2"
		c.me = "http://sandbox.users.itembase.io/v1/me"
		c.root = "http://sandbox.api.itembase.io/v1"
		c.activation = "http://sandbox.solutionservice.itembase.io"
	}

	return &oauth2.Config{
		ClientID:     c.options.ClientID,
		ClientSecret: c.options.ClientSecret,
		Scopes:       c.options.Scopes,
		RedirectURL:  c.options.RedirectURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  endpointURL + "/auth",
			TokenURL: endpointURL + "/token",
		},
	}
}

func (c *client) SaveToken(userID string, token *oauth2.Token) (err error) {
	if c.options.TokenHandler.TokenSaver != nil {
		err = c.options.TokenHandler.TokenSaver(userID, token)
	} else {
		err = errors.New("No Token Store!")
	}
	return
}

func (c *client) GetCachedToken(userID string) (token *oauth2.Token, err error) {
	if c.options.TokenHandler.TokenLoader != nil {
		token, err = c.options.TokenHandler.TokenLoader(userID)
	} else {
		err = errors.New("No Token Cache!")
	}
	return
}

func (c *client) GiveTokenPermissions(authURL string) (authcode string, err error) {
	// add logic for handing retrieving code for oauth exchange and matching state
	// For example throw an error, and send email to user instead with this link

	if c.options.TokenHandler.TokenPermissions != nil {
		if authcode, err = c.options.TokenHandler.TokenPermissions(authURL); err != nil {
			log.Fatal(err)
		}
	} else {
		fmt.Println("Please go to the following URL : ", authURL)
		if _, err := fmt.Scan(&authcode); err != nil {
			log.Fatal(err)
		}
	}

	// Use the authorization code that is pushed to the redirect URL.
	// NewTransportWithCode will do the handshake to retrieve
	// an access token and initiate a Transport that is
	// authorized and authenticated by the retrieved token.
	return
}

// UserOAuthClient returns an oauth2 client for a specific user
func (c *client) UserOAuthClient(ctx context.Context, config *oauth2.Config, userID string) (client *http.Client, err error) {
	var userToken *oauth2.Token

	if userToken, err = c.GetCachedToken(userID); err != nil {
		// if token for user is not cached then go through oauth2 flow
		if userToken, err = c.newUserToken(ctx, config, userID); err != nil {
			return
		}
	}

	if !userToken.Valid() { // if user token is expired
		userToken = &oauth2.Token{RefreshToken: userToken.RefreshToken}
	}

	return config.Client(ctx, userToken), err
}

// TokenOAuthClient returns an oauth2 client for a specific token
func (c *client) TokenOAuthClient(ctx context.Context, config *oauth2.Config, userToken *oauth2.Token) (client *http.Client, err error) {

	if !userToken.Valid() { // if user token is expired
		userToken = &oauth2.Token{RefreshToken: userToken.RefreshToken}
	}

	return config.Client(ctx, userToken), err
}

func (c *client) newUserToken(ctx context.Context, config *oauth2.Config, userID string) (*oauth2.Token, error) {
	stateBytes := make([]byte, 32)
	_, err := rand.Read(stateBytes)
	if err != nil {
		log.Fatalf("Unable to read random bytes: %v", err)
		return nil, err
	}
	state := fmt.Sprintf("%x", stateBytes)
	authURL := config.AuthCodeURL(state, oauth2.AccessTypeOffline)

	authcode, err := c.GiveTokenPermissions(authURL)

	token, err := config.Exchange(oauth2.NoContext, authcode)
	if err != nil {
		log.Fatalf("Exchange error: %v", err)
		return nil, err
	}
	c.SaveToken(userID, token) // save token to datastore

	return token, nil

}

// GetUserIDForToken looks up the corresponding itembase UserID based on
// OAuth2 Token
func (c *client) GetUserIDForToken(token *oauth2.Token) (string, error) {

	log.Println("GetUserIDForToken", token)

	tokenRef := NewClient(c.root, token.AccessToken, c.options, nil)

	me, err := tokenRef.Me()
	if err != nil {
		log.Println("Could not run Me call:", me, err)
		return "", err
	}

	return me.UUID, nil
}

func (c *client) getUserToken(userID string) (token *oauth2.Token) {
	config := c.newConf()
	client, err := c.UserOAuthClient(oauth2.NoContext, config, userID)

	_, err = client.Get(c.me)
	if err == nil {
		fmt.Errorf("Fetch should return an error if no refresh token is set")
	}

	token, err = client.Transport.(*oauth2.Transport).Source.Token()

	if err != nil {
		log.Fatal("Exchange error: %v", err)
	}

	token, err = client.Transport.(*oauth2.Transport).Source.Token()

	if err != nil {
		if token, err = c.newUserToken(oauth2.NoContext, config, userID); err != nil {
			return
		}
		if err != nil {
			log.Fatal("Exchange error:", err)
		}
	} else {
		c.SaveToken(userID, token)
	}

	return
}
<commit_msg>standalone HandleOAuthCode handler added<commit_after>package itembase

import (
	"crypto/rand"
	"errors"
	"fmt"
	"log"
	"net/http"

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

func (c *client) newConf() *oauth2.Config {
	var endpointURL string

	if c.production {
		endpointURL = "https://accounts.itembase.com/oauth/v2"
		c.me = "https://users.itembase.com/v1/me"
		c.root = "https://api.itembase.io/v1"
		c.activation = "https://solutionservice.itembase.com"
	} else {
		endpointURL = "http://sandbox.accounts.itembase.io/oauth/v2"
		c.me = "http://sandbox.users.itembase.io/v1/me"
		c.root = "http://sandbox.api.itembase.io/v1"
		c.activation = "http://sandbox.solutionservice.itembase.io"
	}

	return &oauth2.Config{
		ClientID:     c.options.ClientID,
		ClientSecret: c.options.ClientSecret,
		Scopes:       c.options.Scopes,
		RedirectURL:  c.options.RedirectURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  endpointURL + "/auth",
			TokenURL: endpointURL + "/token",
		},
	}
}

func (c *client) SaveToken(userID string, token *oauth2.Token) (err error) {
	if c.options.TokenHandler.TokenSaver != nil {
		err = c.options.TokenHandler.TokenSaver(userID, token)
	} else {
		err = errors.New("No Token Store!")
	}
	return
}

func (c *client) GetCachedToken(userID string) (token *oauth2.Token, err error) {
	if c.options.TokenHandler.TokenLoader != nil {
		token, err = c.options.TokenHandler.TokenLoader(userID)
	} else {
		err = errors.New("No Token Cache!")
	}
	return
}

func (c *client) GiveTokenPermissions(authURL string) (authcode string, err error) {
	// add logic for handing retrieving code for oauth exchange and matching state
	// For example throw an error, and send email to user instead with this link

	if c.options.TokenHandler.TokenPermissions != nil {
		if authcode, err = c.options.TokenHandler.TokenPermissions(authURL); err != nil {
			log.Fatal(err)
		}
	} else {
		fmt.Println("Please go to the following URL : ", authURL)
		if _, err := fmt.Scan(&authcode); err != nil {
			log.Fatal(err)
		}
	}

	// Use the authorization code that is pushed to the redirect URL.
	// NewTransportWithCode will do the handshake to retrieve
	// an access token and initiate a Transport that is
	// authorized and authenticated by the retrieved token.
	return
}

// UserOAuthClient returns an oauth2 client for a specific user
func (c *client) UserOAuthClient(ctx context.Context, config *oauth2.Config, userID string) (client *http.Client, err error) {
	var userToken *oauth2.Token

	if userToken, err = c.GetCachedToken(userID); err != nil {
		// if token for user is not cached then go through oauth2 flow
		if userToken, err = c.newUserToken(ctx, config, userID); err != nil {
			return
		}
	}

	if !userToken.Valid() { // if user token is expired
		userToken = &oauth2.Token{RefreshToken: userToken.RefreshToken}
	}

	return config.Client(ctx, userToken), err
}

// TokenOAuthClient returns an oauth2 client for a specific token
func (c *client) TokenOAuthClient(ctx context.Context, config *oauth2.Config, userToken *oauth2.Token) (client *http.Client, err error) {

	if !userToken.Valid() { // if user token is expired
		userToken = &oauth2.Token{RefreshToken: userToken.RefreshToken}
	}

	return config.Client(ctx, userToken), err
}

func (c *client) newUserToken(ctx context.Context, config *oauth2.Config, userID string) (*oauth2.Token, error) {
	stateBytes := make([]byte, 32)
	_, err := rand.Read(stateBytes)
	if err != nil {
		log.Fatalf("Unable to read random bytes: %v", err)
		return nil, err
	}
	state := fmt.Sprintf("%x", stateBytes)
	authURL := config.AuthCodeURL(state, oauth2.AccessTypeOffline)

	authcode, err := c.GiveTokenPermissions(authURL)
	if err != nil {
		log.Fatalf("Error when getting token permissions: %v", err)
		return nil, err
	}

	token, err := c.HandleOAuthCode(authcode)
	if err != nil {
		log.Fatalf("Error when handling OAuth Code error: %v", err)
		return nil, err
	}

	c.SaveToken(userID, token) // save token to datastore

	return token, nil
}

// HandleOAuthCode returns a valid token for an OAuth code
func (c *client) HandleOAuthCode(authcode string) (*oauth2.Token, error) {

	config := c.newConf()
	token, err := config.Exchange(oauth2.NoContext, authcode)

	if err != nil {
		log.Fatalf("Exchange error: %v", err)
		return nil, err
	}

	return token, nil

}

// GetUserIDForToken looks up the corresponding itembase UserID based on
// OAuth2 Token
func (c *client) GetUserIDForToken(token *oauth2.Token) (string, error) {

	log.Println("GetUserIDForToken", token)

	tokenRef := NewClient(c.root, token.AccessToken, c.options, nil)

	me, err := tokenRef.Me()
	if err != nil {
		log.Println("Could not run Me call:", me, err)
		return "", err
	}

	return me.UUID, nil
}

func (c *client) getUserToken(userID string) (token *oauth2.Token) {
	config := c.newConf()
	client, err := c.UserOAuthClient(oauth2.NoContext, config, userID)

	_, err = client.Get(c.me)
	if err == nil {
		fmt.Errorf("Fetch should return an error if no refresh token is set")
	}

	token, err = client.Transport.(*oauth2.Transport).Source.Token()

	if err != nil {
		log.Fatal("Exchange error: %v", err)
	}

	token, err = client.Transport.(*oauth2.Transport).Source.Token()

	if err != nil {
		if token, err = c.newUserToken(oauth2.NoContext, config, userID); err != nil {
			return
		}
		if err != nil {
			log.Fatal("Exchange error:", err)
		}
	} else {
		c.SaveToken(userID, token)
	}

	return
}
<|endoftext|>


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

// +build ignore

package main

import (
	"bufio"
	"bytes"
	"flag"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
)

func main() {
	log.SetPrefix("ci: ")
	log.SetFlags(0)

	var (
		race    = flag.Bool("race", false, "enable race detector")
		cover   = flag.String("coverpkg", "", "apply coverage analysis in each test to packages matching the patterns.")
		tags    = flag.String("tags", "", "build tags")
		verbose = flag.Bool("v", false, "enable verbose output")
	)

	flag.Parse()

	out := new(bytes.Buffer)
	cmd := exec.Command("go", "list", "./...")
	cmd.Stdout = out
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin

	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	f, err := os.Create("coverage.txt")
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	args := []string{"test"}

	if *verbose {
		args = append(args, "-v")
	}
	if *cover != "" {
		args = append(args, "-coverprofile=profile.out", "-covermode=atomic", "-coverpkg="+*cover)
	}
	if *tags != "" {
		args = append(args, "-tags="+*tags)
	}
	if *race {
		args = append(args, "-race")
	}
	args = append(args, "")

	scan := bufio.NewScanner(out)
	for scan.Scan() {
		pkg := scan.Text()
		if strings.Contains(pkg, "vendor") {
			continue
		}
		args[len(args)-1] = pkg
		cmd := exec.Command("go", args...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		err := cmd.Run()
		if err != nil {
			log.Fatal(err)
		}
		if *cover != "" {
			profile, err := ioutil.ReadFile("profile.out")
			if err != nil {
				log.Fatal(err)
			}
			_, err = f.Write(profile)
			if err != nil {
				log.Fatal(err)
			}
			os.Remove("profile.out")
		}
	}

	err = f.Close()
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>ci: display elapsed time, refactor pkg-list<commit_after>// Copyright 2018 The go-hep Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build ignore

package main

import (
	"bufio"
	"bytes"
	"flag"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/pkg/errors"
)

func main() {
	log.SetPrefix("ci: ")
	log.SetFlags(0)

	start := time.Now()
	defer func() {
		log.Printf("elapsed time: %v\n", time.Since(start))
	}()

	var (
		race    = flag.Bool("race", false, "enable race detector")
		cover   = flag.String("coverpkg", "", "apply coverage analysis in each test to packages matching the patterns.")
		tags    = flag.String("tags", "", "build tags")
		verbose = flag.Bool("v", false, "enable verbose output")
	)

	flag.Parse()

	pkgs, err := pkgList()
	if err != nil {
		log.Fatal(err)
	}

	f, err := os.Create("coverage.txt")
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	args := []string{"test"}

	if *verbose {
		args = append(args, "-v")
	}
	if *cover != "" {
		args = append(args, "-coverprofile=profile.out", "-covermode=atomic", "-coverpkg="+*cover)
	}
	if *tags != "" {
		args = append(args, "-tags="+*tags)
	}
	if *race {
		args = append(args, "-race")
	}
	args = append(args, "")

	for _, pkg := range pkgs {
		args[len(args)-1] = pkg
		cmd := exec.Command("go", args...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		err := cmd.Run()
		if err != nil {
			log.Fatal(err)
		}
		if *cover != "" {
			profile, err := ioutil.ReadFile("profile.out")
			if err != nil {
				log.Fatal(err)
			}
			_, err = f.Write(profile)
			if err != nil {
				log.Fatal(err)
			}
			os.Remove("profile.out")
		}
	}

	err = f.Close()
	if err != nil {
		log.Fatal(err)
	}
}

func pkgList() ([]string, error) {
	out := new(bytes.Buffer)
	cmd := exec.Command("go", "list", "./...")
	cmd.Stdout = out
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin

	err := cmd.Run()
	if err != nil {
		return nil, errors.Wrapf(err, "could not get package list")
	}

	var pkgs []string
	scan := bufio.NewScanner(out)
	for scan.Scan() {
		pkg := scan.Text()
		if strings.Contains(pkg, "vendor") {
			continue
		}
		pkgs = append(pkgs, pkg)
	}

	return pkgs, nil
}
<|endoftext|>


<commit_before>package auth

import (
	"encoding/json"
	"fmt"
	"github.com/timeredbull/tsuru/db"
	"github.com/timeredbull/tsuru/errors"
	"github.com/timeredbull/tsuru/gitosis"
	"github.com/timeredbull/tsuru/log"
	"io"
	"io/ioutil"
	"launchpad.net/mgo/bson"
	"net/http"
	"strings"
)

func CreateUser(w http.ResponseWriter, r *http.Request) error {
	var u User
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, &u)
	if err != nil {
		return &errors.Http{Code: http.StatusBadRequest, Message: err.Error()}
	}
	err = u.Create()
	if err == nil {
		w.WriteHeader(http.StatusCreated)
		return nil
	}

	if u.Get() == nil {
		err = &errors.Http{Code: http.StatusConflict, Message: "This email is already registered"}
	}

	return err
}

func Login(w http.ResponseWriter, r *http.Request) error {
	var pass map[string]string
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, &pass)
	if err != nil {
		return &errors.Http{Code: http.StatusBadRequest, Message: "Invalid JSON"}
	}

	password, ok := pass["password"]
	if !ok {
		msg := "You must provide a password to login"
		return &errors.Http{Code: http.StatusBadRequest, Message: msg}
	}

	u := User{Email: r.URL.Query().Get(":email")}
	err = u.Get()
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: "User not found"}
	}

	if u.Login(password) {
		t, _ := u.CreateToken()
		fmt.Fprintf(w, `{"token":"%s"}`, t.Token)
		return nil
	}

	msg := "Authentication failed, wrong password"
	return &errors.Http{Code: http.StatusUnauthorized, Message: msg}
}

func applyChangesToKeys(kind int, team *Team, user *User) {
	for _, key := range user.Keys {
		log.Print("adding user ", key.Name, " to ", team.Name)
		ch := gitosis.Change{
			Kind:     kind,
			Args:     map[string]string{"group": team.Name, "member": key.Name},
			Response: make(chan string),
		}
		gitosis.Changes <- ch
		fmt.Println(<-ch.Response)
	}
}

func createTeam(name string, u *User) error {
	team := &Team{Name: name, Users: []*User{u}}
	err := db.Session.Teams().Insert(team)
	if err != nil && strings.Contains(err.Error(), "duplicate key error") {
		return &errors.Http{Code: http.StatusConflict, Message: "This team already exists"}
	}
	ch := gitosis.Change{
		Kind:     gitosis.AddGroup,
		Args:     map[string]string{"group": name},
		Response: make(chan string),
	}
	gitosis.Changes <- ch
	<-ch.Response
	applyChangesToKeys(gitosis.AddMember, team, u)
	return nil
}

func CreateTeam(w http.ResponseWriter, r *http.Request, u *User) error {
	var params map[string]string
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, &params)
	if err != nil {
		return &errors.Http{Code: http.StatusBadRequest, Message: err.Error()}
	}
	name, ok := params["name"]
	if !ok {
		msg := "You must provide the team name"
		return &errors.Http{Code: http.StatusBadRequest, Message: msg}
	}
	return createTeam(name, u)
}

func addUserToTeam(email, teamName string, u *User) error {
	team, user := new(Team), new(User)
	selector := bson.M{"name": teamName}
	err := db.Session.Teams().Find(selector).One(team)
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: "Team not found"}
	}
	if !team.ContainsUser(u) {
		msg := fmt.Sprintf("You are not authorized to add new users to the team %s", team.Name)
		return &errors.Http{Code: http.StatusUnauthorized, Message: msg}
	}
	err = db.Session.Users().Find(bson.M{"email": email}).One(user)
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: "User not found"}
	}
	err = team.AddUser(user)
	if err != nil {
		return &errors.Http{Code: http.StatusConflict, Message: err.Error()}
	}
	err = db.Session.Teams().Update(selector, team)
	if err != nil {
		return err
	}
	applyChangesToKeys(gitosis.AddMember, team, user)
	return nil
}

func AddUserToTeam(w http.ResponseWriter, r *http.Request, u *User) error {
	team := r.URL.Query().Get(":team")
	email := r.URL.Query().Get(":user")
	return addUserToTeam(email, team, u)
}

func removeUserFromTeam(email, teamName string, u *User) error {
	team := new(Team)
	selector := bson.M{"name": teamName}
	err := db.Session.Teams().Find(selector).One(team)
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: "Team not found"}
	}
	if !team.ContainsUser(u) {
		msg := fmt.Sprintf("You are not authorized to remove a member from the team %s", team.Name)
		return &errors.Http{Code: http.StatusUnauthorized, Message: msg}
	}
	if len(team.Users) == 1 {
		msg := "You can not remove this user from this team, because it is the last user within the team, and a team can not be orphaned"
		return &errors.Http{Code: http.StatusForbidden, Message: msg}
	}
	user := User{Email: email}
	err = user.Get()
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: err.Error()}
	}
	err = team.RemoveUser(&user)
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: err.Error()}
	}
	err = db.Session.Teams().Update(selector, team)
	if err != nil {
		return err
	}
	applyChangesToKeys(gitosis.RemoveMember, team, &user)
	return nil
}

func RemoveUserFromTeam(w http.ResponseWriter, r *http.Request, u *User) error {
	email := r.URL.Query().Get(":user")
	team := r.URL.Query().Get(":team")
	return removeUserFromTeam(email, team, u)
}

func getKeyFromBody(b io.Reader) (string, error) {
	var body map[string]string
	content, err := ioutil.ReadAll(b)
	if err != nil {
		return "", err
	}
	err = json.Unmarshal(content, &body)
	if err != nil {
		return "", &errors.Http{Code: http.StatusBadRequest, Message: "Invalid JSON"}
	}
	key, ok := body["key"]
	if !ok || key == "" {
		return "", &errors.Http{Code: http.StatusBadRequest, Message: "Missing key"}
	}
	return key, nil
}

func addKeyToUser(content string, u *User) error {
	key := Key{Content: content}
	if u.hasKey(key) {
		return &errors.Http{Code: http.StatusConflict, Message: "User has this key already"}
	}
	r := make(chan string)
	ch := gitosis.Change{
		Kind:     gitosis.AddKey,
		Args:     map[string]string{"member": u.Email, "key": content},
		Response: r,
	}
	gitosis.Changes <- ch
	var teams []Team
	db.Session.Teams().Find(bson.M{"users.email": u.Email}).All(&teams)
	key.Name = strings.Replace(<-r, ".pub", "", -1)
	for _, team := range teams {
		mch := gitosis.Change{
			Kind: gitosis.AddMember,
			Args: map[string]string{"group": team.Name, "member": key.Name},
		}
		gitosis.Changes <- mch
	}
	u.addKey(key)
	return db.Session.Users().Update(bson.M{"email": u.Email}, u)
}

// AddKeyToUser adds a key to a user.
//
// This function is just an http wrapper around addKeyToUser. The latter function
// exists to be used in other places in the package without the http stuff (request and
// response).
func AddKeyToUser(w http.ResponseWriter, r *http.Request, u *User) error {
	key, err := getKeyFromBody(r.Body)
	if err != nil {
		return err
	}
	return addKeyToUser(key, u)
}

func removeKeyFromUser(content string, u *User) error {
	key, index := u.findKey(Key{Content: content})
	if index < 0 {
		return &errors.Http{Code: http.StatusNotFound, Message: "User does not have this key"}
	}
	u.removeKey(key)
	err := db.Session.Users().Update(bson.M{"email": u.Email}, u)
	if err != nil {
		return err
	}
	ch := gitosis.Change{
		Kind: gitosis.RemoveKey,
		Args: map[string]string{"key": key.Name + ".pub"},
	}
	gitosis.Changes <- ch
	var teams []Team
	db.Session.Teams().Find(bson.M{"users.email": u.Email}).All(&teams)
	for _, team := range teams {
		mch := gitosis.Change{
			Kind: gitosis.RemoveMember,
			Args: map[string]string{"group": team.Name, "member": key.Name},
		}
		gitosis.Changes <- mch
	}
	return nil
}

// RemoveKeyFromUser removes a key from a user.
//
// This function is just an http wrapper around removeKeyFromUser. The latter function
// exists to be used in other places in the package without the http stuff (request and
// response).
func RemoveKeyFromUser(w http.ResponseWriter, r *http.Request, u *User) error {
	key, err := getKeyFromBody(r.Body)
	if err != nil {
		return err
	}
	return removeKeyFromUser(key, u)
}
<commit_msg>api/auth: removed channel receive to make code asynchronous again<commit_after>package auth

import (
	"encoding/json"
	"fmt"
	"github.com/timeredbull/tsuru/db"
	"github.com/timeredbull/tsuru/errors"
	"github.com/timeredbull/tsuru/gitosis"
	"github.com/timeredbull/tsuru/log"
	"io"
	"io/ioutil"
	"launchpad.net/mgo/bson"
	"net/http"
	"strings"
)

func CreateUser(w http.ResponseWriter, r *http.Request) error {
	var u User
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, &u)
	if err != nil {
		return &errors.Http{Code: http.StatusBadRequest, Message: err.Error()}
	}
	err = u.Create()
	if err == nil {
		w.WriteHeader(http.StatusCreated)
		return nil
	}

	if u.Get() == nil {
		err = &errors.Http{Code: http.StatusConflict, Message: "This email is already registered"}
	}

	return err
}

func Login(w http.ResponseWriter, r *http.Request) error {
	var pass map[string]string
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, &pass)
	if err != nil {
		return &errors.Http{Code: http.StatusBadRequest, Message: "Invalid JSON"}
	}

	password, ok := pass["password"]
	if !ok {
		msg := "You must provide a password to login"
		return &errors.Http{Code: http.StatusBadRequest, Message: msg}
	}

	u := User{Email: r.URL.Query().Get(":email")}
	err = u.Get()
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: "User not found"}
	}

	if u.Login(password) {
		t, _ := u.CreateToken()
		fmt.Fprintf(w, `{"token":"%s"}`, t.Token)
		return nil
	}

	msg := "Authentication failed, wrong password"
	return &errors.Http{Code: http.StatusUnauthorized, Message: msg}
}

func applyChangesToKeys(kind int, team *Team, user *User) {
	for _, key := range user.Keys {
		log.Print("adding user ", key.Name, " to ", team.Name)
		ch := gitosis.Change{
			Kind: kind,
			Args: map[string]string{"group": team.Name, "member": key.Name},
		}
		gitosis.Changes <- ch
	}
}

func createTeam(name string, u *User) error {
	team := &Team{Name: name, Users: []*User{u}}
	err := db.Session.Teams().Insert(team)
	if err != nil && strings.Contains(err.Error(), "duplicate key error") {
		return &errors.Http{Code: http.StatusConflict, Message: "This team already exists"}
	}
	ch := gitosis.Change{
		Kind:     gitosis.AddGroup,
		Args:     map[string]string{"group": name},
		Response: make(chan string),
	}
	gitosis.Changes <- ch
	<-ch.Response
	applyChangesToKeys(gitosis.AddMember, team, u)
	return nil
}

func CreateTeam(w http.ResponseWriter, r *http.Request, u *User) error {
	var params map[string]string
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, &params)
	if err != nil {
		return &errors.Http{Code: http.StatusBadRequest, Message: err.Error()}
	}
	name, ok := params["name"]
	if !ok {
		msg := "You must provide the team name"
		return &errors.Http{Code: http.StatusBadRequest, Message: msg}
	}
	return createTeam(name, u)
}

func addUserToTeam(email, teamName string, u *User) error {
	team, user := new(Team), new(User)
	selector := bson.M{"name": teamName}
	err := db.Session.Teams().Find(selector).One(team)
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: "Team not found"}
	}
	if !team.ContainsUser(u) {
		msg := fmt.Sprintf("You are not authorized to add new users to the team %s", team.Name)
		return &errors.Http{Code: http.StatusUnauthorized, Message: msg}
	}
	err = db.Session.Users().Find(bson.M{"email": email}).One(user)
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: "User not found"}
	}
	err = team.AddUser(user)
	if err != nil {
		return &errors.Http{Code: http.StatusConflict, Message: err.Error()}
	}
	err = db.Session.Teams().Update(selector, team)
	if err != nil {
		return err
	}
	applyChangesToKeys(gitosis.AddMember, team, user)
	return nil
}

func AddUserToTeam(w http.ResponseWriter, r *http.Request, u *User) error {
	team := r.URL.Query().Get(":team")
	email := r.URL.Query().Get(":user")
	return addUserToTeam(email, team, u)
}

func removeUserFromTeam(email, teamName string, u *User) error {
	team := new(Team)
	selector := bson.M{"name": teamName}
	err := db.Session.Teams().Find(selector).One(team)
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: "Team not found"}
	}
	if !team.ContainsUser(u) {
		msg := fmt.Sprintf("You are not authorized to remove a member from the team %s", team.Name)
		return &errors.Http{Code: http.StatusUnauthorized, Message: msg}
	}
	if len(team.Users) == 1 {
		msg := "You can not remove this user from this team, because it is the last user within the team, and a team can not be orphaned"
		return &errors.Http{Code: http.StatusForbidden, Message: msg}
	}
	user := User{Email: email}
	err = user.Get()
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: err.Error()}
	}
	err = team.RemoveUser(&user)
	if err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: err.Error()}
	}
	err = db.Session.Teams().Update(selector, team)
	if err != nil {
		return err
	}
	applyChangesToKeys(gitosis.RemoveMember, team, &user)
	return nil
}

func RemoveUserFromTeam(w http.ResponseWriter, r *http.Request, u *User) error {
	email := r.URL.Query().Get(":user")
	team := r.URL.Query().Get(":team")
	return removeUserFromTeam(email, team, u)
}

func getKeyFromBody(b io.Reader) (string, error) {
	var body map[string]string
	content, err := ioutil.ReadAll(b)
	if err != nil {
		return "", err
	}
	err = json.Unmarshal(content, &body)
	if err != nil {
		return "", &errors.Http{Code: http.StatusBadRequest, Message: "Invalid JSON"}
	}
	key, ok := body["key"]
	if !ok || key == "" {
		return "", &errors.Http{Code: http.StatusBadRequest, Message: "Missing key"}
	}
	return key, nil
}

func addKeyToUser(content string, u *User) error {
	key := Key{Content: content}
	if u.hasKey(key) {
		return &errors.Http{Code: http.StatusConflict, Message: "User has this key already"}
	}
	r := make(chan string)
	ch := gitosis.Change{
		Kind:     gitosis.AddKey,
		Args:     map[string]string{"member": u.Email, "key": content},
		Response: r,
	}
	gitosis.Changes <- ch
	var teams []Team
	db.Session.Teams().Find(bson.M{"users.email": u.Email}).All(&teams)
	key.Name = strings.Replace(<-r, ".pub", "", -1)
	for _, team := range teams {
		mch := gitosis.Change{
			Kind: gitosis.AddMember,
			Args: map[string]string{"group": team.Name, "member": key.Name},
		}
		gitosis.Changes <- mch
	}
	u.addKey(key)
	return db.Session.Users().Update(bson.M{"email": u.Email}, u)
}

// AddKeyToUser adds a key to a user.
//
// This function is just an http wrapper around addKeyToUser. The latter function
// exists to be used in other places in the package without the http stuff (request and
// response).
func AddKeyToUser(w http.ResponseWriter, r *http.Request, u *User) error {
	key, err := getKeyFromBody(r.Body)
	if err != nil {
		return err
	}
	return addKeyToUser(key, u)
}

func removeKeyFromUser(content string, u *User) error {
	key, index := u.findKey(Key{Content: content})
	if index < 0 {
		return &errors.Http{Code: http.StatusNotFound, Message: "User does not have this key"}
	}
	u.removeKey(key)
	err := db.Session.Users().Update(bson.M{"email": u.Email}, u)
	if err != nil {
		return err
	}
	ch := gitosis.Change{
		Kind: gitosis.RemoveKey,
		Args: map[string]string{"key": key.Name + ".pub"},
	}
	gitosis.Changes <- ch
	var teams []Team
	db.Session.Teams().Find(bson.M{"users.email": u.Email}).All(&teams)
	for _, team := range teams {
		mch := gitosis.Change{
			Kind: gitosis.RemoveMember,
			Args: map[string]string{"group": team.Name, "member": key.Name},
		}
		gitosis.Changes <- mch
	}
	return nil
}

// RemoveKeyFromUser removes a key from a user.
//
// This function is just an http wrapper around removeKeyFromUser. The latter function
// exists to be used in other places in the package without the http stuff (request and
// response).
func RemoveKeyFromUser(w http.ResponseWriter, r *http.Request, u *User) error {
	key, err := getKeyFromBody(r.Body)
	if err != nil {
		return err
	}
	return removeKeyFromUser(key, u)
}
<|endoftext|>


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

package etcdv3

import (
	"fmt"
	"time"

	"github.com/ligato/cn-infra/core"
	"github.com/ligato/cn-infra/datasync/resync"
	"github.com/ligato/cn-infra/db/keyval/plugin"
	"github.com/ligato/cn-infra/flavors/local"
	"github.com/ligato/cn-infra/health/statuscheck"
	"github.com/ligato/cn-infra/servicelabel"
	"github.com/ligato/cn-infra/utils/safeclose"
)

const (
	// healthCheckProbeKey is a key used to probe Etcd state
	healthCheckProbeKey = "/probe-etcd-connection"
)

// Plugin implements etcdv3 plugin.
type Plugin struct {
	Deps // inject
	*plugin.Skeleton
	disabled        bool
	connection      *BytesConnectionEtcd
	autoCompactDone chan struct{}
	lastConnErr     error
}

// Deps lists dependencies of the etcdv3 plugin.
// If injected, etcd plugin will use StatusCheck to signal the connection status.
type Deps struct {
	local.PluginInfraDeps // inject
	Resync                *resync.Plugin
}

// Init retrieves etcd configuration and establishes a new connection
// with the etcd data store.
// If the configuration file doesn't exist or cannot be read, the returned error
// will be of os.PathError type. An untyped error is returned in case the file
// doesn't contain a valid YAML configuration.
// The function may also return error if TLS connection is selected and the
// CA or client certificate is not accessible(os.PathError)/valid(untyped).
// Check clientv3.New from coreos/etcd for possible errors returned in case
// the connection cannot be established.
func (p *Plugin) Init() (err error) {
	// Init connection
	if p.Skeleton == nil {
		// Retrieve config
		var cfg Config
		if found, err := p.PluginConfig.GetValue(&cfg); err != nil {
			return err
		} else if !found {
			p.Log.Info("etcd config not found ", p.PluginConfig.GetConfigName(), " - skip loading this plugin")
			p.disabled = true
			return nil
		}

		etcdConfig, err := ConfigToClientv3(&cfg)
		if err != nil {
			return err
		}

		p.connection, err = NewEtcdConnectionWithBytes(*etcdConfig, p.Log)
		if err != nil {
			return err
		}

		if cfg.AutoCompact > 0 {
			if cfg.AutoCompact < time.Duration(time.Minute*60) {
				p.Log.Warnf("auto compact option for ETCD is set to less than 60 minutes!")
			}
			p.startPeriodicAutoCompact(cfg.AutoCompact)
		}

		p.Skeleton = plugin.NewSkeleton(p.String(),
			p.ServiceLabel,
			p.connection,
		)
	}

	if err := p.Skeleton.Init(); err != nil {
		return err
	}

	// Register for providing status reports (polling mode).
	if p.StatusCheck != nil {
		p.StatusCheck.Register(core.PluginName(p.String()), func() (statuscheck.PluginState, error) {
			_, _, _, err := p.connection.GetValue(healthCheckProbeKey)
			if err == nil {
				if p.Resync != nil && p.lastConnErr != nil {
					p.Log.Info("Starting resync after ETCD reconnect")
					p.Resync.DoResync()
					p.lastConnErr = nil
				}
				return statuscheck.OK, nil
			}
			p.lastConnErr = err
			return statuscheck.Error, err
		})
	} else {
		p.Log.Warnf("Unable to start status check for etcd")
	}

	return nil
}

// AfterInit registers status polling function with StatusCheck plugin
// (if injected).
func (p *Plugin) AfterInit() error {
	return nil
}

func (p *Plugin) startPeriodicAutoCompact(period time.Duration) {
	p.autoCompactDone = make(chan struct{})
	go func() {
		p.Log.Infof("Starting periodic auto compacting every %v", period)
		for {
			select {
			case <-time.After(period):
				p.Log.Debugf("Executing auto compact")
				if toRev, err := p.connection.Compact(); err != nil {
					p.Log.Errorf("Periodic auto compacting failed: %v", err)
				} else {
					p.Log.Infof("Auto compacting finished (to revision %v)", toRev)
				}
			case <-p.autoCompactDone:
				return
			}
		}
	}()
}

// FromExistingConnection is used mainly for testing of existing connection
// injection into the plugin.
// Note, need to set Deps for returned value!
func FromExistingConnection(connection *BytesConnectionEtcd, sl servicelabel.ReaderAPI) *Plugin {
	skel := plugin.NewSkeleton("testing", sl, connection)
	return &Plugin{Skeleton: skel, connection: connection}
}

// Close shutdowns the connection.
func (p *Plugin) Close() error {
	_, err := safeclose.CloseAll(p.Skeleton, p.autoCompactDone)
	return err
}

// String returns the plugin name from dependencies if injected,
// "kvdbsync" otherwise.
func (p *Plugin) String() string {
	if len(p.Deps.PluginName) == 0 {
		return "kvdbsync"
	}
	return string(p.Deps.PluginName)
}

// Disabled returns *true* if the plugin is not in use due to missing
// etcd configuration.
func (p *Plugin) Disabled() (disabled bool) {
	return p.disabled
}

// PutIfNotExists puts given key-value pair into etcd if there is no value set for the key. If the put was successful
// succeeded is true. If the key already exists succeeded is false and the value for the key is untouched.
func (p *Plugin) PutIfNotExists(key string, value []byte) (succeeded bool, err error) {
	if p.connection != nil {
		return p.connection.PutIfNotExists(key, value)
	}
	return false, fmt.Errorf("connection is not established")
}

// Compact compatcs the ETCD database to the specific revision
func (p *Plugin) Compact(rev ...int64) (toRev int64, err error) {
	if p.connection != nil {
		return p.connection.Compact(rev...)
	}
	return 0, fmt.Errorf("connection is not established")
}
<commit_msg>Temporarly disable resync after etcd reconnect<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package etcdv3

import (
	"fmt"
	"time"

	"github.com/ligato/cn-infra/core"
	"github.com/ligato/cn-infra/datasync/resync"
	"github.com/ligato/cn-infra/db/keyval/plugin"
	"github.com/ligato/cn-infra/flavors/local"
	"github.com/ligato/cn-infra/health/statuscheck"
	"github.com/ligato/cn-infra/servicelabel"
	"github.com/ligato/cn-infra/utils/safeclose"
)

const (
	// healthCheckProbeKey is a key used to probe Etcd state
	healthCheckProbeKey = "/probe-etcd-connection"
)

// Plugin implements etcdv3 plugin.
type Plugin struct {
	Deps // inject
	*plugin.Skeleton
	disabled        bool
	connection      *BytesConnectionEtcd
	autoCompactDone chan struct{}
	lastConnErr     error
}

// Deps lists dependencies of the etcdv3 plugin.
// If injected, etcd plugin will use StatusCheck to signal the connection status.
type Deps struct {
	local.PluginInfraDeps // inject
	Resync                *resync.Plugin
}

// Init retrieves etcd configuration and establishes a new connection
// with the etcd data store.
// If the configuration file doesn't exist or cannot be read, the returned error
// will be of os.PathError type. An untyped error is returned in case the file
// doesn't contain a valid YAML configuration.
// The function may also return error if TLS connection is selected and the
// CA or client certificate is not accessible(os.PathError)/valid(untyped).
// Check clientv3.New from coreos/etcd for possible errors returned in case
// the connection cannot be established.
func (p *Plugin) Init() (err error) {
	// Init connection
	if p.Skeleton == nil {
		// Retrieve config
		var cfg Config
		if found, err := p.PluginConfig.GetValue(&cfg); err != nil {
			return err
		} else if !found {
			p.Log.Info("etcd config not found ", p.PluginConfig.GetConfigName(), " - skip loading this plugin")
			p.disabled = true
			return nil
		}

		etcdConfig, err := ConfigToClientv3(&cfg)
		if err != nil {
			return err
		}

		p.connection, err = NewEtcdConnectionWithBytes(*etcdConfig, p.Log)
		if err != nil {
			return err
		}

		if cfg.AutoCompact > 0 {
			if cfg.AutoCompact < time.Duration(time.Minute*60) {
				p.Log.Warnf("auto compact option for ETCD is set to less than 60 minutes!")
			}
			p.startPeriodicAutoCompact(cfg.AutoCompact)
		}

		p.Skeleton = plugin.NewSkeleton(p.String(),
			p.ServiceLabel,
			p.connection,
		)
	}

	if err := p.Skeleton.Init(); err != nil {
		return err
	}

	// Register for providing status reports (polling mode).
	if p.StatusCheck != nil {
		p.StatusCheck.Register(core.PluginName(p.String()), func() (statuscheck.PluginState, error) {
			_, _, _, err := p.connection.GetValue(healthCheckProbeKey)
			if err == nil {
				if p.Resync != nil && p.lastConnErr != nil {
					p.Log.Info("Starting resync after ETCD reconnect")
					//p.Resync.DoResync()
					p.lastConnErr = nil
				}
				return statuscheck.OK, nil
			}
			p.lastConnErr = err
			return statuscheck.Error, err
		})
	} else {
		p.Log.Warnf("Unable to start status check for etcd")
	}

	return nil
}

// AfterInit registers status polling function with StatusCheck plugin
// (if injected).
func (p *Plugin) AfterInit() error {
	return nil
}

func (p *Plugin) startPeriodicAutoCompact(period time.Duration) {
	p.autoCompactDone = make(chan struct{})
	go func() {
		p.Log.Infof("Starting periodic auto compacting every %v", period)
		for {
			select {
			case <-time.After(period):
				p.Log.Debugf("Executing auto compact")
				if toRev, err := p.connection.Compact(); err != nil {
					p.Log.Errorf("Periodic auto compacting failed: %v", err)
				} else {
					p.Log.Infof("Auto compacting finished (to revision %v)", toRev)
				}
			case <-p.autoCompactDone:
				return
			}
		}
	}()
}

// FromExistingConnection is used mainly for testing of existing connection
// injection into the plugin.
// Note, need to set Deps for returned value!
func FromExistingConnection(connection *BytesConnectionEtcd, sl servicelabel.ReaderAPI) *Plugin {
	skel := plugin.NewSkeleton("testing", sl, connection)
	return &Plugin{Skeleton: skel, connection: connection}
}

// Close shutdowns the connection.
func (p *Plugin) Close() error {
	_, err := safeclose.CloseAll(p.Skeleton, p.autoCompactDone)
	return err
}

// String returns the plugin name from dependencies if injected,
// "kvdbsync" otherwise.
func (p *Plugin) String() string {
	if len(p.Deps.PluginName) == 0 {
		return "kvdbsync"
	}
	return string(p.Deps.PluginName)
}

// Disabled returns *true* if the plugin is not in use due to missing
// etcd configuration.
func (p *Plugin) Disabled() (disabled bool) {
	return p.disabled
}

// PutIfNotExists puts given key-value pair into etcd if there is no value set for the key. If the put was successful
// succeeded is true. If the key already exists succeeded is false and the value for the key is untouched.
func (p *Plugin) PutIfNotExists(key string, value []byte) (succeeded bool, err error) {
	if p.connection != nil {
		return p.connection.PutIfNotExists(key, value)
	}
	return false, fmt.Errorf("connection is not established")
}

// Compact compatcs the ETCD database to the specific revision
func (p *Plugin) Compact(rev ...int64) (toRev int64, err error) {
	if p.connection != nil {
		return p.connection.Compact(rev...)
	}
	return 0, fmt.Errorf("connection is not established")
}
<|endoftext|>


<commit_before>package auth

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/micro/go-micro/v2/util/ctx"

	"github.com/micro/go-micro/v2/api/resolver"
	"github.com/micro/go-micro/v2/api/server"
	"github.com/micro/go-micro/v2/auth"
	"github.com/micro/go-micro/v2/logger"
	"github.com/micro/micro/v2/internal/namespace"
)

// Wrapper wraps a handler and authenticates requests
func Wrapper(r resolver.Resolver, nr *namespace.Resolver) server.Wrapper {
	return func(h http.Handler) http.Handler {
		return authWrapper{
			handler:    h,
			resolver:   r,
			nsResolver: nr,
			auth:       auth.DefaultAuth,
		}
	}
}

type authWrapper struct {
	handler    http.Handler
	auth       auth.Auth
	resolver   resolver.Resolver
	nsResolver *namespace.Resolver
}

func (a authWrapper) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// Determine the namespace and set it in the header
	ns := a.nsResolver.Resolve(req)
	req.Header.Set(namespace.NamespaceKey, ns)

	// Set the metadata so we can access it in micro api / web
	req = req.WithContext(ctx.FromRequest(req))

	// Extract the token from the request
	var token string
	if header := req.Header.Get("Authorization"); len(header) > 0 {
		// Extract the auth token from the request
		if strings.HasPrefix(header, auth.BearerScheme) {
			token = header[len(auth.BearerScheme):]
		}
	} else {
		// Get the token out the cookies if not provided in headers
		if c, err := req.Cookie("micro-token"); err == nil && c != nil {
			token = strings.TrimPrefix(c.Value, auth.TokenCookieName+"=")
			req.Header.Set("Authorization", auth.BearerScheme+token)
		}
	}

	// Get the account using the token, fallback to a blank account
	// since some endpoints can be unauthenticated, so the lack of an
	// account doesn't necesserially mean a forbidden request
	acc, err := a.auth.Inspect(token)
	if err != nil {
		acc = &auth.Account{}
	}

	// Determine the name of the service being requested
	endpoint, err := a.resolver.Resolve(req)
	if err == resolver.ErrInvalidPath || err == resolver.ErrNotFound {
		// a file not served by the resolver has been requested (e.g. favicon.ico)
		endpoint = &resolver.Endpoint{Path: req.URL.Path}
	} else if err != nil {
		logger.Error(err)
		http.Error(w, err.Error(), 500)
		return
	} else {
		// set the endpoint in the context so it can be used to resolve
		// the request later
		ctx := context.WithValue(req.Context(), resolver.Endpoint{}, endpoint)
		*req = *req.Clone(ctx)
	}

	// construct the resource name, e.g. home => go.micro.web.home
	resName := ns
	if len(endpoint.Name) > 0 {
		resName = ns + "." + endpoint.Name
	}

	// determine the resource path. there is an inconsistency in how resolvers
	// use method, some use it as Users.ReadUser (the rpc method), and others
	// use it as the HTTP method, e.g GET. TODO: Refactor this to make it consistent.
	resEndpoint := endpoint.Path
	if len(endpoint.Path) == 0 {
		resEndpoint = endpoint.Method
	}

	// Perform the verification check to see if the account has access to
	// the resource they're requesting
	res := &auth.Resource{Type: "service", Name: resName, Endpoint: resEndpoint, Namespace: ns}
	if err := a.auth.Verify(acc, res); err == nil {
		// The account has the necessary permissions to access the resource
		a.handler.ServeHTTP(w, req)
		return
	}

	// The account is set, but they don't have enough permissions, hence
	// we return a forbidden error.
	if len(acc.ID) > 0 {
		http.Error(w, "Forbidden request", 403)
		return
	}

	// If there is no auth login url set, 401
	loginURL := a.auth.Options().LoginURL
	if loginURL == "" {
		http.Error(w, "unauthorized request", 401)
		return
	}

	// Redirect to the login path
	params := url.Values{"redirect_to": {req.URL.String()}}
	loginWithRedirect := fmt.Sprintf("%v?%v", loginURL, params.Encode())
	http.Redirect(w, req, loginWithRedirect, http.StatusTemporaryRedirect)
}
<commit_msg>Fix auth web redirect bug (#795)<commit_after>package auth

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/micro/go-micro/v2/util/ctx"

	"github.com/micro/go-micro/v2/api/resolver"
	"github.com/micro/go-micro/v2/api/server"
	"github.com/micro/go-micro/v2/auth"
	"github.com/micro/go-micro/v2/logger"
	"github.com/micro/micro/v2/internal/namespace"
)

// Wrapper wraps a handler and authenticates requests
func Wrapper(r resolver.Resolver, nr *namespace.Resolver) server.Wrapper {
	return func(h http.Handler) http.Handler {
		return authWrapper{
			handler:    h,
			resolver:   r,
			nsResolver: nr,
			auth:       auth.DefaultAuth,
		}
	}
}

type authWrapper struct {
	handler    http.Handler
	auth       auth.Auth
	resolver   resolver.Resolver
	nsResolver *namespace.Resolver
}

func (a authWrapper) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// Determine the namespace and set it in the header
	ns := a.nsResolver.Resolve(req)
	req.Header.Set(namespace.NamespaceKey, ns)

	// Set the metadata so we can access it in micro api / web
	req = req.WithContext(ctx.FromRequest(req))

	// Extract the token from the request
	var token string
	if header := req.Header.Get("Authorization"); len(header) > 0 {
		// Extract the auth token from the request
		if strings.HasPrefix(header, auth.BearerScheme) {
			token = header[len(auth.BearerScheme):]
		}
	} else {
		// Get the token out the cookies if not provided in headers
		if c, err := req.Cookie("micro-token"); err == nil && c != nil {
			token = strings.TrimPrefix(c.Value, auth.TokenCookieName+"=")
			req.Header.Set("Authorization", auth.BearerScheme+token)
		}
	}

	// Get the account using the token, fallback to a blank account
	// since some endpoints can be unauthenticated, so the lack of an
	// account doesn't necesserially mean a forbidden request
	acc, err := a.auth.Inspect(token)
	if err != nil {
		acc = &auth.Account{}
	}

	// Determine the name of the service being requested
	endpoint, err := a.resolver.Resolve(req)
	if err == resolver.ErrInvalidPath || err == resolver.ErrNotFound {
		// a file not served by the resolver has been requested (e.g. favicon.ico)
		endpoint = &resolver.Endpoint{Path: req.URL.Path}
	} else if err != nil {
		logger.Error(err)
		http.Error(w, err.Error(), 500)
		return
	} else {
		// set the endpoint in the context so it can be used to resolve
		// the request later
		ctx := context.WithValue(req.Context(), resolver.Endpoint{}, endpoint)
		*req = *req.Clone(ctx)
	}

	// construct the resource name, e.g. home => go.micro.web.home
	resName := a.nsResolver.ResolveWithType(req)
	if len(endpoint.Name) > 0 {
		resName = resName + "." + endpoint.Name
	}

	// determine the resource path. there is an inconsistency in how resolvers
	// use method, some use it as Users.ReadUser (the rpc method), and others
	// use it as the HTTP method, e.g GET. TODO: Refactor this to make it consistent.
	resEndpoint := endpoint.Path
	if len(endpoint.Path) == 0 {
		resEndpoint = endpoint.Method
	}

	// Perform the verification check to see if the account has access to
	// the resource they're requesting
	res := &auth.Resource{Type: "service", Name: resName, Endpoint: resEndpoint, Namespace: ns}
	if err := a.auth.Verify(acc, res); err == nil {
		// The account has the necessary permissions to access the resource
		a.handler.ServeHTTP(w, req)
		return
	}

	// The account is set, but they don't have enough permissions, hence
	// we return a forbidden error.
	if len(acc.ID) > 0 {
		http.Error(w, "Forbidden request", 403)
		return
	}

	// If there is no auth login url set, 401
	loginURL := a.auth.Options().LoginURL
	if loginURL == "" {
		http.Error(w, "unauthorized request", 401)
		return
	}

	// Redirect to the login path
	params := url.Values{"redirect_to": {req.URL.String()}}
	loginWithRedirect := fmt.Sprintf("%v?%v", loginURL, params.Encode())
	http.Redirect(w, req, loginWithRedirect, http.StatusTemporaryRedirect)
}
<|endoftext|>


<commit_before>package elasticsearch

import (
	"context"
	"errors"
	"fmt"

	"github.com/korfuri/goref"
	log "github.com/sirupsen/logrus"
	elastic "gopkg.in/olivere/elastic.v5"
)

const (
	// Max number of errors reported in one call to
	// LoadGraphToElastic
	maxErrorsReported = 20

	// Types in the Elastic search index
	packageType = "package"
	refType     = "ref"
)

// PackageExists returns whether the provided loadpath + version tuple
// exists in this index.
func PackageExists(loadpath string, version int64, client *elastic.Client, index string) bool {
	ctx := context.Background()
	docID := fmt.Sprintf("v1@%d@%s", version, loadpath)
	pkgDoc, _ := client.Get().
		Index(index).
		Type(packageType).
		Id(docID).
		Do(ctx)
	// TODO: handle errors better. Right now we assume that any
	// error is a 404 and can be ignored safely.
	return pkgDoc != nil
}

// LoadGraphToElastic loads all Packages and Refs from a PackageGraph
// to the provided ES index.
func LoadGraphToElastic(pg goref.PackageGraph, client *elastic.Client, index string) ([]*goref.Ref, error) {
	ctx := context.Background()
	missedRefs := make([]*goref.Ref, 0)
	errs := make([]error, 0)

	for _, p := range pg.Packages {
		log.Infof("Processing package %s", p.Path)

		if PackageExists(p.Path, p.Version, client, index) {
			log.Infof("Package %s already exists in this index.", p)
			continue
		}

		log.Infof("Creating Package %s in the index", p)
		if _, err := client.Index().
			Index(index).
			Type(packageType).
			Id(p.DocumentID()).
			BodyJson(p).
			Do(ctx); err != nil {
			return nil, err
		}

		for _, r := range p.OutRefs {
			log.Infof("Creating Ref document [%s] in the index", r)
			refDoc, err := client.Index().
				Index(index).
				Type(refType).
				BodyJson(r).
				Do(ctx)
			if err != nil {
				missedRefs = append(missedRefs, r)
				errs = append(errs, err)
				log.Infof("Create Ref document failed with err:[%s] for Ref:[%s]", err, r)
			} else {
				log.Infof("Created Ref document with docID:[%s] for Ref:[%s]", refDoc.Id, r)
			}
		}
	}
	if len(missedRefs) > 0 {
		errStr := fmt.Sprintf("%d refs couldn't be imported. Errors were:\n", len(missedRefs))
		c := 0
		for _, e := range errs {
			errStr = errStr + e.Error() + "\n"
			c = c + 1
			if c >= maxErrorsReported {
				break
			}
		}
		return missedRefs, errors.New(errStr)
	}
	return nil, nil
}
<commit_msg>Remove or downgrade debug logging to log.Debug.<commit_after>package elasticsearch

import (
	"context"
	"errors"
	"fmt"

	"github.com/korfuri/goref"
	log "github.com/sirupsen/logrus"
	elastic "gopkg.in/olivere/elastic.v5"
)

const (
	// Max number of errors reported in one call to
	// LoadGraphToElastic
	maxErrorsReported = 20

	// Types in the Elastic search index
	packageType = "package"
	refType     = "ref"
)

// PackageExists returns whether the provided loadpath + version tuple
// exists in this index.
func PackageExists(loadpath string, version int64, client *elastic.Client, index string) bool {
	ctx := context.Background()
	docID := fmt.Sprintf("v1@%d@%s", version, loadpath)
	pkgDoc, _ := client.Get().
		Index(index).
		Type(packageType).
		Id(docID).
		Do(ctx)
	// TODO: handle errors better. Right now we assume that any
	// error is a 404 and can be ignored safely.
	return pkgDoc != nil
}

// LoadGraphToElastic loads all Packages and Refs from a PackageGraph
// to the provided ES index.
func LoadGraphToElastic(pg goref.PackageGraph, client *elastic.Client, index string) ([]*goref.Ref, error) {
	ctx := context.Background()
	missedRefs := make([]*goref.Ref, 0)
	errs := make([]error, 0)

	for _, p := range pg.Packages {
		if PackageExists(p.Path, p.Version, client, index) {
			log.Infof("Package %s already exists in this index.", p)
			continue
		}

		log.Debugf("Creating Package %s in the index", p)
		if _, err := client.Index().
			Index(index).
			Type(packageType).
			Id(p.DocumentID()).
			BodyJson(p).
			Do(ctx); err != nil {
			return nil, err
		}

		for _, r := range p.OutRefs {
			refDoc, err := client.Index().
				Index(index).
				Type(refType).
				BodyJson(r).
				Do(ctx)
			if err != nil {
				missedRefs = append(missedRefs, r)
				errs = append(errs, err)
				log.Debugf("Create Ref document failed with err:[%s] for Ref:[%s]", err, r)
			} else {
				log.Debugf("Created Ref document with docID:[%s] for Ref:[%s]", refDoc.Id, r)
			}
		}
	}
	if len(missedRefs) > 0 {
		errStr := fmt.Sprintf("%d refs couldn't be imported. Errors were:\n", len(missedRefs))
		c := 0
		for _, e := range errs {
			errStr = errStr + e.Error() + "\n"
			c = c + 1
			if c >= maxErrorsReported {
				break
			}
		}
		return missedRefs, errors.New(errStr)
	}
	return nil, nil
}
<|endoftext|>