<commit_before>// session.go - mixnet client session
// Copyright (C) 2018  David Stainton.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package client

import (
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	mrand "math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/katzenpost/client/config"
	cConstants "github.com/katzenpost/client/constants"
	"github.com/katzenpost/client/internal/pkiclient"
	"github.com/katzenpost/client/utils"
	coreConstants "github.com/katzenpost/core/constants"
	"github.com/katzenpost/core/crypto/ecdh"
	"github.com/katzenpost/core/log"
	"github.com/katzenpost/core/pki"
	"github.com/katzenpost/core/sphinx"
	sConstants "github.com/katzenpost/core/sphinx/constants"
	"github.com/katzenpost/core/worker"
	"github.com/katzenpost/minclient"
	"gopkg.in/eapache/channels.v1"
	"gopkg.in/op/go-logging.v1"
)

// Session is the struct type that keeps state for a given session.
type Session struct {
	worker.Worker

	cfg       *config.Config
	pkiClient pki.Client
	minclient *minclient.Client
	log       *logging.Logger

	fatalErrCh chan error
	opCh       chan workerOp

	eventCh   channels.Channel
	EventSink chan Event

	linkKey   *ecdh.PrivateKey
	onlineAt  time.Time
	hasPKIDoc bool

	egressQueue EgressQueue

	surbIDMap        sync.Map // [sConstants.SURBIDLength]byte -> *Message
	sentWaitChanMap  sync.Map // MessageID -> chan *Message
	replyWaitChanMap sync.Map // MessageID -> chan []byte

	decoyLoopTally uint64
}

// New establishes a session with provider using key.
// This method will block until session is connected to the Provider.
func NewSession(
	ctx context.Context,
	fatalErrCh chan error,
	logBackend *log.Backend,
	cfg *config.Config,
	linkKey *ecdh.PrivateKey) (*Session, error) {
	var err error

	// create a pkiclient for our own client lookups
	// AND create a pkiclient for minclient's use
	proxyCfg := cfg.UpstreamProxyConfig()
	pkiClient, err := cfg.NewPKIClient(logBackend, proxyCfg)
	if err != nil {
		return nil, err
	}

	// create a pkiclient for minclient's use
	pkiClient2, err := cfg.NewPKIClient(logBackend, proxyCfg)
	if err != nil {
		return nil, err
	}
	pkiCacheClient := pkiclient.New(pkiClient2)

	clientLog := logBackend.GetLogger(fmt.Sprintf("%s@%s_client", cfg.Account.User, cfg.Account.Provider))

	s := &Session{
		cfg:         cfg,
		linkKey:     linkKey,
		pkiClient:   pkiClient,
		log:         clientLog,
		fatalErrCh:  fatalErrCh,
		eventCh:     channels.NewInfiniteChannel(),
		EventSink:   make(chan Event),
		opCh:        make(chan workerOp, 8),
		egressQueue: new(Queue),
	}
	// Configure and bring up the minclient instance.
	clientCfg := &minclient.ClientConfig{
		User:                cfg.Account.User,
		Provider:            cfg.Account.Provider,
		ProviderKeyPin:      cfg.Account.ProviderKeyPin,
		LinkKey:             s.linkKey,
		LogBackend:          logBackend,
		PKIClient:           pkiCacheClient,
		OnConnFn:            s.onConnection,
		OnMessageFn:         s.onMessage,
		OnACKFn:             s.onACK,
		OnDocumentFn:        s.onDocument,
		DialContextFn:       proxyCfg.ToDialContext("authority"),
		PreferedTransports:  cfg.Debug.PreferedTransports,
		MessagePollInterval: time.Duration(cfg.Debug.PollingInterval) * time.Second,
		EnableTimeSync:      false, // Be explicit about it.
	}

	s.Go(s.eventSinkWorker)
	s.Go(s.garbageCollectionWorker)

	s.minclient, err = minclient.New(clientCfg)
	if err != nil {
		return nil, err
	}

	// block until we get the first PKI document
	// and then set our timers accordingly
	err = s.awaitFirstPKIDoc(ctx)
	if err != nil {
		return nil, err
	}
	s.Go(s.worker)
	return s, nil
}

func (s *Session) eventSinkWorker() {
	for {
		select {
		case <-s.HaltCh():
			s.log.Debugf("Event sink worker terminating gracefully.")
			return
		case e := <-s.eventCh.Out():
			select {
			case s.EventSink <- e.(Event):
			case <-s.HaltCh():
				s.log.Debugf("Event sink worker terminating gracefully.")
				return
			}
		}
	}
}

func (s *Session) garbageCollectionWorker() {
	timer := time.NewTimer(cConstants.GarbageCollectionInterval)
	defer timer.Stop()
	for {
		select {
		case <-s.HaltCh():
			s.log.Debugf("Garbage collection worker terminating gracefully.")
			return
		case <-timer.C:
			s.garbageCollect()
			timer.Reset(cConstants.GarbageCollectionInterval)
		}
	}
}

func (s *Session) garbageCollect() {
	s.log.Debug("Running garbage collection process.")
	// [sConstants.SURBIDLength]byte -> *Message
	surbIDMapRange := func(rawSurbID, rawMessage interface{}) bool {
		surbID := rawSurbID.([sConstants.SURBIDLength]byte)
		message := rawMessage.(*Message)
		if time.Now().After(message.SentAt.Add(message.ReplyETA).Add(cConstants.RoundTripTimeSlop)) {
			s.log.Debug("Garbage collecting SURB ID Map entry for Message ID %x", message.ID)
			s.surbIDMap.Delete(surbID)
			s.eventCh.In() <- &MessageIDGarbageCollected{
				MessageID: message.ID,
			}
		}
		return true
	}
	s.surbIDMap.Range(surbIDMapRange)
}

func (s *Session) awaitFirstPKIDoc(ctx context.Context) error {
	for {
		var qo workerOp
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-s.HaltCh():
			s.log.Debugf("Await first pki doc worker terminating gracefully")
			return errors.New("terminating gracefully")
		case <-time.After(time.Duration(s.cfg.Debug.InitialMaxPKIRetrievalDelay) * time.Second):
			return errors.New("timeout failure awaiting first PKI document")
		case qo = <-s.opCh:
		}
		switch op := qo.(type) {
		case opNewDocument:
			// Determine if PKI doc is valid. If not then abort.
			err := s.isDocValid(op.doc)
			if err != nil {
				s.fatalErrCh <- fmt.Errorf("aborting, PKI doc is not valid for our decoy traffic use case: %v", err)
				return err
			}
			return nil
		default:
			continue
		}
	}
	// NOT REACHED
}

// GetService returns a randomly selected service
// matching the specified service name
func (s *Session) GetService(serviceName string) (*utils.ServiceDescriptor, error) {
	doc := s.minclient.CurrentDocument()
	if doc == nil {
		return nil, errors.New("pki doc is nil")
	}
	serviceDescriptors := utils.FindServices(serviceName, doc)
	if len(serviceDescriptors) == 0 {
		return nil, errors.New("error, GetService failure, service not found in pki doc")
	}
	return &serviceDescriptors[mrand.Intn(len(serviceDescriptors))], nil
}

// OnConnection will be called by the minclient api
// upon connection change status to the Provider
func (s *Session) onConnection(err error) {
	s.log.Debugf("onConnection %v", err)
	s.eventCh.In() <- &ConnectionStatusEvent{
		IsConnected: err == nil,
		Err:         err,
	}
	s.opCh <- opConnStatusChanged{
		isConnected: err == nil,
	}
}

// OnMessage will be called by the minclient api
// upon receiving a message
func (s *Session) onMessage(ciphertextBlock []byte) error {
	s.log.Debugf("OnMessage")
	return nil
}

func (s *Session) incrementDecoyLoopTally() {
	atomic.AddUint64(&s.decoyLoopTally, 1)
}

func (s *Session) decrementDecoyLoopTally() {
	atomic.AddUint64(&s.decoyLoopTally, ^uint64(0))
}

// OnACK is called by the minclient api when we receive a SURB reply message.
func (s *Session) onACK(surbID *[sConstants.SURBIDLength]byte, ciphertext []byte) error {
	idStr := fmt.Sprintf("[%v]", hex.EncodeToString(surbID[:]))
	s.log.Infof("OnACK with SURBID %x", idStr)

	rawMessage, ok := s.surbIDMap.Load(*surbID)
	if !ok {
		s.log.Debug("Strange, received reply with unexpected SURBID")
		return nil
	}
	s.surbIDMap.Delete(*surbID)
	msg := rawMessage.(*Message)
	plaintext, err := sphinx.DecryptSURBPayload(ciphertext, msg.Key)
	if err != nil {
		s.log.Infof("Discarding SURB Reply, decryption failure: %s", err)
		return nil
	}
	if len(plaintext) != coreConstants.ForwardPayloadLength {
		s.log.Warningf("Discarding SURB %v: Invalid payload size: %v", idStr, len(plaintext))
		return nil
	}
	if msg.WithSURB && msg.IsDecoy {
		s.decrementDecoyLoopTally()
		return nil
	}

	if msg.IsBlocking {
		replyWaitChanRaw, ok := s.replyWaitChanMap.Load(*msg.ID)
		if !ok {
			//XXX: this can happen if a SURB-ACK arrives after a call to BlockingSendUnreliableMessage has timed-out
			// because the session.surbIDMap has not been deleted or garbage collected
			s.log.Warningf("Discarding surb %v for blocking message %x : caller likely timed-out", idStr, msg.ID)
			return nil
		}
		replyWaitChan := replyWaitChanRaw.(chan []byte)
		replyWaitChan <- plaintext[2:]
	} else {
		s.eventCh.In() <- &MessageReplyEvent{
			MessageID: msg.ID,
			Payload:   plaintext[2:],
			Err:       nil,
		}
	}
	return nil
}

func (s *Session) onDocument(doc *pki.Document) {
	s.log.Debugf("onDocument(): Epoch %v", doc.Epoch)
	s.hasPKIDoc = true
	s.opCh <- opNewDocument{
		doc: doc,
	}
	s.eventCh.In() <- &NewDocumentEvent{
		Document: doc,
	}
}

func (s *Session) CurrentDocument() *pki.Document {
	return s.minclient.CurrentDocument()
}

func (s *Session) GetReunionConfig() *config.Reunion {
	return s.cfg.Reunion
}

func (s *Session) GetPandaConfig() *config.Panda {
	return s.cfg.Panda
}

func (s *Session) Shutdown() {
	s.Halt()
	s.minclient.Shutdown()
	s.minclient.Wait()
}
<commit_msg>allow shorter polling intervals<commit_after>// session.go - mixnet client session
// Copyright (C) 2018  David Stainton.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package client

import (
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	mrand "math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/katzenpost/client/config"
	cConstants "github.com/katzenpost/client/constants"
	"github.com/katzenpost/client/internal/pkiclient"
	"github.com/katzenpost/client/utils"
	coreConstants "github.com/katzenpost/core/constants"
	"github.com/katzenpost/core/crypto/ecdh"
	"github.com/katzenpost/core/log"
	"github.com/katzenpost/core/pki"
	"github.com/katzenpost/core/sphinx"
	sConstants "github.com/katzenpost/core/sphinx/constants"
	"github.com/katzenpost/core/worker"
	"github.com/katzenpost/minclient"
	"gopkg.in/eapache/channels.v1"
	"gopkg.in/op/go-logging.v1"
)

// Session is the struct type that keeps state for a given session.
type Session struct {
	worker.Worker

	cfg       *config.Config
	pkiClient pki.Client
	minclient *minclient.Client
	log       *logging.Logger

	fatalErrCh chan error
	opCh       chan workerOp

	eventCh   channels.Channel
	EventSink chan Event

	linkKey   *ecdh.PrivateKey
	onlineAt  time.Time
	hasPKIDoc bool

	egressQueue EgressQueue

	surbIDMap        sync.Map // [sConstants.SURBIDLength]byte -> *Message
	sentWaitChanMap  sync.Map // MessageID -> chan *Message
	replyWaitChanMap sync.Map // MessageID -> chan []byte

	decoyLoopTally uint64
}

// New establishes a session with provider using key.
// This method will block until session is connected to the Provider.
func NewSession(
	ctx context.Context,
	fatalErrCh chan error,
	logBackend *log.Backend,
	cfg *config.Config,
	linkKey *ecdh.PrivateKey) (*Session, error) {
	var err error

	// create a pkiclient for our own client lookups
	// AND create a pkiclient for minclient's use
	proxyCfg := cfg.UpstreamProxyConfig()
	pkiClient, err := cfg.NewPKIClient(logBackend, proxyCfg)
	if err != nil {
		return nil, err
	}

	// create a pkiclient for minclient's use
	pkiClient2, err := cfg.NewPKIClient(logBackend, proxyCfg)
	if err != nil {
		return nil, err
	}
	pkiCacheClient := pkiclient.New(pkiClient2)

	clientLog := logBackend.GetLogger(fmt.Sprintf("%s@%s_client", cfg.Account.User, cfg.Account.Provider))

	s := &Session{
		cfg:         cfg,
		linkKey:     linkKey,
		pkiClient:   pkiClient,
		log:         clientLog,
		fatalErrCh:  fatalErrCh,
		eventCh:     channels.NewInfiniteChannel(),
		EventSink:   make(chan Event),
		opCh:        make(chan workerOp, 8),
		egressQueue: new(Queue),
	}
	// Configure and bring up the minclient instance.
	clientCfg := &minclient.ClientConfig{
		User:                cfg.Account.User,
		Provider:            cfg.Account.Provider,
		ProviderKeyPin:      cfg.Account.ProviderKeyPin,
		LinkKey:             s.linkKey,
		LogBackend:          logBackend,
		PKIClient:           pkiCacheClient,
		OnConnFn:            s.onConnection,
		OnMessageFn:         s.onMessage,
		OnACKFn:             s.onACK,
		OnDocumentFn:        s.onDocument,
		DialContextFn:       proxyCfg.ToDialContext("authority"),
		PreferedTransports:  cfg.Debug.PreferedTransports,
		MessagePollInterval: time.Duration(cfg.Debug.PollingInterval) * time.Millisecond,
		EnableTimeSync:      false, // Be explicit about it.
	}

	s.Go(s.eventSinkWorker)
	s.Go(s.garbageCollectionWorker)

	s.minclient, err = minclient.New(clientCfg)
	if err != nil {
		return nil, err
	}

	// block until we get the first PKI document
	// and then set our timers accordingly
	err = s.awaitFirstPKIDoc(ctx)
	if err != nil {
		return nil, err
	}
	s.Go(s.worker)
	return s, nil
}

func (s *Session) eventSinkWorker() {
	for {
		select {
		case <-s.HaltCh():
			s.log.Debugf("Event sink worker terminating gracefully.")
			return
		case e := <-s.eventCh.Out():
			select {
			case s.EventSink <- e.(Event):
			case <-s.HaltCh():
				s.log.Debugf("Event sink worker terminating gracefully.")
				return
			}
		}
	}
}

func (s *Session) garbageCollectionWorker() {
	timer := time.NewTimer(cConstants.GarbageCollectionInterval)
	defer timer.Stop()
	for {
		select {
		case <-s.HaltCh():
			s.log.Debugf("Garbage collection worker terminating gracefully.")
			return
		case <-timer.C:
			s.garbageCollect()
			timer.Reset(cConstants.GarbageCollectionInterval)
		}
	}
}

func (s *Session) garbageCollect() {
	s.log.Debug("Running garbage collection process.")
	// [sConstants.SURBIDLength]byte -> *Message
	surbIDMapRange := func(rawSurbID, rawMessage interface{}) bool {
		surbID := rawSurbID.([sConstants.SURBIDLength]byte)
		message := rawMessage.(*Message)
		if time.Now().After(message.SentAt.Add(message.ReplyETA).Add(cConstants.RoundTripTimeSlop)) {
			s.log.Debug("Garbage collecting SURB ID Map entry for Message ID %x", message.ID)
			s.surbIDMap.Delete(surbID)
			s.eventCh.In() <- &MessageIDGarbageCollected{
				MessageID: message.ID,
			}
		}
		return true
	}
	s.surbIDMap.Range(surbIDMapRange)
}

func (s *Session) awaitFirstPKIDoc(ctx context.Context) error {
	for {
		var qo workerOp
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-s.HaltCh():
			s.log.Debugf("Await first pki doc worker terminating gracefully")
			return errors.New("terminating gracefully")
		case <-time.After(time.Duration(s.cfg.Debug.InitialMaxPKIRetrievalDelay) * time.Second):
			return errors.New("timeout failure awaiting first PKI document")
		case qo = <-s.opCh:
		}
		switch op := qo.(type) {
		case opNewDocument:
			// Determine if PKI doc is valid. If not then abort.
			err := s.isDocValid(op.doc)
			if err != nil {
				s.fatalErrCh <- fmt.Errorf("aborting, PKI doc is not valid for our decoy traffic use case: %v", err)
				return err
			}
			return nil
		default:
			continue
		}
	}
	// NOT REACHED
}

// GetService returns a randomly selected service
// matching the specified service name
func (s *Session) GetService(serviceName string) (*utils.ServiceDescriptor, error) {
	doc := s.minclient.CurrentDocument()
	if doc == nil {
		return nil, errors.New("pki doc is nil")
	}
	serviceDescriptors := utils.FindServices(serviceName, doc)
	if len(serviceDescriptors) == 0 {
		return nil, errors.New("error, GetService failure, service not found in pki doc")
	}
	return &serviceDescriptors[mrand.Intn(len(serviceDescriptors))], nil
}

// OnConnection will be called by the minclient api
// upon connection change status to the Provider
func (s *Session) onConnection(err error) {
	s.log.Debugf("onConnection %v", err)
	s.eventCh.In() <- &ConnectionStatusEvent{
		IsConnected: err == nil,
		Err:         err,
	}
	s.opCh <- opConnStatusChanged{
		isConnected: err == nil,
	}
}

// OnMessage will be called by the minclient api
// upon receiving a message
func (s *Session) onMessage(ciphertextBlock []byte) error {
	s.log.Debugf("OnMessage")
	return nil
}

func (s *Session) incrementDecoyLoopTally() {
	atomic.AddUint64(&s.decoyLoopTally, 1)
}

func (s *Session) decrementDecoyLoopTally() {
	atomic.AddUint64(&s.decoyLoopTally, ^uint64(0))
}

// OnACK is called by the minclient api when we receive a SURB reply message.
func (s *Session) onACK(surbID *[sConstants.SURBIDLength]byte, ciphertext []byte) error {
	idStr := fmt.Sprintf("[%v]", hex.EncodeToString(surbID[:]))
	s.log.Infof("OnACK with SURBID %x", idStr)

	rawMessage, ok := s.surbIDMap.Load(*surbID)
	if !ok {
		s.log.Debug("Strange, received reply with unexpected SURBID")
		return nil
	}
	s.surbIDMap.Delete(*surbID)
	msg := rawMessage.(*Message)
	plaintext, err := sphinx.DecryptSURBPayload(ciphertext, msg.Key)
	if err != nil {
		s.log.Infof("Discarding SURB Reply, decryption failure: %s", err)
		return nil
	}
	if len(plaintext) != coreConstants.ForwardPayloadLength {
		s.log.Warningf("Discarding SURB %v: Invalid payload size: %v", idStr, len(plaintext))
		return nil
	}
	if msg.WithSURB && msg.IsDecoy {
		s.decrementDecoyLoopTally()
		return nil
	}

	if msg.IsBlocking {
		replyWaitChanRaw, ok := s.replyWaitChanMap.Load(*msg.ID)
		if !ok {
			//XXX: this can happen if a SURB-ACK arrives after a call to BlockingSendUnreliableMessage has timed-out
			// because the session.surbIDMap has not been deleted or garbage collected
			s.log.Warningf("Discarding surb %v for blocking message %x : caller likely timed-out", idStr, msg.ID)
			return nil
		}
		replyWaitChan := replyWaitChanRaw.(chan []byte)
		replyWaitChan <- plaintext[2:]
	} else {
		s.eventCh.In() <- &MessageReplyEvent{
			MessageID: msg.ID,
			Payload:   plaintext[2:],
			Err:       nil,
		}
	}
	return nil
}

func (s *Session) onDocument(doc *pki.Document) {
	s.log.Debugf("onDocument(): Epoch %v", doc.Epoch)
	s.hasPKIDoc = true
	s.opCh <- opNewDocument{
		doc: doc,
	}
	s.eventCh.In() <- &NewDocumentEvent{
		Document: doc,
	}
}

func (s *Session) CurrentDocument() *pki.Document {
	return s.minclient.CurrentDocument()
}

func (s *Session) GetReunionConfig() *config.Reunion {
	return s.cfg.Reunion
}

func (s *Session) GetPandaConfig() *config.Panda {
	return s.cfg.Panda
}

func (s *Session) Shutdown() {
	s.Halt()
	s.minclient.Shutdown()
	s.minclient.Wait()
}
<|endoftext|>


<commit_before>package revel

import (
	"fmt"
	"github.com/streadway/simpleuuid"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

// A signed cookie (and thus limited to 4kb in size).
// Restriction: Keys may not have a colon in them.
type Session map[string]string

const (
	SESSION_ID_KEY = "_ID"
	TS_KEY         = "_TS"
)

var expireAfterDuration time.Duration

func init() {
	// Set expireAfterDuration, default to 30 days if no value in config
	OnAppStart(func() {
		var err error
		if expiresString, ok := Config.String("session.expires"); !ok {
			expireAfterDuration = 30 * 24 * time.Hour
		} else if expireAfterDuration, err = time.ParseDuration(expiresString); err != nil {
			panic(fmt.Errorf("session.expires invalid: %s", err))
		}
	})
}

// Return a UUID identifying this session.
func (s Session) Id() string {
	if uuidStr, ok := s[SESSION_ID_KEY]; ok {
		return uuidStr
	}

	uuid, err := simpleuuid.NewTime(time.Now())
	if err != nil {
		panic(err) // I don't think this can actually happen.
	}
	s[SESSION_ID_KEY] = uuid.String()
	return s[SESSION_ID_KEY]
}

// Return a time.Time with session expiration date
func getSessionExpiration() time.Time {
	return time.Now().Add(expireAfterDuration)
}

// Returns an http.Cookie containing the signed session.
func (s Session) cookie() *http.Cookie {
	var sessionValue string
	ts := getSessionExpiration()
	s[TS_KEY] = getSessionExpirationCookie(ts)
	for key, value := range s {
		if strings.ContainsAny(key, ":\x00") {
			panic("Session keys may not have colons or null bytes")
		}
		if strings.Contains(value, "\x00") {
			panic("Session values may not have null bytes")
		}
		sessionValue += "\x00" + key + ":" + value + "\x00"
	}

	sessionData := url.QueryEscape(sessionValue)
	return &http.Cookie{
		Name:     CookiePrefix + "_SESSION",
		Value:    Sign(sessionData) + "-" + sessionData,
		Path:     "/",
		HttpOnly: CookieHttpOnly,
		Secure:   CookieSecure,
		Expires:  ts.UTC(),
	}
}

func sessionTimeoutExpiredOrMissing(session Session) bool {
	if exp, present := session[TS_KEY]; !present {
		return true
	} else if expInt, _ := strconv.Atoi(exp); int64(expInt) < time.Now().Unix() {
		return true
	}
	return false
}

// Returns a Session pulled from signed cookie.
func getSessionFromCookie(cookie *http.Cookie) Session {
	session := make(Session)

	// Separate the data from the signature.
	hyphen := strings.Index(cookie.Value, "-")
	if hyphen == -1 || hyphen >= len(cookie.Value)-1 {
		return session
	}
	sig, data := cookie.Value[:hyphen], cookie.Value[hyphen+1:]

	// Verify the signature.
	if !Verify(data, sig) {
		INFO.Println("Session cookie signature failed")
		return session
	}

	ParseKeyValueCookie(data, func(key, val string) {
		session[key] = val
	})

	if sessionTimeoutExpiredOrMissing(session) {
		session = make(Session)
	}

	return session
}

func SessionFilter(c *Controller, fc []Filter) {
	c.Session = restoreSession(c.Request.Request)

	fc[0](c, fc[1:])

	// Store the session (and sign it).
	c.SetCookie(c.Session.cookie())
}

func restoreSession(req *http.Request) Session {
	session := make(Session)
	cookie, err := req.Cookie(CookiePrefix + "_SESSION")
	if err != nil {
		return session
	}

	return getSessionFromCookie(cookie)
}

func getSessionExpirationCookie(t time.Time) string {
	return strconv.FormatInt(t.Unix(), 10)
}
<commit_msg>Session variables are now copied to RenderArgs["session"]<commit_after>package revel

import (
	"fmt"
	"github.com/streadway/simpleuuid"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

// A signed cookie (and thus limited to 4kb in size).
// Restriction: Keys may not have a colon in them.
type Session map[string]string

const (
	SESSION_ID_KEY = "_ID"
	TS_KEY         = "_TS"
)

var expireAfterDuration time.Duration

func init() {
	// Set expireAfterDuration, default to 30 days if no value in config
	OnAppStart(func() {
		var err error
		if expiresString, ok := Config.String("session.expires"); !ok {
			expireAfterDuration = 30 * 24 * time.Hour
		} else if expireAfterDuration, err = time.ParseDuration(expiresString); err != nil {
			panic(fmt.Errorf("session.expires invalid: %s", err))
		}
	})
}

// Return a UUID identifying this session.
func (s Session) Id() string {
	if uuidStr, ok := s[SESSION_ID_KEY]; ok {
		return uuidStr
	}

	uuid, err := simpleuuid.NewTime(time.Now())
	if err != nil {
		panic(err) // I don't think this can actually happen.
	}
	s[SESSION_ID_KEY] = uuid.String()
	return s[SESSION_ID_KEY]
}

// Return a time.Time with session expiration date
func getSessionExpiration() time.Time {
	return time.Now().Add(expireAfterDuration)
}

// Returns an http.Cookie containing the signed session.
func (s Session) cookie() *http.Cookie {
	var sessionValue string
	ts := getSessionExpiration()
	s[TS_KEY] = getSessionExpirationCookie(ts)
	for key, value := range s {
		if strings.ContainsAny(key, ":\x00") {
			panic("Session keys may not have colons or null bytes")
		}
		if strings.Contains(value, "\x00") {
			panic("Session values may not have null bytes")
		}
		sessionValue += "\x00" + key + ":" + value + "\x00"
	}

	sessionData := url.QueryEscape(sessionValue)
	return &http.Cookie{
		Name:     CookiePrefix + "_SESSION",
		Value:    Sign(sessionData) + "-" + sessionData,
		Path:     "/",
		HttpOnly: CookieHttpOnly,
		Secure:   CookieSecure,
		Expires:  ts.UTC(),
	}
}

func sessionTimeoutExpiredOrMissing(session Session) bool {
	if exp, present := session[TS_KEY]; !present {
		return true
	} else if expInt, _ := strconv.Atoi(exp); int64(expInt) < time.Now().Unix() {
		return true
	}
	return false
}

// Returns a Session pulled from signed cookie.
func getSessionFromCookie(cookie *http.Cookie) Session {
	session := make(Session)

	// Separate the data from the signature.
	hyphen := strings.Index(cookie.Value, "-")
	if hyphen == -1 || hyphen >= len(cookie.Value)-1 {
		return session
	}
	sig, data := cookie.Value[:hyphen], cookie.Value[hyphen+1:]

	// Verify the signature.
	if !Verify(data, sig) {
		INFO.Println("Session cookie signature failed")
		return session
	}

	ParseKeyValueCookie(data, func(key, val string) {
		session[key] = val
	})

	if sessionTimeoutExpiredOrMissing(session) {
		session = make(Session)
	}

	return session
}

func SessionFilter(c *Controller, fc []Filter) {
	c.Session = restoreSession(c.Request.Request)
	// Make session vars available in templates as {{.session.xyz}}
	c.RenderArgs["session"] = c.Session

	fc[0](c, fc[1:])

	// Store the session (and sign it).
	c.SetCookie(c.Session.cookie())
}

func restoreSession(req *http.Request) Session {
	session := make(Session)
	cookie, err := req.Cookie(CookiePrefix + "_SESSION")
	if err != nil {
		return session
	}

	return getSessionFromCookie(cookie)
}

func getSessionExpirationCookie(t time.Time) string {
	return strconv.FormatInt(t.Unix(), 10)
}
<|endoftext|>


<commit_before>package session

import (
	"log"
	"net/http"

	"github.com/gorilla/sessions"
	"github.com/labstack/echo"
	"github.com/labstack/echo/engine/standard"
)

const (
	DefaultKey  = "github.com/ipfans/echo-session"
	errorFormat = "[sessions] ERROR! %s\n"
)

type Store interface {
	sessions.Store
	Options(Options)
	MaxAge(int)
}

// Options stores configuration for a session or session store.
// Fields are a subset of http.Cookie fields.
type Options struct {
	Path   string
	Domain string
	// MaxAge=0 means no 'Max-Age' attribute specified.
	// MaxAge<0 means delete cookie now, equivalently 'Max-Age: 0'.
	// MaxAge>0 means Max-Age attribute present and given in seconds.
	MaxAge   int
	Secure   bool
	HttpOnly bool
}

// Wraps thinly gorilla-session methods.
// Session stores the values and optional configuration for a session.
type Session interface {
	// Get returns the session value associated to the given key.
	Get(key interface{}) interface{}
	// Set sets the session value associated to the given key.
	Set(key interface{}, val interface{})
	// Delete removes the session value associated to the given key.
	Delete(key interface{})
	// Clear deletes all values in the session.
	Clear()
	// AddFlash adds a flash message to the session.
	// A single variadic argument is accepted, and it is optional: it defines the flash key.
	// If not defined "_flash" is used by default.
	AddFlash(value interface{}, vars ...string)
	// Flashes returns a slice of flash messages from the session.
	// A single variadic argument is accepted, and it is optional: it defines the flash key.
	// If not defined "_flash" is used by default.
	Flashes(vars ...string) []interface{}
	// Options sets confuguration for a session.
	Options(Options)
	// Save saves all sessions used during the current request.
	Save() error
}

func Sessions(name string, store Store) echo.MiddlewareFunc {
	return func(h echo.HandlerFunc) echo.HandlerFunc {
		return func(ctx echo.Context) error {
			rq := ctx.Request().(*standard.Request)
			rs := ctx.Response().(*standard.Response)

			s := &session{name, rq.Request, store, nil, false, rs.ResponseWriter}
			ctx.Set(DefaultKey, s)
			return h(ctx)
		}
	}
}

type session struct {
	name    string
	request *http.Request
	store   Store
	session *sessions.Session
	written bool
	writer  http.ResponseWriter
}

func (s *session) Get(key interface{}) interface{} {
	return s.Session().Values[key]
}

func (s *session) Set(key interface{}, val interface{}) {
	s.Session().Values[key] = val
	s.written = true
}

func (s *session) Delete(key interface{}) {
	delete(s.Session().Values, key)
	s.written = true
}

func (s *session) Clear() {
	for key := range s.Session().Values {
		s.Delete(key)
	}
}

func (s *session) AddFlash(value interface{}, vars ...string) {
	s.Session().AddFlash(value, vars...)
	s.written = true
}

func (s *session) Flashes(vars ...string) []interface{} {
	s.written = true
	return s.Session().Flashes(vars...)
}

func (s *session) Options(options Options) {
	s.Session().Options = &sessions.Options{
		Path:     options.Path,
		Domain:   options.Domain,
		MaxAge:   options.MaxAge,
		Secure:   options.Secure,
		HttpOnly: options.HttpOnly,
	}
}

func (s *session) Save() error {
	if s.Written() {
		e := s.Session().Save(s.request, s.writer)
		if e == nil {
			s.written = false
		}
		return e
	}
	return nil
}

func (s *session) Session() *sessions.Session {
	if s.session == nil {
		var err error
		s.session, err = s.store.Get(s.request, s.name)
		if err != nil {
			log.Printf(errorFormat, err)
		}
	}
	return s.session
}

func (s *session) Written() bool {
	return s.written
}

// shortcut to get session
func Default(ctx echo.Context) Session {
	session := ctx.Get(DefaultKey)
	if session == nil {
		return nil
	}
	return ctx.Get(DefaultKey).(Session)
}
<commit_msg>Compatibility for echo v3<commit_after>package session

import (
	"log"
	"net/http"

	"github.com/gorilla/sessions"
	"github.com/labstack/echo"
)

const (
	DefaultKey  = "github.com/ipfans/echo-session"
	errorFormat = "[sessions] ERROR! %s\n"
)

type Store interface {
	sessions.Store
	Options(Options)
	MaxAge(int)
}

// Options stores configuration for a session or session store.
// Fields are a subset of http.Cookie fields.
type Options struct {
	Path   string
	Domain string
	// MaxAge=0 means no 'Max-Age' attribute specified.
	// MaxAge<0 means delete cookie now, equivalently 'Max-Age: 0'.
	// MaxAge>0 means Max-Age attribute present and given in seconds.
	MaxAge   int
	Secure   bool
	HttpOnly bool
}

// Wraps thinly gorilla-session methods.

// Session stores the values and optional configuration for a session.
type Session interface {
	// Get returns the session value associated to the given key.
	Get(key interface{}) interface{}
	// Set sets the session value associated to the given key.
	Set(key interface{}, val interface{})
	// Delete removes the session value associated to the given key.
	Delete(key interface{})
	// Clear deletes all values in the session.
	Clear()
	// AddFlash adds a flash message to the session.
	// A single variadic argument is accepted, and it is optional: it defines the flash key.
	// If not defined "_flash" is used by default.
	AddFlash(value interface{}, vars ...string)
	// Flashes returns a slice of flash messages from the session.
	// A single variadic argument is accepted, and it is optional: it defines the flash key.
	// If not defined "_flash" is used by default.
	Flashes(vars ...string) []interface{}
	// Options sets confuguration for a session.
	Options(Options)
	// Save saves all sessions used during the current request.
	Save() error
}

func Sessions(name string, store Store) echo.MiddlewareFunc {
	return func(h echo.HandlerFunc) echo.HandlerFunc {
		return func(ctx echo.Context) error {
			rq := ctx.Request()
			rs := ctx.Response()
			s := &session{name, rq, store, nil, false, rs.Writer()}
			ctx.Set(DefaultKey, s)
			return h(ctx)
		}
	}
}

type session struct {
	name    string
	request *http.Request
	store   Store
	session *sessions.Session
	written bool
	writer  http.ResponseWriter
}

func (s *session) Get(key interface{}) interface{} {
	return s.Session().Values[key]
}

func (s *session) Set(key interface{}, val interface{}) {
	s.Session().Values[key] = val
	s.written = true
}

func (s *session) Delete(key interface{}) {
	delete(s.Session().Values, key)
	s.written = true
}

func (s *session) Clear() {
	for key := range s.Session().Values {
		s.Delete(key)
	}
}

func (s *session) AddFlash(value interface{}, vars ...string) {
	s.Session().AddFlash(value, vars...)
	s.written = true
}

func (s *session) Flashes(vars ...string) []interface{} {
	s.written = true
	return s.Session().Flashes(vars...)
}

func (s *session) Options(options Options) {
	s.Session().Options = &sessions.Options{
		Path:     options.Path,
		Domain:   options.Domain,
		MaxAge:   options.MaxAge,
		Secure:   options.Secure,
		HttpOnly: options.HttpOnly,
	}
}

func (s *session) Save() error {
	if s.Written() {
		e := s.Session().Save(s.request, s.writer)
		if e == nil {
			s.written = false
		}
		return e
	}
	return nil
}

func (s *session) Session() *sessions.Session {
	if s.session == nil {
		var err error
		s.session, err = s.store.Get(s.request, s.name)
		if err != nil {
			log.Printf(errorFormat, err)
		}
	}
	return s.session
}

func (s *session) Written() bool {
	return s.written
}

// shortcut to get session
func Default(ctx echo.Context) Session {
	session := ctx.Get(DefaultKey)
	if session == nil {
		return nil
	}
	return ctx.Get(DefaultKey).(Session)
}
<|endoftext|>


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

// The following enables go generate to generate the doc.go file.
//go:generate go run $JIRI_ROOT/release/go/src/v.io/x/lib/cmdline/testdata/gendoc.go . -help

package main

import (
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	goauth2 "google.golang.org/api/oauth2/v2"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/security"
	"v.io/v23/vom"
	"v.io/x/lib/cmdline"
	"v.io/x/ref"
	lsecurity "v.io/x/ref/lib/security"
	"v.io/x/ref/lib/v23cmd"
	"v.io/x/ref/services/agent/internal/ipc"
	"v.io/x/ref/services/agent/internal/server"

	_ "v.io/x/ref/runtime/factories/roaming"
)

var (
	keyFileFlag      string
	oauthBlesserFlag string
)

func main() {
	cmdGcreds.Flags.StringVar(&keyFileFlag, "key-file", "", "The JSON file containing the Google Cloud credentials to use. If empty, the default credentials will be used.")
	cmdGcreds.Flags.StringVar(&oauthBlesserFlag, "oauth-blesser", "https://dev.v.io/auth/google/bless", "The URL of the OAuthBlesser service.")

	cmdline.HideGlobalFlagsExcept()
	cmdline.Main(cmdGcreds)
}

var cmdGcreds = &cmdline.Command{
	Runner: v23cmd.RunnerFunc(runGcreds),
	Name:   "gcreds",
	Short:  "Runs a command with Google Cloud blessings",
	Long: `
Command gcreds runs a command with Google Cloud Blessings.

The Google Cloud credentials can be specified with the --key-file flag. If
--key-file is not set, the default credentials will be used. These credentials
are exchanged for Vanadium blessings using the OAuthBlesser service specified
with the --oauth-blesser flag.

The command is executed with an ephemeral principal that has the blessings
received from the OAuthBlesser service. The principal and its blessings are held
in memory and do not persist after gcreds exits.

See https://developers.google.com/identity/protocols/application-default-credentials
for more information on Google Cloud credentials.
`,
	ArgsName: "<command>",
}

func runGcreds(ctx *context.T, env *cmdline.Env, args []string) error {
	if len(args) == 0 {
		return env.UsageErrorf("missing command")
	}
	p, err := lsecurity.NewPrincipal()
	if err != nil {
		return err
	}
	if ctx, err = v23.WithPrincipal(ctx, p); err != nil {
		return err
	}
	if err = getGoogleCloudBlessings(ctx); err != nil {
		return err
	}
	// Refresh the blessings periodically.
	go refreshCreds(ctx)

	workDir, err := ioutil.TempDir("", "gcreds-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(workDir)
	socketPath := filepath.Join(workDir, "agent.sock")
	keyPath := filepath.Join(workDir, "keys")

	// Run the server.
	i := ipc.NewIPC()
	defer i.Close()
	if err = server.ServeAgent(i, p); err != nil {
		return err
	}
	if err = server.ServeKeyManager(i, keyPath, nil); err != nil {
		return err
	}

	if err = i.Listen(socketPath); err != nil {
		return err
	}

	ref.EnvClearCredentials()
	if err := os.Setenv(ref.EnvAgentPath, socketPath); err != nil {
		return err
	}
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdin = env.Stdin
	cmd.Stdout = env.Stdout
	cmd.Stderr = env.Stderr
	return cmd.Run()
}

func refreshCreds(ctx *context.T) {
	for {
		p := v23.GetPrincipal(ctx)
		exp := p.BlessingStore().Default().Expiry()
		if exp.IsZero() {
			// Nothing to do.
			return
		}
		delay := exp.Add(-time.Minute).Sub(time.Now())
		if delay < time.Second {
			delay = time.Second
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(delay):
			ctx.Info("refreshing credentials")
			if err := getGoogleCloudBlessings(ctx); err != nil {
				ctx.Errorf("getGoogleCloudBlessings: %v", err)
			}
		}
	}
}

// getGoogleCloudBlessings gets blessings using Google Cloud OAuth credentials.
func getGoogleCloudBlessings(ctx *context.T) error {
	const scope = goauth2.UserinfoEmailScope
	var tokSource oauth2.TokenSource
	if len(keyFileFlag) > 0 {
		data, err := ioutil.ReadFile(keyFileFlag)
		if err != nil {
			return err
		}
		conf, err := google.JWTConfigFromJSON(data, scope)
		if err != nil {
			return err
		}
		tokSource = conf.TokenSource(oauth2.NoContext)
	} else {
		var err error
		if tokSource, err = google.DefaultTokenSource(oauth2.NoContext, scope); err != nil {
			return err
		}
	}
	token, err := tokSource.Token()
	if err != nil {
		return err
	}

	principal := v23.GetPrincipal(ctx)
	bytes, err := principal.PublicKey().MarshalBinary()
	if err != nil {
		return err
	}
	expiry, err := security.NewExpiryCaveat(time.Now().Add(10 * time.Minute))
	if err != nil {
		return err
	}
	caveats, err := base64VomEncode([]security.Caveat{expiry})
	if err != nil {
		return err
	}
	resp, err := http.PostForm(
		oauthBlesserFlag,
		// This interface is defined in:
		// https://godoc.org/v.io/x/ref/services/identity/internal/handlers#NewOAuthBlessingHandler
		url.Values{
			"public_key":    {base64.URLEncoding.EncodeToString(bytes)},
			"token":         {token.AccessToken},
			"caveats":       {caveats},
			"output_format": {"base64vom"},
		},
	)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("got %s", resp.Status)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var blessings security.Blessings
	if err := base64VomDecode(string(body), &blessings); err != nil {
		return err
	}
	if err := lsecurity.SetDefaultBlessings(principal, blessings); err != nil {
		return err
	}
	return nil
}

func base64VomEncode(i interface{}) (string, error) {
	v, err := vom.Encode(i)
	return base64.URLEncoding.EncodeToString(v), err
}

func base64VomDecode(s string, i interface{}) error {
	b, err := base64.URLEncoding.DecodeString(s)
	if err != nil {
		return err
	}
	return vom.Decode(b, i)
}
<commit_msg>TBR: services/agent/gcreds: Retry requests to oauth-blesser<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The following enables go generate to generate the doc.go file.
//go:generate go run $JIRI_ROOT/release/go/src/v.io/x/lib/cmdline/testdata/gendoc.go . -help

package main

import (
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	goauth2 "google.golang.org/api/oauth2/v2"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/security"
	"v.io/v23/vom"
	"v.io/x/lib/cmdline"
	"v.io/x/ref"
	lsecurity "v.io/x/ref/lib/security"
	"v.io/x/ref/lib/v23cmd"
	"v.io/x/ref/services/agent/internal/ipc"
	"v.io/x/ref/services/agent/internal/server"

	_ "v.io/x/ref/runtime/factories/roaming"
)

var (
	keyFileFlag      string
	oauthBlesserFlag string
)

func main() {
	cmdGcreds.Flags.StringVar(&keyFileFlag, "key-file", "", "The JSON file containing the Google Cloud credentials to use. If empty, the default credentials will be used.")
	cmdGcreds.Flags.StringVar(&oauthBlesserFlag, "oauth-blesser", "https://dev.v.io/auth/google/bless", "The URL of the OAuthBlesser service.")

	cmdline.HideGlobalFlagsExcept()
	cmdline.Main(cmdGcreds)
}

var cmdGcreds = &cmdline.Command{
	Runner: v23cmd.RunnerFunc(runGcreds),
	Name:   "gcreds",
	Short:  "Runs a command with Google Cloud blessings",
	Long: `
Command gcreds runs a command with Google Cloud Blessings.

The Google Cloud credentials can be specified with the --key-file flag. If
--key-file is not set, the default credentials will be used. These credentials
are exchanged for Vanadium blessings using the OAuthBlesser service specified
with the --oauth-blesser flag.

The command is executed with an ephemeral principal that has the blessings
received from the OAuthBlesser service. The principal and its blessings are held
in memory and do not persist after gcreds exits.

See https://developers.google.com/identity/protocols/application-default-credentials
for more information on Google Cloud credentials.
`,
	ArgsName: "<command>",
}

func runGcreds(ctx *context.T, env *cmdline.Env, args []string) error {
	if len(args) == 0 {
		return env.UsageErrorf("missing command")
	}
	p, err := lsecurity.NewPrincipal()
	if err != nil {
		return err
	}
	if ctx, err = v23.WithPrincipal(ctx, p); err != nil {
		return err
	}
	if err = getGoogleCloudBlessings(ctx); err != nil {
		return err
	}
	// Refresh the blessings periodically.
	go refreshCreds(ctx)

	workDir, err := ioutil.TempDir("", "gcreds-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(workDir)
	socketPath := filepath.Join(workDir, "agent.sock")
	keyPath := filepath.Join(workDir, "keys")

	// Run the server.
	i := ipc.NewIPC()
	defer i.Close()
	if err = server.ServeAgent(i, p); err != nil {
		return err
	}
	if err = server.ServeKeyManager(i, keyPath, nil); err != nil {
		return err
	}

	if err = i.Listen(socketPath); err != nil {
		return err
	}

	ref.EnvClearCredentials()
	if err := os.Setenv(ref.EnvAgentPath, socketPath); err != nil {
		return err
	}
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdin = env.Stdin
	cmd.Stdout = env.Stdout
	cmd.Stderr = env.Stderr
	return cmd.Run()
}

func refreshCreds(ctx *context.T) {
	for {
		p := v23.GetPrincipal(ctx)
		exp := p.BlessingStore().Default().Expiry()
		if exp.IsZero() {
			// Nothing to do.
			return
		}
		delay := exp.Add(-time.Minute).Sub(time.Now())
		if delay < time.Second {
			delay = time.Second
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(delay):
			ctx.Info("refreshing credentials")
			if err := getGoogleCloudBlessings(ctx); err != nil {
				ctx.Errorf("getGoogleCloudBlessings: %v", err)
			}
		}
	}
}

// getGoogleCloudBlessings gets blessings using Google Cloud OAuth credentials.
func getGoogleCloudBlessings(ctx *context.T) error {
	const scope = goauth2.UserinfoEmailScope
	var tokSource oauth2.TokenSource
	if len(keyFileFlag) > 0 {
		data, err := ioutil.ReadFile(keyFileFlag)
		if err != nil {
			return err
		}
		conf, err := google.JWTConfigFromJSON(data, scope)
		if err != nil {
			return err
		}
		tokSource = conf.TokenSource(oauth2.NoContext)
	} else {
		var err error
		if tokSource, err = google.DefaultTokenSource(oauth2.NoContext, scope); err != nil {
			return err
		}
	}
	token, err := tokSource.Token()
	if err != nil {
		return err
	}

	principal := v23.GetPrincipal(ctx)
	bytes, err := principal.PublicKey().MarshalBinary()
	if err != nil {
		return err
	}
	expiry, err := security.NewExpiryCaveat(time.Now().Add(10 * time.Minute))
	if err != nil {
		return err
	}
	caveats, err := base64VomEncode([]security.Caveat{expiry})
	if err != nil {
		return err
	}
	// This interface is defined in:
	// https://godoc.org/v.io/x/ref/services/identity/internal/handlers#NewOAuthBlessingHandler
	v := url.Values{
		"public_key":    {base64.URLEncoding.EncodeToString(bytes)},
		"token":         {token.AccessToken},
		"caveats":       {caveats},
		"output_format": {"base64vom"},
	}
	for attempt := 0; attempt < 30; attempt++ {
		if attempt > 0 {
			ctx.Infof("retrying")
			time.Sleep(time.Second)
		}
		if body, err := postBlessRequest(v); err == nil {
			var blessings security.Blessings
			if err := base64VomDecode(string(body), &blessings); err != nil {
				return err
			}
			return lsecurity.SetDefaultBlessings(principal, blessings)
		} else {
			ctx.Infof("error from oauth-blesser: %v", err)
		}
	}
	return fmt.Errorf("too many failures")
}

func postBlessRequest(values url.Values) ([]byte, error) {
	resp, err := http.PostForm(oauthBlesserFlag, values)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("got %s", resp.Status)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return body, nil
}

func base64VomEncode(i interface{}) (string, error) {
	v, err := vom.Encode(i)
	return base64.URLEncoding.EncodeToString(v), err
}

func base64VomDecode(s string, i interface{}) error {
	b, err := base64.URLEncoding.DecodeString(s)
	if err != nil {
		return err
	}
	return vom.Decode(b, i)
}
<|endoftext|>


<commit_before>package services

import (
	"encoding/json"
	"log"
	"net/http"
	"time"

	"github.com/khades/servbot/models"
	"github.com/khades/servbot/repos"
)

type stream struct {
	CreatedAt time.Time `json:"created_at"`
	Channel   channelInfo
}

type channelInfo struct {
	Status string
	Game   string
}

type responseBody struct {
	Stream stream
}

// CheckStreamStatus runs gettting all data from all channels bot applied to
func CheckStreamStatus() {
	for _, value := range repos.Config.Channels {
		go getStatus(value)
	}
}

func getStatus(channel string) {
	var responseBody = new(responseBody)
	url := "https://api.twitch.tv/kraken/streams/" + channel + "?client_id=c5w8oif66lg711pfrh86h8od3sek43d"
	resp, respError := http.Get(url)
	var status = models.StreamStatus{
		Online: false}
	if respError != nil {
		return
	}
	defer resp.Body.Close()

	marshallError := json.NewDecoder(resp.Body).Decode(responseBody)
	log.Println(responseBody)
	if marshallError != nil {
		return
	}
	if !responseBody.Stream.CreatedAt.IsZero() {
		status = models.StreamStatus{
			Online:      true,
			Description: "я кот",
			Game:        responseBody.Stream.Channel.Game,
			Title:       responseBody.Stream.Channel.Status,
			Start:       responseBody.Stream.CreatedAt}
	}
	// else {
	// 	repos.SetStreamStatusOffline(channel)
	// }
	repos.PushStreamStatus(channel, status)

}

// func getOfflineStatus(channel string) {
// 	var responseBody = new(channelInfo)
// 	url := "https://api.twitch.tv/kraken/channels/" + channel
// 	resp, respError := http.Get(url)

// 	if respError != nil {
// 		return
// 	}
// 	defer resp.Body.Close()

// 	marshallError := json.NewDecoder(resp.Body).Decode(responseBody)

// 	if marshallError != nil {
// 		return
// 	}

// 	status := models.StreamStatus{
// 		Online: false,
// 		Game:   responseBody.Game,
// 		Title:  responseBody.Status}
// 	repos.PushStreamStatus(channel, status)

// }
<commit_msg>asfsa<commit_after>package services

import (
	"encoding/json"
	"log"
	"net/http"
	"time"

	"github.com/khades/servbot/models"
	"github.com/khades/servbot/repos"
)

type stream struct {
	CreatedAt time.Time `json:"created_at"`
	Channel   channelInfo
}

type channelInfo struct {
	Status string
	Game   string
}

type responseBody struct {
	Stream stream
}

// CheckStreamStatus runs gettting all data from all channels bot applied to
func CheckStreamStatus() {
	for _, value := range repos.Config.Channels {
		go getStatus(value)
	}
}

func getStatus(channel string) {
	log.Println("Started check")
	var responseBody = new(responseBody)
	url := "https://api.twitch.tv/kraken/streams/" + channel + "?client_id=c5w8oif66lg711pfrh86h8od3sek43d"
	resp, respError := http.Get(url)
	var status = models.StreamStatus{
		Online: false}
	if respError != nil {
		return
	}
	defer resp.Body.Close()

	marshallError := json.NewDecoder(resp.Body).Decode(responseBody)
	log.Println(responseBody)
	if marshallError != nil {
		return
	}
	if !responseBody.Stream.CreatedAt.IsZero() {
		status = models.StreamStatus{
			Online:      true,
			Description: "я кот",
			Game:        responseBody.Stream.Channel.Game,
			Title:       responseBody.Stream.Channel.Status,
			Start:       responseBody.Stream.CreatedAt}
	}
	// else {
	// 	repos.SetStreamStatusOffline(channel)
	// }
	repos.PushStreamStatus(channel, status)

}

// func getOfflineStatus(channel string) {
// 	var responseBody = new(channelInfo)
// 	url := "https://api.twitch.tv/kraken/channels/" + channel
// 	resp, respError := http.Get(url)

// 	if respError != nil {
// 		return
// 	}
// 	defer resp.Body.Close()

// 	marshallError := json.NewDecoder(resp.Body).Decode(responseBody)

// 	if marshallError != nil {
// 		return
// 	}

// 	status := models.StreamStatus{
// 		Online: false,
// 		Game:   responseBody.Game,
// 		Title:  responseBody.Status}
// 	repos.PushStreamStatus(channel, status)

// }
<|endoftext|>


<commit_before>package filestorage

import (
	"encoding/json"
	"fmt"
	"os"
	"time"

	"github.com/adelethalialopez/predict/api"
)

type FileStorage struct {
	Filename string
}

func (fs FileStorage) SavePrediction(p *api.Prediction) error {
	file, err := os.OpenFile(fs.Filename, os.O_APPEND|os.O_WRONLY, os.ModeAppend)
	if err != nil {
		file, err = os.Create(fs.Filename)
		fmt.Printf("History file not found, so creating: %s\n", fs.Filename)
	}

	defer file.Close()

	if p.CreatedAt == nil {
		now := time.Now()
		p.CreatedAt = &now
	}

	bytes, err := json.Marshal(p)
	if err != nil {
		return err
	}
	_, err = file.Write(bytes)
	if err != nil {
		return err
	}

	return nil
}

func (fs FileStorage) UpdatePrediction(p *api.Prediction) (*api.Prediction, error) {
	return nil, nil
}

func (fs FileStorage) GetPrediction(id string) (*api.Prediction, error) {
	return nil, nil
}

func (fs FileStorage) GetPredictions(p *api.Prediction) ([]api.Prediction, error) {
	return nil, nil
}
<commit_msg>start fleshing out update prediction file storage<commit_after>package filestorage

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"time"

	"github.com/adelethalialopez/predict/api"
	"github.com/pborman/uuid"
)

type FileStorage struct {
	Filename string
}

func (fs FileStorage) SavePrediction(p *api.Prediction) error {
	file, err := os.OpenFile(fs.Filename, os.O_APPEND|os.O_WRONLY, os.ModeAppend)
	if err != nil {
		file, err = os.Create(fs.Filename)
		fmt.Printf("History file not found, so creating: %s\n", fs.Filename)
	}

	defer file.Close()

	p.ID = uuid.NewRandom().String()

	if p.CreatedAt == nil {
		now := time.Now()
		p.CreatedAt = &now
	}

	bytes, err := json.Marshal(p)
	if err != nil {
		return err
	}
	_, err = file.Write(bytes)
	if err != nil {
		return err
	}

	return nil
}

func (fs FileStorage) UpdatePrediction(p *api.Prediction) (*api.Prediction, error) {
	file, err := os.OpenFile(fs.Filename, os.O_APPEND|os.O_RDWR, os.ModeAppend)
	if err != nil {
		fmt.Printf("No history file at: %s\n", fs.Filename)
	}

	dec := json.NewDecoder(file)
	ps := make([]api.Prediction, 0, 20)

	for {
		var np api.Prediction

		if err := dec.Decode(&np); err == io.EOF {
			break
		} else if err != nil {
			log.Fatal(err)
		}
		ps = append(ps, np)
	}
	return nil, nil
}

func (fs FileStorage) GetPrediction(id string) (*api.Prediction, error) {
	return nil, nil
}

func (fs FileStorage) GetPredictions(p *api.Prediction) ([]api.Prediction, error) {
	return nil, nil
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 Robin Engel
 *
 * 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 context

import (
	"gopkg.in/mgo.v2"
	"time"
	"github.com/bluedevel/mosel/api"
	"gopkg.in/mgo.v2/bson"
	"log"
)

type nodeDoc struct {
	Name     string `bson:"name"`
	Diagrams map[string]diagramDoc `bson:"diagrams"`
}

type diagramDoc struct {
	Name   string `bson:"name"`
	Graphs map[string]graphDoc `bson:"graphs"`
}

type graphDoc struct {
	Name       string `bson:"name"`
	DataPoints []dataPointDoc `bson:"dataPoints"`
}

type dataPointDoc struct {
	Time  time.Time `bson:"time"`
	Value string `bson:"value"`
}

type mongoDataPersistence struct {
	session  *mgo.Session
	database *mgo.Database
}

func NewMongoDataPersistence(session *mgo.Session) *mongoDataPersistence {
	return &mongoDataPersistence{
		session:session,
	}
}

func (pers *mongoDataPersistence) Init() error {
	pers.database = pers.session.DB("")
	return nil
}

func (pers *mongoDataPersistence) Add(nodeName string, t time.Time, info api.NodeInfo) {
	coll := pers.database.C("nodes")

	selector := bson.M{"name": nodeName}

	var node nodeDoc
	itr := coll.Find(selector).Iter()

	doUpdate := itr.Next(&node)

	if err := itr.Err(); err != nil {
		log.Fatal(err)
		return
	}

	initNodeIfNil(nodeName, &node)

	for diagramName, graphs := range info {
		var diagram diagramDoc
		initDiagramIfNil(diagramName, &node, &diagram)

		for graphName, value := range graphs {
			var graph graphDoc
			initGraphIfNil(graphName, &diagram, &graph)

			point := dataPointDoc{
				Time:t,
				Value:value,
			}

			graph.DataPoints = append(graph.DataPoints, point)
		}
	}

	if doUpdate {
		coll.Update(selector, node)
	} else {
		coll.Insert(node)
	}
}

func (pers *mongoDataPersistence) GetAll() (DataCacheStorage, error) {
	return nil, nil
}

func (pers *mongoDataPersistence) GetAllSince(since time.Duration) (DataCacheStorage, error) {
	return nil, nil
}

func initNodeIfNil(name string, node *nodeDoc) {
	if node.Diagrams == nil {
		node.Diagrams = make(map[string]diagramDoc)
	}
}

func initDiagramIfNil(name string, node *nodeDoc, dia *diagramDoc) {
	if _, ok := node.Diagrams[name]; !ok {
		dia.Name = name
		dia.Graphs = make(map[string]graphDoc)
		node.Diagrams[name] = *dia
	} else {
		(*dia) = node.Diagrams[name]
	}
}

func initGraphIfNil(name string, dia *diagramDoc, gr *graphDoc) {
	if _, ok := dia.Graphs[name]; !ok {
		gr.Name = name
		gr.DataPoints = make([]dataPointDoc, 0)
		dia.Graphs[name] = *gr
	} else {
		(*gr) = dia.Graphs[name]
	}
}<commit_msg>fix missing node name in mongo<commit_after>/*
 * Copyright 2017 Robin Engel
 *
 * 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 context

import (
	"gopkg.in/mgo.v2"
	"time"
	"github.com/bluedevel/mosel/api"
	"gopkg.in/mgo.v2/bson"
	"log"
)

type nodeDoc struct {
	Name     string `bson:"name"`
	Diagrams map[string]diagramDoc `bson:"diagrams"`
}

type diagramDoc struct {
	Name   string `bson:"name"`
	Graphs map[string]graphDoc `bson:"graphs"`
}

type graphDoc struct {
	Name       string `bson:"name"`
	DataPoints []dataPointDoc `bson:"dataPoints"`
}

type dataPointDoc struct {
	Time  time.Time `bson:"time"`
	Value string `bson:"value"`
}

type mongoDataPersistence struct {
	session  *mgo.Session
	database *mgo.Database
}

func NewMongoDataPersistence(session *mgo.Session) *mongoDataPersistence {
	return &mongoDataPersistence{
		session:session,
	}
}

func (pers *mongoDataPersistence) Init() error {
	pers.database = pers.session.DB("")
	return nil
}

func (pers *mongoDataPersistence) Add(nodeName string, t time.Time, info api.NodeInfo) {
	coll := pers.database.C("nodes")

	selector := bson.M{"name": nodeName}

	var node nodeDoc
	itr := coll.Find(selector).Iter()

	doUpdate := itr.Next(&node)

	if err := itr.Err(); err != nil {
		log.Fatal(err)
		return
	}

	initNodeIfNil(nodeName, &node)

	for diagramName, graphs := range info {
		var diagram diagramDoc
		initDiagramIfNil(diagramName, &node, &diagram)

		for graphName, value := range graphs {
			var graph graphDoc
			initGraphIfNil(graphName, &diagram, &graph)

			point := dataPointDoc{
				Time:t,
				Value:value,
			}

			graph.DataPoints = append(graph.DataPoints, point)
		}
	}

	if doUpdate {
		coll.Update(selector, node)
	} else {
		coll.Insert(node)
	}
}

func (pers *mongoDataPersistence) GetAll() (DataCacheStorage, error) {
	return nil, nil
}

func (pers *mongoDataPersistence) GetAllSince(since time.Duration) (DataCacheStorage, error) {
	return nil, nil
}

func initNodeIfNil(name string, node *nodeDoc) {
	if node.Name == "" {
		node.Name = name
	}
	if node.Diagrams == nil {
		node.Diagrams = make(map[string]diagramDoc)
	}
}

func initDiagramIfNil(name string, node *nodeDoc, dia *diagramDoc) {
	if _, ok := node.Diagrams[name]; !ok {
		dia.Name = name
		dia.Graphs = make(map[string]graphDoc)
		node.Diagrams[name] = *dia
	} else {
		(*dia) = node.Diagrams[name]
	}
}

func initGraphIfNil(name string, dia *diagramDoc, gr *graphDoc) {
	if _, ok := dia.Graphs[name]; !ok {
		gr.Name = name
		gr.DataPoints = make([]dataPointDoc, 0)
		dia.Graphs[name] = *gr
	} else {
		(*gr) = dia.Graphs[name]
	}
}<|endoftext|>


<commit_before>/*
 * Copyright 2017 Robin Engel
 *
 * 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 context

import (
	"gopkg.in/mgo.v2"
	"time"
	"github.com/bluedevel/mosel/api"
	"gopkg.in/mgo.v2/bson"
	"log"
)

type nodeDoc struct {
	Name     string `bson:"name"`
	Diagrams map[string]diagramDoc `bson:"diagrams"`
}

type diagramDoc struct {
	Name   string `bson:"name"`
	Graphs map[string]graphDoc `bson:"graphs"`
}

type graphDoc struct {
	Name       string `bson:"name"`
	DataPoints []dataPointDoc `bson:"dataPoints"`
}

type dataPointDoc struct {
	Time  time.Time `bson:"time"`
	Value string `bson:"value"`
}

type mongoDataPersistence struct {
	session  *mgo.Session
	database *mgo.Database
}

func NewMongoDataPersistence(session *mgo.Session) *mongoDataPersistence {
	return &mongoDataPersistence{
		session:session,
	}
}

func (pers *mongoDataPersistence) Init() error {
	pers.database = pers.session.DB("")
	return nil
}

func (pers *mongoDataPersistence) Add(nodeName string, t time.Time, info api.NodeInfo) {
	coll := pers.database.C("nodes")

	selector := bson.M{"name": nodeName}

	var node nodeDoc
	itr := coll.Find(selector).Iter()

	doUpdate := itr.Next(&node)

	if err := itr.Err(); err != nil {
		log.Fatal(err)
		return
	}

	initNodeIfNil(nodeName, &node)

	for diagramName, graphs := range info {
		var diagram diagramDoc
		initDiagramIfNil(diagramName, &node, &diagram)

		for graphName, value := range graphs {
			var graph graphDoc

			if g, ok := diagram.Graphs[graphName]; ok {
				graph = g
			} else {
				graph = graphDoc{
					Name:graphName,
					DataPoints:make([]dataPointDoc, 0),
				}
			}

			point := dataPointDoc{
				Time:t,
				Value:value,
			}

			graph.DataPoints = append(graph.DataPoints, point)
			diagram.Graphs[graphName] = graph
		}
	}

	if doUpdate {
		coll.Update(selector, node)
	} else {
		coll.Insert(node)
	}
}

func (pers *mongoDataPersistence) GetAll() (DataCacheStorage, error) {
	return nil, nil
}

func (pers *mongoDataPersistence) GetAllSince(since time.Duration) (DataCacheStorage, error) {
	return nil, nil
}

func initNodeIfNil(name string, node *nodeDoc) {
	if node.Name == "" {
		node.Name = name
	}
	if node.Diagrams == nil {
		node.Diagrams = make(map[string]diagramDoc)
	}
}

func initDiagramIfNil(name string, node *nodeDoc, dia *diagramDoc) {
	if _, ok := node.Diagrams[name]; !ok {
		dia.Name = name
		dia.Graphs = make(map[string]graphDoc)
		node.Diagrams[name] = *dia
	} else {
		(*dia) = node.Diagrams[name]
	}
}<commit_msg>refactored mongo data model<commit_after>/*
 * Copyright 2017 Robin Engel
 *
 * 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 context

import (
	"gopkg.in/mgo.v2"
	"time"
	"github.com/bluedevel/mosel/api"
	"gopkg.in/mgo.v2/bson"
	"log"
)

type nodeDoc struct {
	Name     string `bson:"name"`
	Diagrams []diagramDoc `bson:"diagrams"`
}

type diagramDoc struct {
	Name   string `bson:"name"`
	Graphs []graphDoc `bson:"graphs"`
}

type graphDoc struct {
	Name       string `bson:"name"`
	DataPoints []dataPointDoc `bson:"dataPoints"`
}

type dataPointDoc struct {
	Time  time.Time `bson:"time"`
	Value string `bson:"value"`
}

type mongoDataPersistence struct {
	session  *mgo.Session
	database *mgo.Database
}

func NewMongoDataPersistence(session *mgo.Session) *mongoDataPersistence {
	return &mongoDataPersistence{
		session:session,
	}
}

func (pers *mongoDataPersistence) Init() error {
	pers.database = pers.session.DB("")
	return nil
}

func (pers *mongoDataPersistence) Add(nodeName string, t time.Time, info api.NodeInfo) {
	coll := pers.database.C("nodes")

	selector := bson.M{"name": nodeName}

	var node nodeDoc
	itr := coll.Find(selector).Iter()

	doUpdate := itr.Next(&node)

	if err := itr.Err(); err != nil {
		log.Fatal(err)
		return
	}

	if node.Name == "" {
		node.Name = nodeName
	}

	if node.Diagrams == nil {
		node.Diagrams = make([]diagramDoc, len(info))
	}

	for diagramName, graphs := range info {
		diaIndex := findDiagramByName(diagramName, node.Diagrams)

		var dia diagramDoc
		if diaIndex == -1 {
			dia = diagramDoc{
				Name:diagramName,
				Graphs: make([]graphDoc, 0),
			}
		} else {
			dia = node.Diagrams[diaIndex]
		}

		for graphName, value := range graphs {
			graphIndex := findGraphByName(graphName, dia.Graphs)

			var graph graphDoc
			if graphIndex == -1 {
				graph = graphDoc{
					Name:graphName,
					DataPoints: make([]dataPointDoc, 0),
				}
			} else {
				graph = dia.Graphs[graphIndex]
			}

			point := dataPointDoc{
				Time:t,
				Value:value,
			}

			graph.DataPoints = append(graph.DataPoints, point)

			if graphIndex == -1 {
				dia.Graphs = append(dia.Graphs, graph)
			} else {
				dia.Graphs[graphIndex] = graph
			}
		}

		if diaIndex == -1 {
			node.Diagrams = append(node.Diagrams, dia)
		} else {
			node.Diagrams[diaIndex] = dia
		}
	}

	if doUpdate {
		coll.Update(selector, node)
	} else {
		coll.Insert(node)
	}
}

func (pers *mongoDataPersistence) GetAll() (DataCacheStorage, error) {
	return nil, nil
}

func (pers *mongoDataPersistence) GetAllSince(since time.Duration) (DataCacheStorage, error) {
	return nil, nil
}

func findDiagramByName(name string, diagrams []diagramDoc) int {
	for i, dia := range diagrams {
		if dia.Name == name {
			return i
		}
	}
	return -1
}

func findGraphByName(name string, graphs []graphDoc) int {
	for i, gr := range graphs {
		if gr.Name == name {
			return i
		}
	}
	return -1
}<|endoftext|>


<commit_before>/*
 * Copyright 2017 Robin Engel
 *
 * 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 context

import (
	"gopkg.in/mgo.v2"
	"time"
	"github.com/bluedevel/mosel/api"
	"gopkg.in/mgo.v2/bson"
	"log"
)

type nodeDoc struct {
	Name     string `bson:"name"`
	Diagrams map[string]diagramDoc `bson:"diagrams"`
}

type diagramDoc struct {
	Name   string `bson:"name"`
	Graphs map[string]graphDoc `bson:"graphs"`
}

type graphDoc struct {
	Name       string `bson:"name"`
	DataPoints []dataPointDoc `bson:"dataPoints"`
}

type dataPointDoc struct {
	Time  time.Time `bson:"time"`
	Value string `bson:"value"`
}

type mongoDataPersistence struct {
	session  *mgo.Session
	database *mgo.Database
}

func NewMongoDataPersistence(session *mgo.Session) *mongoDataPersistence {
	return &mongoDataPersistence{
		session:session,
	}
}

func (pers *mongoDataPersistence) Init() error {
	pers.database = pers.session.DB("")
	return nil
}

func (pers *mongoDataPersistence) Add(nodeName string, t time.Time, info api.NodeInfo) {
	coll := pers.database.C("nodes")

	selector := bson.M{"name": nodeName}

	var node nodeDoc
	itr := coll.Find(selector).Iter()

	doUpdate := itr.Next(&node)

	if err := itr.Err(); err != nil {
		log.Fatal(err)
		return
	}

	initNodeIfNil(nodeName, &node)

	for diagramName, graphs := range info {
		var diagram diagramDoc
		initDiagramIfNil(diagramName, &node, &diagram)

		for graphName, value := range graphs {
			var graph graphDoc
			initGraphIfNil(graphName, &diagram, &graph)

			point := dataPointDoc{
				Time:t,
				Value:value,
			}

			graph.DataPoints = append(graph.DataPoints, point)
		}
	}

	if doUpdate {
		coll.Update(selector, node)
	} else {
		coll.Insert(node)
	}
}

func (pers *mongoDataPersistence) GetAll() (DataCacheStorage, error) {
	return nil, nil
}

func (pers *mongoDataPersistence) GetAllSince(since time.Duration) (DataCacheStorage, error) {
	return nil, nil
}

func initNodeIfNil(name string, node *nodeDoc) {
	if node.Name == "" {
		node.Name = name
	}
	if node.Diagrams == nil {
		node.Diagrams = make(map[string]diagramDoc)
	}
}

func initDiagramIfNil(name string, node *nodeDoc, dia *diagramDoc) {
	if _, ok := node.Diagrams[name]; !ok {
		dia.Name = name
		dia.Graphs = make(map[string]graphDoc)
		node.Diagrams[name] = *dia
	} else {
		(*dia) = node.Diagrams[name]
	}
}

func initGraphIfNil(name string, dia *diagramDoc, gr *graphDoc) {
	if _, ok := dia.Graphs[name]; !ok {
		gr.Name = name
		gr.DataPoints = make([]dataPointDoc, 0)
		dia.Graphs[name] = *gr
	} else {
		(*gr) = dia.Graphs[name]
	}
}<commit_msg>inserting data to mongodb!!! :)<commit_after>/*
 * Copyright 2017 Robin Engel
 *
 * 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 context

import (
	"gopkg.in/mgo.v2"
	"time"
	"github.com/bluedevel/mosel/api"
	"gopkg.in/mgo.v2/bson"
	"log"
)

type nodeDoc struct {
	Name     string `bson:"name"`
	Diagrams map[string]diagramDoc `bson:"diagrams"`
}

type diagramDoc struct {
	Name   string `bson:"name"`
	Graphs map[string]graphDoc `bson:"graphs"`
}

type graphDoc struct {
	Name       string `bson:"name"`
	DataPoints []dataPointDoc `bson:"dataPoints"`
}

type dataPointDoc struct {
	Time  time.Time `bson:"time"`
	Value string `bson:"value"`
}

type mongoDataPersistence struct {
	session  *mgo.Session
	database *mgo.Database
}

func NewMongoDataPersistence(session *mgo.Session) *mongoDataPersistence {
	return &mongoDataPersistence{
		session:session,
	}
}

func (pers *mongoDataPersistence) Init() error {
	pers.database = pers.session.DB("")
	return nil
}

func (pers *mongoDataPersistence) Add(nodeName string, t time.Time, info api.NodeInfo) {
	coll := pers.database.C("nodes")

	selector := bson.M{"name": nodeName}

	var node nodeDoc
	itr := coll.Find(selector).Iter()

	doUpdate := itr.Next(&node)

	if err := itr.Err(); err != nil {
		log.Fatal(err)
		return
	}

	initNodeIfNil(nodeName, &node)

	for diagramName, graphs := range info {
		var diagram diagramDoc
		initDiagramIfNil(diagramName, &node, &diagram)

		for graphName, value := range graphs {
			var graph graphDoc

			if g, ok := diagram.Graphs[graphName]; ok {
				graph = g
			} else {
				graph = graphDoc{
					Name:graphName,
					DataPoints:make([]dataPointDoc, 0),
				}
			}

			point := dataPointDoc{
				Time:t,
				Value:value,
			}

			graph.DataPoints = append(graph.DataPoints, point)
			diagram.Graphs[graphName] = graph
		}
	}

	if doUpdate {
		coll.Update(selector, node)
	} else {
		coll.Insert(node)
	}
}

func (pers *mongoDataPersistence) GetAll() (DataCacheStorage, error) {
	return nil, nil
}

func (pers *mongoDataPersistence) GetAllSince(since time.Duration) (DataCacheStorage, error) {
	return nil, nil
}

func initNodeIfNil(name string, node *nodeDoc) {
	if node.Name == "" {
		node.Name = name
	}
	if node.Diagrams == nil {
		node.Diagrams = make(map[string]diagramDoc)
	}
}

func initDiagramIfNil(name string, node *nodeDoc, dia *diagramDoc) {
	if _, ok := node.Diagrams[name]; !ok {
		dia.Name = name
		dia.Graphs = make(map[string]graphDoc)
		node.Diagrams[name] = *dia
	} else {
		(*dia) = node.Diagrams[name]
	}
}<|endoftext|>


<commit_before>package model

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
)

type StatusCode int

func (c StatusCode) Value() (driver.Value, error) {
	return []byte(c.String()), nil
}

func (c *StatusCode) Scan(src interface{}) error {
	switch t := src.(type) {
	case []byte:
		code, err := statusStringToCode(string(t))
		if err != nil {
			return err
		}
		*c = code
	case nil:
		*c = StatusInit
	default:
		return errors.New("Incompatible type for StatusCode")
	}
	return nil
}

func (c StatusCode) String() string {
	return statusCodeMap[c]
}

func (c StatusCode) MarshalJSON() ([]byte, error) {
	return json.Marshal(c.String())
}

func (c *StatusCode) UnmarshalJSON(data []byte) error {
	if c == nil {
		return errors.New("StatusCode: UnmarshalJSON on nil pointer")
	}
	code, err := statusStringToCode(string(data))
	if err != nil {
		return err
	}
	*c = code
	return nil
}

func statusStringToCode(status string) (StatusCode, error) {
	for k, v := range statusCodeMap {
		if v == status {
			return k, nil
		}
	}
	return 0, fmt.Errorf("StatusCode %s not exists", status)
}

const (
	// Default status, This should not be exported to client
	StatusInit StatusCode = iota
	// Received your API request to send a message
	StatusAccepted
	// The message is queued to be sent out
	StatusQueued
	// The message is in the process of dispatching to the upstream carrier
	StatusSending
	// The message could not be sent to the upstream carrier
	StatusFailed
	// The message was successfully accepted by the upstream carrie
	StatusSent
	// Received an undocumented status code from the upstream carrier
	StatusUnknown
	// Received that the message was not delivered from the upstream carrier
	StatusUndelivered
	// Received confirmation of message delivery from the upstream carrier
	StatusDelivered
)

var statusCodeMap = map[StatusCode]string{
	StatusInit:        "init",
	StatusAccepted:    "accepted",
	StatusQueued:      "queued",
	StatusSending:     "sending",
	StatusFailed:      "failed",
	StatusSent:        "sent",
	StatusUnknown:     "unknown",
	StatusUndelivered: "undelivered",
	StatusDelivered:   "delivered",
}
<commit_msg>Fix statusStringToCode<commit_after>package model

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
)

type StatusCode int

func (c StatusCode) Value() (driver.Value, error) {
	return []byte(c.String()), nil
}

func (c *StatusCode) Scan(src interface{}) error {
	switch t := src.(type) {
	case []byte:
		code, err := statusStringToCode(string(t))
		if err != nil {
			return err
		}
		*c = code
	case nil:
		*c = StatusInit
	default:
		return errors.New("Incompatible type for StatusCode")
	}
	return nil
}

func (c StatusCode) String() string {
	return statusCodeMap[c]
}

func (c StatusCode) MarshalJSON() ([]byte, error) {
	return json.Marshal(c.String())
}

func (c *StatusCode) UnmarshalJSON(data []byte) error {
	if c == nil {
		return errors.New("StatusCode: UnmarshalJSON on nil pointer")
	}
	var code string
	err := json.Unmarshal(data, &code)
	if err != nil {
		return err
	}
	statusCode, err := statusStringToCode(code)
	if err != nil {
		return err
	}
	*c = statusCode
	return nil
}

func statusStringToCode(status string) (StatusCode, error) {
	for k, v := range statusCodeMap {
		if v == status {
			return k, nil
		}
	}
	return 0, fmt.Errorf("StatusCode %s not exists", status)
}

const (
	// Default status, This should not be exported to client
	StatusInit StatusCode = iota
	// Received your API request to send a message
	StatusAccepted
	// The message is queued to be sent out
	StatusQueued
	// The message is in the process of dispatching to the upstream carrier
	StatusSending
	// The message could not be sent to the upstream carrier
	StatusFailed
	// The message was successfully accepted by the upstream carrie
	StatusSent
	// Received an undocumented status code from the upstream carrier
	StatusUnknown
	// Received that the message was not delivered from the upstream carrier
	StatusUndelivered
	// Received confirmation of message delivery from the upstream carrier
	StatusDelivered
)

var statusCodeMap = map[StatusCode]string{
	StatusInit:        "init",
	StatusAccepted:    "accepted",
	StatusQueued:      "queued",
	StatusSending:     "sending",
	StatusFailed:      "failed",
	StatusSent:        "sent",
	StatusUnknown:     "unknown",
	StatusUndelivered: "undelivered",
	StatusDelivered:   "delivered",
}
<|endoftext|>


<commit_before>package graph

import (
	"erx"
	"runtime"
	"strings"
)

type NodeId uint

type Nodes []NodeId

type Connection struct {
	Tail NodeId
	Head NodeId
}

type ConnectionsIterable interface {
	ConnectionsIter() <-chan Connection
}

type NodesIterable interface {
	NodesIter() <-chan NodeId
}

type GraphNodesWriter interface {
	// Adding single node to graph
	AddNode(node NodeId) erx.Error
}

type GraphNodesReader interface {
	// Getting nodes count in graph
	NodesCnt() int
}

type GraphNodesRemover interface {
	// Removing node from graph
	RemoveNode(node NodeId) erx.Error
}

type DirectedGraphArcsWriter interface {
	// Adding directed arc to graph
	AddArc(from, to NodeId) erx.Error
}

type DirectedGraphArcsRemover interface {
	// Removding directed arc
	RemoveArc(from, to NodeId) erx.Error
}

type DirectedGraphReader interface {
	// Getting arcs count in graph
	ArcsCnt() int

	// Getting all graph sources.
	GetSources() (Nodes, erx.Error)

	// Getting all graph sinks.
	GetSinks() (Nodes, erx.Error)

	// Getting node accessors
	GetAccessors(node NodeId) (Nodes, erx.Error)

	// Getting node predecessors
	GetPredecessors(node NodeId) (Nodes, erx.Error)

	// Checking arrow existance between node1 and node2
	//
	// node1 and node2 must exist in graph or error will be returned
	CheckArc(node1, node2 NodeId) (bool, erx.Error)

}

// Interface representing directed graph
type DirectedGraph interface {
	GraphNodesWriter
	GraphNodesReader
	GraphNodesRemover

	ConnectionsIterable
	NodesIterable
	DirectedGraphArcsWriter
	DirectedGraphArcsRemover
	DirectedGraphReader
}

type UndirectedGraphReader interface {
	// Arrows count in graph
	EdgesCnt() int

	// Checking edge existance between node1 and node2
	//
	// node1 and node2 must exist in graph or error will be returned
	CheckEdge(node1, node2 NodeId) (bool, erx.Error)

	// Getting all nodes, connected to given one
	GetNeighbours(node NodeId) (Nodes, erx.Error)
}

type UndirectedGraphEdgesWriter interface {
	// Adding new edge to graph
	AddEdge(node1, node2 NodeId) (erx.Error)
}

type UndirectedGraphEdgesRemover interface {
	// Removing edge, connecting node1 and node2
	RemoveEdge(node1, node2 NodeId) erx.Error
}

// Interface representing undirected graph
type UndirectedGraph interface {
	GraphNodesWriter
	GraphNodesReader
	GraphNodesRemover

	ConnectionsIterable
	NodesIterable

	UndirectedGraphEdgesWriter
	UndirectedGraphEdgesRemover
	UndirectedGraphReader
}

type MixedGraph interface {
	GraphNodesWriter
	GraphNodesReader
	GraphNodesRemover

	ConnectionsIterable
	NodesIterable

	UndirectedGraphEdgesWriter
	UndirectedGraphEdgesRemover
	UndirectedGraphReader

	DirectedGraphArcsWriter
	DirectedGraphArcsRemover
	DirectedGraphReader

	ConnectionsIterable

	// Iterate over only undirected edges
	EdgesIter() ConnectionsIterable

	// Iterate over only directed arcs
	ArcsIter() ConnectionsIterable
}

func init() {
	// adding to erx directory prefix to cut from file names
	_, file, _, _ := runtime.Caller(0)
	dirName := file[0:strings.LastIndex(file, "/")]
	prevDirName := dirName[0:strings.LastIndex(dirName, "/")+1]
	erx.AddPathCut(prevDirName)
}
<commit_msg>small change in MixedGraph interface<commit_after>package graph

import (
	"erx"
	"runtime"
	"strings"
)

type NodeId uint

type Nodes []NodeId

type Connection struct {
	Tail NodeId
	Head NodeId
}

type ConnectionsIterable interface {
	ConnectionsIter() <-chan Connection
}

type NodesIterable interface {
	NodesIter() <-chan NodeId
}

type GraphNodesWriter interface {
	// Adding single node to graph
	AddNode(node NodeId) erx.Error
}

type GraphNodesReader interface {
	// Getting nodes count in graph
	NodesCnt() int
}

type GraphNodesRemover interface {
	// Removing node from graph
	RemoveNode(node NodeId) erx.Error
}

type DirectedGraphArcsWriter interface {
	// Adding directed arc to graph
	AddArc(from, to NodeId) erx.Error
}

type DirectedGraphArcsRemover interface {
	// Removding directed arc
	RemoveArc(from, to NodeId) erx.Error
}

type DirectedGraphReader interface {
	// Getting arcs count in graph
	ArcsCnt() int

	// Getting all graph sources.
	GetSources() (Nodes, erx.Error)

	// Getting all graph sinks.
	GetSinks() (Nodes, erx.Error)

	// Getting node accessors
	GetAccessors(node NodeId) (Nodes, erx.Error)

	// Getting node predecessors
	GetPredecessors(node NodeId) (Nodes, erx.Error)

	// Checking arrow existance between node1 and node2
	//
	// node1 and node2 must exist in graph or error will be returned
	CheckArc(node1, node2 NodeId) (bool, erx.Error)

}

// Interface representing directed graph
type DirectedGraph interface {
	GraphNodesWriter
	GraphNodesReader
	GraphNodesRemover

	ConnectionsIterable
	NodesIterable
	DirectedGraphArcsWriter
	DirectedGraphArcsRemover
	DirectedGraphReader
}

type UndirectedGraphReader interface {
	// Arrows count in graph
	EdgesCnt() int

	// Checking edge existance between node1 and node2
	//
	// node1 and node2 must exist in graph or error will be returned
	CheckEdge(node1, node2 NodeId) (bool, erx.Error)

	// Getting all nodes, connected to given one
	GetNeighbours(node NodeId) (Nodes, erx.Error)
}

type UndirectedGraphEdgesWriter interface {
	// Adding new edge to graph
	AddEdge(node1, node2 NodeId) (erx.Error)
}

type UndirectedGraphEdgesRemover interface {
	// Removing edge, connecting node1 and node2
	RemoveEdge(node1, node2 NodeId) erx.Error
}

// Interface representing undirected graph
type UndirectedGraph interface {
	GraphNodesWriter
	GraphNodesReader
	GraphNodesRemover

	ConnectionsIterable
	NodesIterable

	UndirectedGraphEdgesWriter
	UndirectedGraphEdgesRemover
	UndirectedGraphReader
}

type MixedGraph interface {
	GraphNodesWriter
	GraphNodesReader
	GraphNodesRemover

	ConnectionsIterable
	NodesIterable

	UndirectedGraphEdgesWriter
	UndirectedGraphEdgesRemover
	UndirectedGraphReader

	DirectedGraphArcsWriter
	DirectedGraphArcsRemover
	DirectedGraphReader

	// Iterate over only undirected edges
	EdgesIter() ConnectionsIterable

	// Iterate over only directed arcs
	ArcsIter() ConnectionsIterable
}

func init() {
	// adding to erx directory prefix to cut from file names
	_, file, _, _ := runtime.Caller(0)
	dirName := file[0:strings.LastIndex(file, "/")]
	prevDirName := dirName[0:strings.LastIndex(dirName, "/")+1]
	erx.AddPathCut(prevDirName)
}
<|endoftext|>


<commit_before>package commands

import (
	"log"
	"os"
	"sort"
	"strconv"

	"github.com/concourse/atc"
	"github.com/concourse/fly/atcclient"
	"github.com/concourse/fly/rc"
	"github.com/concourse/fly/ui"
	"github.com/fatih/color"
)

type ContainersCommand struct{}

var containersCommand ContainersCommand

func init() {
	containers, err := Parser.AddCommand(
		"containers",
		"Print the running containers",
		"",
		&containersCommand,
	)
	if err != nil {
		panic(err)
	}

	containers.Aliases = []string{"cs"}
}

func (command *ContainersCommand) Execute([]string) error {
	target, err := rc.SelectTarget(globalOptions.Target, globalOptions.Insecure)
	if err != nil {
		log.Fatalln(err)
	}

	client, err := atcclient.NewClient(*target)
	if err != nil {
		log.Fatalln(err)
	}
	handler := atcclient.NewAtcHandler(client)

	containers, err := handler.ListContainers()
	if err != nil {
		log.Fatalln(err)
	}

	headers := ui.TableRow{
		{Contents: "handle", Color: color.New(color.Bold)},
		{Contents: "name", Color: color.New(color.Bold)},
		{Contents: "pipeline", Color: color.New(color.Bold)},
		{Contents: "type", Color: color.New(color.Bold)},
		{Contents: "build id", Color: color.New(color.Bold)},
		{Contents: "worker", Color: color.New(color.Bold)},
	}

	table := ui.Table{Data: []ui.TableRow{headers}}

	sort.Sort(containersByHandle(containers))

	for _, c := range containers {
		row := ui.TableRow{
			{Contents: c.ID},
			{Contents: c.Name},
			{Contents: c.PipelineName},
			{Contents: c.Type},
			buildIDOrNone(c.BuildID),
			{Contents: c.WorkerName},
		}

		table.Data = append(table.Data, row)
	}

	return table.Render(os.Stdout)
}

type containersByHandle []atc.Container

func (cs containersByHandle) Len() int               { return len(cs) }
func (cs containersByHandle) Swap(i int, j int)      { cs[i], cs[j] = cs[j], cs[i] }
func (cs containersByHandle) Less(i int, j int) bool { return cs[i].ID < cs[j].ID }

func buildIDOrNone(id int) ui.TableCell {
	var column ui.TableCell

	if id == 0 {
		column.Contents = "none"
		column.Color = color.New(color.Faint)
	} else {
		column.Contents = strconv.Itoa(id)
	}

	return column
}
<commit_msg>fix 'containers' command not using headers<commit_after>package commands

import (
	"log"
	"os"
	"sort"
	"strconv"

	"github.com/concourse/atc"
	"github.com/concourse/fly/atcclient"
	"github.com/concourse/fly/rc"
	"github.com/concourse/fly/ui"
	"github.com/fatih/color"
)

type ContainersCommand struct{}

var containersCommand ContainersCommand

func init() {
	containers, err := Parser.AddCommand(
		"containers",
		"Print the running containers",
		"",
		&containersCommand,
	)
	if err != nil {
		panic(err)
	}

	containers.Aliases = []string{"cs"}
}

func (command *ContainersCommand) Execute([]string) error {
	target, err := rc.SelectTarget(globalOptions.Target, globalOptions.Insecure)
	if err != nil {
		log.Fatalln(err)
	}

	client, err := atcclient.NewClient(*target)
	if err != nil {
		log.Fatalln(err)
	}
	handler := atcclient.NewAtcHandler(client)

	containers, err := handler.ListContainers()
	if err != nil {
		log.Fatalln(err)
	}

	table := ui.Table{
		Headers: ui.TableRow{
			{Contents: "handle", Color: color.New(color.Bold)},
			{Contents: "name", Color: color.New(color.Bold)},
			{Contents: "pipeline", Color: color.New(color.Bold)},
			{Contents: "type", Color: color.New(color.Bold)},
			{Contents: "build id", Color: color.New(color.Bold)},
			{Contents: "worker", Color: color.New(color.Bold)},
		},
	}

	sort.Sort(containersByHandle(containers))

	for _, c := range containers {
		row := ui.TableRow{
			{Contents: c.ID},
			{Contents: c.Name},
			{Contents: c.PipelineName},
			{Contents: c.Type},
			buildIDOrNone(c.BuildID),
			{Contents: c.WorkerName},
		}

		table.Data = append(table.Data, row)
	}

	return table.Render(os.Stdout)
}

type containersByHandle []atc.Container

func (cs containersByHandle) Len() int               { return len(cs) }
func (cs containersByHandle) Swap(i int, j int)      { cs[i], cs[j] = cs[j], cs[i] }
func (cs containersByHandle) Less(i int, j int) bool { return cs[i].ID < cs[j].ID }

func buildIDOrNone(id int) ui.TableCell {
	var column ui.TableCell

	if id == 0 {
		column.Contents = "none"
		column.Color = color.New(color.Faint)
	} else {
		column.Contents = strconv.Itoa(id)
	}

	return column
}
<|endoftext|>


<commit_before>package runner

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/grafana/metrictank/cmd/mt-simulate-lock-congestion/metricname"
	"github.com/grafana/metrictank/cmd/mt-simulate-lock-congestion/query"

	"github.com/grafana/metrictank/idx"
	"github.com/grafana/metrictank/idx/memory"
	"github.com/raintank/schema"
	"golang.org/x/sync/errgroup"
)

// TestRun represents one test run, including all the structures used for the test
type TestRun struct {
	index                 idx.MetricIndex
	metricNameGenerator   metricname.NameGenerator
	queryPatternGenerator query.QueryGenerator
	stats                 runStats
	addsPerSec            uint32 // how many adds per second we want to execute. 0 means unlimited, as many as possible
	addThreads            uint32 // number of concurrent add threads
	addSampleFactor       uint32 // how often to print info about a completed add
	addDelay              uint32 // once the benchmark starts the querying threads start, if addDelay is >0 then the adding into the index is delayed by the given number of seconds
	initialIndexSize      uint32 // prepopulate the index with the defined number of entries before starting the actual test run
	queriesPerSec         uint32 // how many queries per second we want to execute. 0 means unlimited, as many as possible
	querySampleFactor     uint32 // how often to print info about a completed query
	totalQueryCount       uint32 // the total number of queries we will start
	runDuration           time.Duration
}

const orgID = 1

// NewTestRun Instantiates a new test run
func NewTestRun(nameGenerator metricname.NameGenerator, queryGenerator query.QueryGenerator, addDelay, addsPerSec, addThreads, addSampleFactor, initialIndexSize, queriesPerSec, querySampleFactor uint32, runDuration time.Duration) *TestRun {
	totalQueryCount := (queriesPerSec + addDelay) * uint32(runDuration.Seconds())

	runner := TestRun{
		stats:                 newRunStats(totalQueryCount),
		totalQueryCount:       totalQueryCount,
		index:                 memory.New(),
		metricNameGenerator:   nameGenerator,
		queryPatternGenerator: queryGenerator,
		addsPerSec:            addsPerSec,
		addThreads:            addThreads,
		addSampleFactor:       addSampleFactor,
		addDelay:              addDelay,
		initialIndexSize:      initialIndexSize,
		queriesPerSec:         queriesPerSec,
		querySampleFactor:     querySampleFactor,
		runDuration:           runDuration,
	}

	return &runner
}

// Run executes the run
func (t *TestRun) Run() {
	mainCtx, cancel := context.WithCancel(context.Background())

	log.Printf("Starting the metric name generator")
	t.metricNameGenerator.Start(mainCtx, t.addThreads)

	log.Printf("Starting the query generator")
	t.queryPatternGenerator.Start()

	log.Printf("Prepopulating the index with %d entries", t.initialIndexSize)
	t.prepopulateIndex()

	// wait group to wait until all routines have been started
	startedWg := sync.WaitGroup{}
	workerThreads, workerCtx := errgroup.WithContext(mainCtx)

	startedWg.Add(1)
	workerThreads.Go(t.queryRoutine(workerCtx, &startedWg))

	<-time.After(time.Second * time.Duration(t.addDelay))

	startedWg.Add(int(t.addThreads))
	for i := uint32(0); i < t.addThreads; i++ {
		workerThreads.Go(t.addRoutine(workerCtx, &startedWg, i))
	}
	startedWg.Wait()
	log.Printf("Benchmark has started")

	// as soon as all routines have started, we start the timer
	<-time.After(t.runDuration)
	cancel()
	workerThreads.Wait()
}

// PrintStats writes all the statistics in human readable format into stdout
func (t *TestRun) PrintStats() {
	t.stats.Print(uint32(t.runDuration.Seconds()))
}

func (t *TestRun) prepopulateIndex() {
	for i := uint32(0); i < t.initialIndexSize; i++ {
		md := schema.MetricData{OrgId: orgID, Interval: 1, Value: 2, Time: int64(time.Now().Unix()), Mtype: "gauge"}
		md.Name = t.metricNameGenerator.GetNewMetricName()
		md.SetId()
		key, err := schema.MKeyFromString(md.Id)
		if err != nil {
			log.Fatalf("Unexpected error when generating MKey from ID string: %s", err)
		}

		t.index.AddOrUpdate(key, &md, int32(i%t.addThreads))
	}
}

func (t *TestRun) addRoutine(ctx context.Context, startedWg *sync.WaitGroup, routineID uint32) func() error {
	partitionID := int32(routineID)

	addEntryToIndex := func(time int) error {
		md := schema.MetricData{OrgId: orgID, Interval: 1, Value: 2, Time: int64(time), Mtype: "gauge"}
		md.Name = t.metricNameGenerator.GetNewMetricName()
		md.SetId()
		key, err := schema.MKeyFromString(md.Id)
		if err != nil {
			return fmt.Errorf("Unexpected error when generating MKey from ID string: %s", err)
		}

		t.index.AddOrUpdate(key, &md, partitionID)
		adds := t.stats.incAddsCompleted()
		if adds%t.addSampleFactor == 0 {
			log.Printf("Sample: added metric name to index %s", md.Name)
		}

		return nil
	}

	if t.addsPerSec > 0 {
		interval := time.Duration(int64(1000000000) * int64(t.addThreads) / int64(t.addsPerSec))

		return func() error {
			ticker := time.NewTicker(interval)
			go func() {
				<-ctx.Done()
				ticker.Stop()
			}()

			startedWg.Done()
			for {
				select {
				case <-ctx.Done():
					return nil
				case tick := <-ticker.C:
					if err := addEntryToIndex(tick.Second()); err != nil {
						return err
					}
				}
			}
		}
	}

	return func() error {
		startedWg.Done()

		for i := 0; ; i++ {
			select {
			case <-ctx.Done():
				return nil
			default:
				if err := addEntryToIndex(i); err != nil {
					return err
				}
			}
		}
	}
}

func (t *TestRun) queryRoutine(ctx context.Context, startedWg *sync.WaitGroup) func() error {
	return func() error {
		ticker := time.NewTicker(time.Duration(int64(float64(1000000000) / float64(t.queriesPerSec))))
		startedWg.Done()
		for {
			select {
			case <-ctx.Done():
				return nil
			case <-ticker.C:
				go t.runQuery()
			}
		}
	}
}

func (t *TestRun) runQuery() {
	pattern := t.queryPatternGenerator.GetPattern()
	pre := time.Now()

	queryStartedID := t.stats.incQueriesStarted()
	if queryStartedID > t.totalQueryCount {
		return
	}

	_, err := t.index.Find(orgID, pattern, 0)
	if err != nil {
		log.Printf("Warning: Query failed with error: %s", err)
	}

	queryCompletedID := t.stats.incQueriesCompleted()
	t.stats.addQueryTime(queryCompletedID, time.Now().Sub(pre))
	if queryCompletedID%t.querySampleFactor == 0 {
		log.Printf("Sample: queried for pattern %s", pattern)
	}
}
<commit_msg>print result count<commit_after>package runner

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/grafana/metrictank/cmd/mt-simulate-lock-congestion/metricname"
	"github.com/grafana/metrictank/cmd/mt-simulate-lock-congestion/query"

	"github.com/grafana/metrictank/idx"
	"github.com/grafana/metrictank/idx/memory"
	"github.com/raintank/schema"
	"golang.org/x/sync/errgroup"
)

// TestRun represents one test run, including all the structures used for the test
type TestRun struct {
	index                 idx.MetricIndex
	metricNameGenerator   metricname.NameGenerator
	queryPatternGenerator query.QueryGenerator
	stats                 runStats
	addsPerSec            uint32 // how many adds per second we want to execute. 0 means unlimited, as many as possible
	addThreads            uint32 // number of concurrent add threads
	addSampleFactor       uint32 // how often to print info about a completed add
	addDelay              uint32 // once the benchmark starts the querying threads start, if addDelay is >0 then the adding into the index is delayed by the given number of seconds
	initialIndexSize      uint32 // prepopulate the index with the defined number of entries before starting the actual test run
	queriesPerSec         uint32 // how many queries per second we want to execute. 0 means unlimited, as many as possible
	querySampleFactor     uint32 // how often to print info about a completed query
	totalQueryCount       uint32 // the total number of queries we will start
	runDuration           time.Duration
}

const orgID = 1

// NewTestRun Instantiates a new test run
func NewTestRun(nameGenerator metricname.NameGenerator, queryGenerator query.QueryGenerator, addDelay, addsPerSec, addThreads, addSampleFactor, initialIndexSize, queriesPerSec, querySampleFactor uint32, runDuration time.Duration) *TestRun {
	totalQueryCount := (queriesPerSec + addDelay) * uint32(runDuration.Seconds())

	runner := TestRun{
		stats:                 newRunStats(totalQueryCount),
		totalQueryCount:       totalQueryCount,
		index:                 memory.New(),
		metricNameGenerator:   nameGenerator,
		queryPatternGenerator: queryGenerator,
		addsPerSec:            addsPerSec,
		addThreads:            addThreads,
		addSampleFactor:       addSampleFactor,
		addDelay:              addDelay,
		initialIndexSize:      initialIndexSize,
		queriesPerSec:         queriesPerSec,
		querySampleFactor:     querySampleFactor,
		runDuration:           runDuration,
	}

	return &runner
}

// Run executes the run
func (t *TestRun) Run() {
	mainCtx, cancel := context.WithCancel(context.Background())

	log.Printf("Starting the metric name generator")
	t.metricNameGenerator.Start(mainCtx, t.addThreads)

	log.Printf("Starting the query generator")
	t.queryPatternGenerator.Start()

	log.Printf("Prepopulating the index with %d entries", t.initialIndexSize)
	t.prepopulateIndex()

	// wait group to wait until all routines have been started
	startedWg := sync.WaitGroup{}
	workerThreads, workerCtx := errgroup.WithContext(mainCtx)

	startedWg.Add(1)
	workerThreads.Go(t.queryRoutine(workerCtx, &startedWg))

	<-time.After(time.Second * time.Duration(t.addDelay))

	startedWg.Add(int(t.addThreads))
	for i := uint32(0); i < t.addThreads; i++ {
		workerThreads.Go(t.addRoutine(workerCtx, &startedWg, i))
	}
	startedWg.Wait()
	log.Printf("Benchmark has started")

	// as soon as all routines have started, we start the timer
	<-time.After(t.runDuration)
	cancel()
	workerThreads.Wait()
}

// PrintStats writes all the statistics in human readable format into stdout
func (t *TestRun) PrintStats() {
	t.stats.Print(uint32(t.runDuration.Seconds()))
}

func (t *TestRun) prepopulateIndex() {
	for i := uint32(0); i < t.initialIndexSize; i++ {
		md := schema.MetricData{OrgId: orgID, Interval: 1, Value: 2, Time: int64(time.Now().Unix()), Mtype: "gauge"}
		md.Name = t.metricNameGenerator.GetNewMetricName()
		md.SetId()
		key, err := schema.MKeyFromString(md.Id)
		if err != nil {
			log.Fatalf("Unexpected error when generating MKey from ID string: %s", err)
		}

		t.index.AddOrUpdate(key, &md, int32(i%t.addThreads))
	}
}

func (t *TestRun) addRoutine(ctx context.Context, startedWg *sync.WaitGroup, routineID uint32) func() error {
	partitionID := int32(routineID)

	addEntryToIndex := func(time int) error {
		md := schema.MetricData{OrgId: orgID, Interval: 1, Value: 2, Time: int64(time), Mtype: "gauge"}
		md.Name = t.metricNameGenerator.GetNewMetricName()
		md.SetId()
		key, err := schema.MKeyFromString(md.Id)
		if err != nil {
			return fmt.Errorf("Unexpected error when generating MKey from ID string: %s", err)
		}

		t.index.AddOrUpdate(key, &md, partitionID)
		adds := t.stats.incAddsCompleted()
		if adds%t.addSampleFactor == 0 {
			log.Printf("Sample: added metric name to index %s", md.Name)
		}

		return nil
	}

	if t.addsPerSec > 0 {
		interval := time.Duration(int64(1000000000) * int64(t.addThreads) / int64(t.addsPerSec))

		return func() error {
			ticker := time.NewTicker(interval)
			go func() {
				<-ctx.Done()
				ticker.Stop()
			}()

			startedWg.Done()
			for {
				select {
				case <-ctx.Done():
					return nil
				case tick := <-ticker.C:
					if err := addEntryToIndex(tick.Second()); err != nil {
						return err
					}
				}
			}
		}
	}

	return func() error {
		startedWg.Done()

		for i := 0; ; i++ {
			select {
			case <-ctx.Done():
				return nil
			default:
				if err := addEntryToIndex(i); err != nil {
					return err
				}
			}
		}
	}
}

func (t *TestRun) queryRoutine(ctx context.Context, startedWg *sync.WaitGroup) func() error {
	return func() error {
		ticker := time.NewTicker(time.Duration(int64(float64(1000000000) / float64(t.queriesPerSec))))
		startedWg.Done()
		for {
			select {
			case <-ctx.Done():
				return nil
			case <-ticker.C:
				go t.runQuery()
			}
		}
	}
}

func (t *TestRun) runQuery() {
	pattern := t.queryPatternGenerator.GetPattern()
	pre := time.Now()

	queryStartedID := t.stats.incQueriesStarted()
	if queryStartedID > t.totalQueryCount {
		return
	}

	res, err := t.index.Find(orgID, pattern, 0)
	if err != nil {
		log.Printf("Warning: Query failed with error: %s", err)
	}

	queryCompletedID := t.stats.incQueriesCompleted()
	t.stats.addQueryTime(queryCompletedID, time.Now().Sub(pre))
	if queryCompletedID%t.querySampleFactor == 0 {
		log.Printf("Sample: query got %d results for pattern %s", len(res), pattern)
	}
}
<|endoftext|>