<commit_before>package agent

import (
	log "github.com/coreos/fleet/third_party/github.com/golang/glog"

	"github.com/coreos/fleet/event"
	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/machine"
)

type EventHandler struct {
	agent *Agent
}

func NewEventHandler(agent *Agent) *EventHandler {
	return &EventHandler{agent}
}

func (eh *EventHandler) HandleEventJobOffered(ev event.Event) {
	jo := ev.Payload.(job.JobOffer)
	log.V(1).Infof("EventJobOffered(%s): verifying ability to run Job", jo.Job.Name)

	// Everything we check against could change over time, so we track all
	// offers starting here for future bidding even if we can't bid now
	eh.agent.TrackOffer(jo)

	if eh.agent.AbleToRun(&jo.Job) {
		log.Infof("EventJobOffered(%s): passed all criteria, submitting JobBid", jo.Job.Name)
		eh.agent.Bid(jo.Job.Name)
	} else {
		log.V(1).Infof("EventJobOffered(%s): not all criteria met, not bidding", jo.Job.Name)
	}
}

func (eh *EventHandler) HandleEventJobScheduled(ev event.Event) {
	jobName := ev.Payload.(string)
	log.V(1).Infof("EventJobScheduled(%s): Dropping outstanding offers and bids", jobName)

	eh.agent.OfferResolved(jobName)

	if ev.Context.(machine.MachineState).BootId != eh.agent.Machine().State().BootId {
		log.V(1).Infof("EventJobScheduled(%s): Job not scheduled to this Agent, checking unbade offers", jobName)
		eh.agent.BidForPossibleJobs()
		return
	}

	log.V(1).Infof("EventJobScheduled(%s): Job scheduled to this Agent", jobName)

	j := eh.agent.FetchJob(jobName)
	if j == nil {
		log.Errorf("EventJobScheduled(%s): Failed to fetch Job", jobName)
		return
	}

	if ok := eh.agent.VerifyJob(j); !ok {
		log.Errorf("EventJobScheduled(%s): Failed to verify job", j.Name)
		return
	}

	if !eh.agent.AbleToRun(j) {
		log.V(1).Infof("EventJobScheduled(%s): Unable to run scheduled Job, rescheduling.", jobName)
		eh.agent.RescheduleJob(j)
		return
	}

	log.V(1).Infof("EventJobScheduled(%s): Starting Job", j.Name)
	eh.agent.StartJob(j)

	log.V(1).Infof("EventJobScheduled(%s): Bidding for all possible peers of Job", j.Name)
	eh.agent.BidForPossiblePeers(j.Name)
}

func (eh *EventHandler) HandleEventJobStopped(ev event.Event) {
	//TODO(bcwaldon): We should check the context of the event before
	// making any changes to local systemd or the registry

	jobName := ev.Payload.(string)
	log.Infof("EventJobStopped(%s): stopping corresponding unit", jobName)
	eh.agent.StopJob(jobName)

	log.V(1).Infof("EventJobStopped(%s): revisiting unresolved JobOffers", jobName)
	eh.agent.BidForPossibleJobs()
}

func (eh *EventHandler) HandleEventJobUpdated(ev event.Event) {
	j := ev.Payload.(job.Job)

	localBootId := eh.agent.Machine().State().BootId
	targetBootId := ev.Context.(string)

	if targetBootId != localBootId {
		log.V(1).Infof("EventJobUpdated(%s): Job not scheduled to Agent %s, skipping", j.Name, localBootId)
		return
	}

	log.V(1).Infof("EventJobUpdated(%s): Starting Job", j.Name)
	eh.agent.StartJob(&j)
}

func (eh *EventHandler) HandleEventJobStateUpdated(ev event.Event) {
	jobName := ev.Context.(string)
	state := ev.Payload.(*job.JobState)

	if state == nil {
		log.V(1).Infof("EventJobStateUpdated(%s): received nil JobState object", jobName)
	} else {
		log.V(1).Infof("EventJobStateUpdated(%s): pushing state (loadState=%s, activeState=%s, subState=%s) to Registry", jobName, state.LoadState, state.ActiveState, state.SubState)

		// FIXME: This should probably be set in the underlying event-generation code
		ms := eh.agent.Machine().State()
		state.MachineState = &ms
	}

	eh.agent.ReportJobState(jobName, state)
}

func (eh *EventHandler) HandleEventMachineCreated(ev event.Event) {
	mach := ev.Payload.(machine.MachineState)
	if mach.BootId != eh.agent.Machine().State().BootId {
		log.V(1).Infof("EventMachineCreated(%s): references other Machine, discarding event", mach.BootId)
	}

	for _, jo := range eh.agent.UnresolvedJobOffers() {
		log.V(1).Infof("EventMachineCreated(%s): verifying ability to run Job(%s)", mach.BootId, jo.Job.Name)

		// Everything we check against could change over time, so we track all
		// offers starting here for future bidding even if we can't bid now
		eh.agent.TrackOffer(jo)

		if eh.agent.AbleToRun(&jo.Job) {
			log.Infof("EventMachineCreated(%s): passed all criteria, submitting JobBid(%s)", mach.BootId, jo.Job.Name)
			eh.agent.Bid(jo.Job.Name)
		} else {
			log.V(1).Infof("EventMachineCreated(%s): not all criteria met, not bidding for Job(%s)", mach.BootId, jo.Job.Name)
		}
	}
}
<commit_msg>fix(fleet/VerifyJob): check job from EventJobUpdated<commit_after>package agent

import (
	log "github.com/coreos/fleet/third_party/github.com/golang/glog"

	"github.com/coreos/fleet/event"
	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/machine"
)

type EventHandler struct {
	agent *Agent
}

func NewEventHandler(agent *Agent) *EventHandler {
	return &EventHandler{agent}
}

func (eh *EventHandler) HandleEventJobOffered(ev event.Event) {
	jo := ev.Payload.(job.JobOffer)
	log.V(1).Infof("EventJobOffered(%s): verifying ability to run Job", jo.Job.Name)

	// Everything we check against could change over time, so we track all
	// offers starting here for future bidding even if we can't bid now
	eh.agent.TrackOffer(jo)

	if eh.agent.AbleToRun(&jo.Job) {
		log.Infof("EventJobOffered(%s): passed all criteria, submitting JobBid", jo.Job.Name)
		eh.agent.Bid(jo.Job.Name)
	} else {
		log.V(1).Infof("EventJobOffered(%s): not all criteria met, not bidding", jo.Job.Name)
	}
}

func (eh *EventHandler) HandleEventJobScheduled(ev event.Event) {
	jobName := ev.Payload.(string)
	log.V(1).Infof("EventJobScheduled(%s): Dropping outstanding offers and bids", jobName)

	eh.agent.OfferResolved(jobName)

	if ev.Context.(machine.MachineState).BootId != eh.agent.Machine().State().BootId {
		log.V(1).Infof("EventJobScheduled(%s): Job not scheduled to this Agent, checking unbade offers", jobName)
		eh.agent.BidForPossibleJobs()
		return
	}

	log.V(1).Infof("EventJobScheduled(%s): Job scheduled to this Agent", jobName)

	j := eh.agent.FetchJob(jobName)
	if j == nil {
		log.Errorf("EventJobScheduled(%s): Failed to fetch Job", jobName)
		return
	}

	if !eh.agent.VerifyJob(j) {
		log.Errorf("EventJobScheduled(%s): Failed to verify job", j.Name)
		return
	}

	if !eh.agent.AbleToRun(j) {
		log.V(1).Infof("EventJobScheduled(%s): Unable to run scheduled Job, rescheduling.", jobName)
		eh.agent.RescheduleJob(j)
		return
	}

	log.V(1).Infof("EventJobScheduled(%s): Starting Job", j.Name)
	eh.agent.StartJob(j)

	log.V(1).Infof("EventJobScheduled(%s): Bidding for all possible peers of Job", j.Name)
	eh.agent.BidForPossiblePeers(j.Name)
}

func (eh *EventHandler) HandleEventJobStopped(ev event.Event) {
	//TODO(bcwaldon): We should check the context of the event before
	// making any changes to local systemd or the registry

	jobName := ev.Payload.(string)
	log.Infof("EventJobStopped(%s): stopping corresponding unit", jobName)
	eh.agent.StopJob(jobName)

	log.V(1).Infof("EventJobStopped(%s): revisiting unresolved JobOffers", jobName)
	eh.agent.BidForPossibleJobs()
}

func (eh *EventHandler) HandleEventJobUpdated(ev event.Event) {
	j := ev.Payload.(job.Job)

	localBootId := eh.agent.Machine().State().BootId
	targetBootId := ev.Context.(string)

	if targetBootId != localBootId {
		log.V(1).Infof("EventJobUpdated(%s): Job not scheduled to Agent %s, skipping", j.Name, localBootId)
		return
	}

	if ok := eh.agent.VerifyJob(&j); !ok {
		log.Errorf("EventJobUpdated(%s): Failed to verify job", j.Name)
		return
	}

	log.V(1).Infof("EventJobUpdated(%s): Starting Job", j.Name)
	eh.agent.StartJob(&j)
}

func (eh *EventHandler) HandleEventJobStateUpdated(ev event.Event) {
	jobName := ev.Context.(string)
	state := ev.Payload.(*job.JobState)

	if state == nil {
		log.V(1).Infof("EventJobStateUpdated(%s): received nil JobState object", jobName)
	} else {
		log.V(1).Infof("EventJobStateUpdated(%s): pushing state (loadState=%s, activeState=%s, subState=%s) to Registry", jobName, state.LoadState, state.ActiveState, state.SubState)

		// FIXME: This should probably be set in the underlying event-generation code
		ms := eh.agent.Machine().State()
		state.MachineState = &ms
	}

	eh.agent.ReportJobState(jobName, state)
}

func (eh *EventHandler) HandleEventMachineCreated(ev event.Event) {
	mach := ev.Payload.(machine.MachineState)
	if mach.BootId != eh.agent.Machine().State().BootId {
		log.V(1).Infof("EventMachineCreated(%s): references other Machine, discarding event", mach.BootId)
	}

	for _, jo := range eh.agent.UnresolvedJobOffers() {
		log.V(1).Infof("EventMachineCreated(%s): verifying ability to run Job(%s)", mach.BootId, jo.Job.Name)

		// Everything we check against could change over time, so we track all
		// offers starting here for future bidding even if we can't bid now
		eh.agent.TrackOffer(jo)

		if eh.agent.AbleToRun(&jo.Job) {
			log.Infof("EventMachineCreated(%s): passed all criteria, submitting JobBid(%s)", mach.BootId, jo.Job.Name)
			eh.agent.Bid(jo.Job.Name)
		} else {
			log.V(1).Infof("EventMachineCreated(%s): not all criteria met, not bidding for Job(%s)", mach.BootId, jo.Job.Name)
		}
	}
}
<|endoftext|>


<commit_before>/*
 A Naive Bayesian Classifier
 Jake Brukhman <jbrukh@gmail.com>

 BAYESIAN CLASSIFICATION REFRESHER: suppose you have a set
 of classes (e.g. categories) C := {C_1, ..., C_n}, and a
 document D consisting of words D := {W_1, ..., W_k}.
 We wish to ascertain the probability that the document
 belongs to some class C_j given some set of training data
 associating documents and classes.

 By Bayes' Theorem, we have that

    P(C_j|D) = P(D|C_j)*P(C_j)/P(D).

 The LHS is the probability that the document belongs to class
 C_j given the document itself (by which is meant, in practice,
 the word frequencies occurring in this document), and our program
 will calculate this probability for each j and spit out the
 most likely class for this document.

 P(C_j) is referred to as the "prior" probability, or the
 probability that a document belongs to C_j in general, without
 seeing the document first. P(D|C_j) is the probability of seeing
 such a document, given that it belongs to C_j. Here, by assuming
 that words appear independently in documents (this being the
 "naive" assumption), we can estimate

    P(D|C_j) ~= P(W_1|C_j)*...*P(W_k|C_j)

 where P(W_i|C_j) is the probability of seeing the given word
 in a document of the given class. Finally, P(D) can be seen as
 merely a scaling factor and is not strictly relevant to
 classificiation, unless you want to normalize the resulting
 scores and actually see probabilities. In this case, note that

    P(D) = SUM_j(P(D|C_j)*P(C_j))

 One practical issue with performing these calculations is the
 possibility of float64 underflow when calculating P(D|C_j), as
 individual word probabilities can be arbitrarily small, and
 a document can have an arbitrarily large number of them. A
 typical method for dealing with this case is to transform the
 probability to the log domain and perform additions instead
 of multiplications:

   log P(C_j|D) ~ log(P(C_j)) + SUM_i(log P(W_i|C_j))

 where i = 1, ..., k. Note that by doing this, we are discarding
 the scaling factor P(D) and our scores are no longer
 probabilities; however, the monotonic relationship of the
 scores is preserved by the log function.
*/
package bayesian

import (
	"encoding/gob"
	"io"
	"math"
	"os"
	"errors"
)

// defaultProb is the tiny non-zero probability that a word
// we have not seen before appears in the class.
const defaultProb = 0.00000000001

// Class defines a set of classes that the classifier will
// filter: C = {C_1, ..., C_n}. You should define your classes
// as a set of constants, for example as follows:
//
//    const (
//        Good Class = "Good"
//        Bad Class = "Bad
//    )
//
// Class values should be unique.
type Class string

// Classifier implements the Naive Bayesian Classifier.
type Classifier struct {
	Classes []Class
	learned int // docs learned
	seen    int // docs seen
	datas   map[Class]*classData
}

// serializableClassifier represents a container for
// Classifier objects whose fields are modifiable by
// reflection and are therefore writeable by gob.
type serializableClassifier struct {
	Classes []Class
	Learned int
	Seen    int
	Datas   map[Class]*classData
}

// classData holds the frequency data for words in a
// particular class. In the future, we may replace this
// structure with a trie-like structure for more
// efficient storage.
type classData struct {
	Freqs map[string]int
	Total int
}

// newClassData creates a new empty classData node.
func newClassData() *classData {
	return &classData{
		Freqs: make(map[string]int),
	}
}

// getWordProb returns P(W|C_j) -- the probability of seeing
// a particular word W in a document of this class.
func (d *classData) getWordProb(word string) float64 {
	value, ok := d.Freqs[word]
	if !ok {
		return defaultProb
	}
	return float64(value) / float64(d.Total)
}

// getWordsProb returns P(D|C_j) -- the probability of seeing
// this set of words in a document of this class.
//
// Note that words should not be empty, and this method of
// calulation is prone to underflow if there are many words
// and their individual probabilties are small.
func (d *classData) getWordsProb(words []string) (prob float64) {
	prob = 1
	for _, word := range words {
		prob *= d.getWordProb(word)
	}
	return
}

// NewClassifier returns a new classifier. The classes the provided
// should be at least 2 in number and unique, or this method will
// panic.
func NewClassifier(classes ...Class) (c *Classifier) {
	n := len(classes)

	// check size
	if n < 2 {
		panic("provide at least two classes")
	}

	// check uniqueness
	check := make(map[Class]bool, n)
	for _, class := range classes {
		check[class] = true
	}
	if len(check) != n {
		panic("classes must be unique")
	}
	// create the classifier
	c = &Classifier{
		Classes: classes,
		datas:   make(map[Class]*classData, n),
	}
	for _, class := range classes {
		c.datas[class] = newClassData()
	}
	return
}

// NewClassifierFromFile loads an existing classifier from
// file. The classifier was previously saved with a call
// to c.WriteToFile(string).
func NewClassifierFromFile(name string) (c *Classifier, err error) {
	file, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	return NewClassifierFromReader(file)
}

//This actually does the deserializing of a Gob encoded classifier
func NewClassifierFromReader(r io.Reader) (c *Classifier, err error) {
	dec := gob.NewDecoder(r)
	w := new(serializableClassifier)
	err = dec.Decode(w)

	return &Classifier{w.Classes, w.Learned, w.Seen, w.Datas}, err
}

// getPriors returns the prior probabilities for the
// classes provided -- P(C_j).
//
// TODO: There is a way to smooth priors, currently
// not implemented here.
func (c *Classifier) getPriors() (priors []float64) {
	n := len(c.Classes)
	priors = make([]float64, n, n)
	sum := 0
	for index, class := range c.Classes {
		total := c.datas[class].Total
		priors[index] = float64(total)
		sum += total
	}
	if sum != 0 {
		for i := 0; i < n; i++ {
			priors[i] /= float64(sum)
		}
	}
	return
}

// Learned returns the number of documents ever learned
// in the lifetime of this classifier.
func (c *Classifier) Learned() int {
	return c.learned
}

// Seen returns the number of documents ever classified
// in the lifetime of this classifier.
func (c *Classifier) Seen() int {
	return c.seen
}

// WordCount returns the number of words counted for
// each class in the lifetime of the classifier.
func (c *Classifier) WordCount() (result []int) {
	result = make([]int, len(c.Classes))
	for inx, class := range c.Classes {
		data := c.datas[class]
		result[inx] = data.Total
	}
	return
}

// Learn will accept new training documents for
// supervised learning.
func (c *Classifier) Learn(document []string, which Class) {
	data := c.datas[which]
	for _, word := range document {
		data.Freqs[word]++
		data.Total++
	}
	c.learned++
}

// LogScores produces "log-likelihood"-like scores that can
// be used to classify documents into classes.
//
// The value of the score is proportional to the likelihood,
// as determined by the classifier, that the given document
// belongs to the given class. This is true even when scores
// returned are negative, which they will be (since we are
// taking logs of probabilities).
//
// The index j of the score corresponds to the class given
// by c.Classes[j].
//
// Additionally returned are "inx" and "strict" values. The
// inx corresponds to the maximum score in the array. If more
// than one of the scores holds the maximum values, then
// strict is false.
//
// Unlike c.Probabilities(), this function is not prone to
// floating point underflow and is relatively safe to use.
func (c *Classifier) LogScores(document []string) (scores []float64, inx int, strict bool) {
	n := len(c.Classes)
	scores = make([]float64, n, n)
	priors := c.getPriors()

	// calculate the score for each class
	for index, class := range c.Classes {
		data := c.datas[class]
		// c is the sum of the logarithms
		// as outlined in the refresher
		score := math.Log(priors[index])
		for _, word := range document {
			score += math.Log(data.getWordProb(word))
		}
		scores[index] = score
	}
	inx, strict = findMax(scores)
	c.seen++
	return scores, inx, strict
}

// ProbScores works the same as LogScores, but delivers
// actual probabilities as discussed above. Note that float64
// underflow is possible if the word list contains too
// many words that have probabilities very close to 0.
//
// Notes on underflow: underflow is going to occur when you're
// trying to assess large numbers of words that you have
// never seen before. Depending on the application, this
// may or may not be a concern. Consider using SafeProbScores()
// instead.
func (c *Classifier) ProbScores(doc []string) (scores []float64, inx int, strict bool) {
	n := len(c.Classes)
	scores = make([]float64, n, n)
	priors := c.getPriors()
	sum := float64(0)
	// calculate the score for each class
	for index, class := range c.Classes {
		data := c.datas[class]
		// c is the sum of the logarithms
		// as outlined in the refresher
		score := priors[index]
		for _, word := range doc {
			score *= data.getWordProb(word)
		}
		scores[index] = score
		sum += score
	}
	for i := 0; i < n; i++ {
		scores[i] /= sum
	}
	inx, strict = findMax(scores)
	c.seen++
	return scores, inx, strict
}

// SafeProbScores works the same as ProbScores, but is
// able to detect underflow in those cases where underflow
// results in the reverse classification. If an underflow is detected,
// this method panics, allowing the user to recover as
// necessary. Note that underflow, under certain rare circumstances,
// may still result in incorrect probabilities being returned,
// but this method guarantees that all panic-less invokations
// are properly classified.
//
// Underflow detection is more costly because it also
// has to make additional log score calculations.
func (c *Classifier) SafeProbScores(doc []string) (scores []float64, inx int, strict bool, err error) {
	n := len(c.Classes)
	scores = make([]float64, n, n)
	logScores := make([]float64, n, n)
	priors := c.getPriors()
	sum := float64(0)
	// calculate the score for each class
	for index, class := range c.Classes {
		data := c.datas[class]
		// c is the sum of the logarithms
		// as outlined in the refresher
		score := priors[index]
		logScore := math.Log(priors[index])
		for _, word := range doc {
			p := data.getWordProb(word)
			score *= p
			logScore += math.Log(p)
		}
		scores[index] = score
		logScores[index] = logScore
		sum += score
	}
	for i := 0; i < n; i++ {
		scores[i] /= sum
	}
	inx, strict = findMax(scores)
	logInx, logStrict := findMax(logScores)

	// detect underflow -- the size
	// relation between scores and logScores
	// must be preserved or something is wrong
	if inx != logInx || strict != logStrict {
		err = errors.New("possible underflow detected")
	}
	c.seen++
	return scores, inx, strict, err
}

// WordFrequencies returns a matrix of word frequencies that currently
// exist in the classifier for each class state for the given input
// words. In other words, if you obtain the frequencies
//
//    freqs := c.WordFrequencies(/* [j]string */)
//
// then the expression freq[i][j] represents the frequency of the j-th
// word within the i-th class.
func (c *Classifier) WordFrequencies(words []string) (freqMatrix [][]float64) {
	n, l := len(c.Classes), len(words)
	freqMatrix = make([][]float64, n)
	for i, _ := range freqMatrix {
		arr := make([]float64, l)
		data := c.datas[c.Classes[i]]
		for j, _ := range arr {
			arr[j] = data.getWordProb(words[j])
		}
		freqMatrix[i] = arr
	}
	return
}

// Serialize this classifier to a file.
func (c *Classifier) WriteToFile(name string) (err error) {
	file, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	return c.WriteTo(file)
}

//Serialize this classifier to GOB and write to Writer
func (c *Classifier) WriteTo(w io.Writer) (err error) {
	enc := gob.NewEncoder(w)
	err = enc.Encode(&serializableClassifier{c.Classes, c.learned, c.seen, c.datas})
	return
}

// findMax finds the maximum of a set of scores; if the
// maximum is strict -- that is, it is the single unique
// maximum from the set -- then strict has return value
// true. Otherwise it is false.
func findMax(scores []float64) (inx int, strict bool) {
	inx = 0
	strict = true
	for i := 1; i < len(scores); i++ {
		if scores[inx] < scores[i] {
			inx = i
			strict = true
		} else if scores[inx] == scores[i] {
			strict = false
		}
	}
	return
}
<commit_msg>Added field tags so Classifier Classes can be returned as JSON<commit_after>/*
 A Naive Bayesian Classifier
 Jake Brukhman <jbrukh@gmail.com>

 BAYESIAN CLASSIFICATION REFRESHER: suppose you have a set
 of classes (e.g. categories) C := {C_1, ..., C_n}, and a
 document D consisting of words D := {W_1, ..., W_k}.
 We wish to ascertain the probability that the document
 belongs to some class C_j given some set of training data
 associating documents and classes.

 By Bayes' Theorem, we have that

    P(C_j|D) = P(D|C_j)*P(C_j)/P(D).

 The LHS is the probability that the document belongs to class
 C_j given the document itself (by which is meant, in practice,
 the word frequencies occurring in this document), and our program
 will calculate this probability for each j and spit out the
 most likely class for this document.

 P(C_j) is referred to as the "prior" probability, or the
 probability that a document belongs to C_j in general, without
 seeing the document first. P(D|C_j) is the probability of seeing
 such a document, given that it belongs to C_j. Here, by assuming
 that words appear independently in documents (this being the
 "naive" assumption), we can estimate

    P(D|C_j) ~= P(W_1|C_j)*...*P(W_k|C_j)

 where P(W_i|C_j) is the probability of seeing the given word
 in a document of the given class. Finally, P(D) can be seen as
 merely a scaling factor and is not strictly relevant to
 classificiation, unless you want to normalize the resulting
 scores and actually see probabilities. In this case, note that

    P(D) = SUM_j(P(D|C_j)*P(C_j))

 One practical issue with performing these calculations is the
 possibility of float64 underflow when calculating P(D|C_j), as
 individual word probabilities can be arbitrarily small, and
 a document can have an arbitrarily large number of them. A
 typical method for dealing with this case is to transform the
 probability to the log domain and perform additions instead
 of multiplications:

   log P(C_j|D) ~ log(P(C_j)) + SUM_i(log P(W_i|C_j))

 where i = 1, ..., k. Note that by doing this, we are discarding
 the scaling factor P(D) and our scores are no longer
 probabilities; however, the monotonic relationship of the
 scores is preserved by the log function.
*/
package bayesian

import (
	"encoding/gob"
	"errors"
	"io"
	"math"
	"os"
)

// defaultProb is the tiny non-zero probability that a word
// we have not seen before appears in the class.
const defaultProb = 0.00000000001

// Class defines a set of classes that the classifier will
// filter: C = {C_1, ..., C_n}. You should define your classes
// as a set of constants, for example as follows:
//
//    const (
//        Good Class = "Good"
//        Bad Class = "Bad
//    )
//
// Class values should be unique.
type Class string

// Classifier implements the Naive Bayesian Classifier.
type Classifier struct {
	Classes []Class `json:"classes"`
	learned int     // docs learned
	seen    int     // docs seen
	datas   map[Class]*classData
}

// serializableClassifier represents a container for
// Classifier objects whose fields are modifiable by
// reflection and are therefore writeable by gob.
type serializableClassifier struct {
	Classes []Class
	Learned int
	Seen    int
	Datas   map[Class]*classData
}

// classData holds the frequency data for words in a
// particular class. In the future, we may replace this
// structure with a trie-like structure for more
// efficient storage.
type classData struct {
	Freqs map[string]int
	Total int
}

// newClassData creates a new empty classData node.
func newClassData() *classData {
	return &classData{
		Freqs: make(map[string]int),
	}
}

// getWordProb returns P(W|C_j) -- the probability of seeing
// a particular word W in a document of this class.
func (d *classData) getWordProb(word string) float64 {
	value, ok := d.Freqs[word]
	if !ok {
		return defaultProb
	}
	return float64(value) / float64(d.Total)
}

// getWordsProb returns P(D|C_j) -- the probability of seeing
// this set of words in a document of this class.
//
// Note that words should not be empty, and this method of
// calulation is prone to underflow if there are many words
// and their individual probabilties are small.
func (d *classData) getWordsProb(words []string) (prob float64) {
	prob = 1
	for _, word := range words {
		prob *= d.getWordProb(word)
	}
	return
}

// NewClassifier returns a new classifier. The classes the provided
// should be at least 2 in number and unique, or this method will
// panic.
func NewClassifier(classes ...Class) (c *Classifier) {
	n := len(classes)

	// check size
	if n < 2 {
		panic("provide at least two classes")
	}

	// check uniqueness
	check := make(map[Class]bool, n)
	for _, class := range classes {
		check[class] = true
	}
	if len(check) != n {
		panic("classes must be unique")
	}
	// create the classifier
	c = &Classifier{
		Classes: classes,
		datas:   make(map[Class]*classData, n),
	}
	for _, class := range classes {
		c.datas[class] = newClassData()
	}
	return
}

// NewClassifierFromFile loads an existing classifier from
// file. The classifier was previously saved with a call
// to c.WriteToFile(string).
func NewClassifierFromFile(name string) (c *Classifier, err error) {
	file, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	return NewClassifierFromReader(file)
}

//This actually does the deserializing of a Gob encoded classifier
func NewClassifierFromReader(r io.Reader) (c *Classifier, err error) {
	dec := gob.NewDecoder(r)
	w := new(serializableClassifier)
	err = dec.Decode(w)

	return &Classifier{w.Classes, w.Learned, w.Seen, w.Datas}, err
}

// getPriors returns the prior probabilities for the
// classes provided -- P(C_j).
//
// TODO: There is a way to smooth priors, currently
// not implemented here.
func (c *Classifier) getPriors() (priors []float64) {
	n := len(c.Classes)
	priors = make([]float64, n, n)
	sum := 0
	for index, class := range c.Classes {
		total := c.datas[class].Total
		priors[index] = float64(total)
		sum += total
	}
	if sum != 0 {
		for i := 0; i < n; i++ {
			priors[i] /= float64(sum)
		}
	}
	return
}

// Learned returns the number of documents ever learned
// in the lifetime of this classifier.
func (c *Classifier) Learned() int {
	return c.learned
}

// Seen returns the number of documents ever classified
// in the lifetime of this classifier.
func (c *Classifier) Seen() int {
	return c.seen
}

// WordCount returns the number of words counted for
// each class in the lifetime of the classifier.
func (c *Classifier) WordCount() (result []int) {
	result = make([]int, len(c.Classes))
	for inx, class := range c.Classes {
		data := c.datas[class]
		result[inx] = data.Total
	}
	return
}

// Learn will accept new training documents for
// supervised learning.
func (c *Classifier) Learn(document []string, which Class) {
	data := c.datas[which]
	for _, word := range document {
		data.Freqs[word]++
		data.Total++
	}
	c.learned++
}

// LogScores produces "log-likelihood"-like scores that can
// be used to classify documents into classes.
//
// The value of the score is proportional to the likelihood,
// as determined by the classifier, that the given document
// belongs to the given class. This is true even when scores
// returned are negative, which they will be (since we are
// taking logs of probabilities).
//
// The index j of the score corresponds to the class given
// by c.Classes[j].
//
// Additionally returned are "inx" and "strict" values. The
// inx corresponds to the maximum score in the array. If more
// than one of the scores holds the maximum values, then
// strict is false.
//
// Unlike c.Probabilities(), this function is not prone to
// floating point underflow and is relatively safe to use.
func (c *Classifier) LogScores(document []string) (scores []float64, inx int, strict bool) {
	n := len(c.Classes)
	scores = make([]float64, n, n)
	priors := c.getPriors()

	// calculate the score for each class
	for index, class := range c.Classes {
		data := c.datas[class]
		// c is the sum of the logarithms
		// as outlined in the refresher
		score := math.Log(priors[index])
		for _, word := range document {
			score += math.Log(data.getWordProb(word))
		}
		scores[index] = score
	}
	inx, strict = findMax(scores)
	c.seen++
	return scores, inx, strict
}

// ProbScores works the same as LogScores, but delivers
// actual probabilities as discussed above. Note that float64
// underflow is possible if the word list contains too
// many words that have probabilities very close to 0.
//
// Notes on underflow: underflow is going to occur when you're
// trying to assess large numbers of words that you have
// never seen before. Depending on the application, this
// may or may not be a concern. Consider using SafeProbScores()
// instead.
func (c *Classifier) ProbScores(doc []string) (scores []float64, inx int, strict bool) {
	n := len(c.Classes)
	scores = make([]float64, n, n)
	priors := c.getPriors()
	sum := float64(0)
	// calculate the score for each class
	for index, class := range c.Classes {
		data := c.datas[class]
		// c is the sum of the logarithms
		// as outlined in the refresher
		score := priors[index]
		for _, word := range doc {
			score *= data.getWordProb(word)
		}
		scores[index] = score
		sum += score
	}
	for i := 0; i < n; i++ {
		scores[i] /= sum
	}
	inx, strict = findMax(scores)
	c.seen++
	return scores, inx, strict
}

// SafeProbScores works the same as ProbScores, but is
// able to detect underflow in those cases where underflow
// results in the reverse classification. If an underflow is detected,
// this method panics, allowing the user to recover as
// necessary. Note that underflow, under certain rare circumstances,
// may still result in incorrect probabilities being returned,
// but this method guarantees that all panic-less invokations
// are properly classified.
//
// Underflow detection is more costly because it also
// has to make additional log score calculations.
func (c *Classifier) SafeProbScores(doc []string) (scores []float64, inx int, strict bool, err error) {
	n := len(c.Classes)
	scores = make([]float64, n, n)
	logScores := make([]float64, n, n)
	priors := c.getPriors()
	sum := float64(0)
	// calculate the score for each class
	for index, class := range c.Classes {
		data := c.datas[class]
		// c is the sum of the logarithms
		// as outlined in the refresher
		score := priors[index]
		logScore := math.Log(priors[index])
		for _, word := range doc {
			p := data.getWordProb(word)
			score *= p
			logScore += math.Log(p)
		}
		scores[index] = score
		logScores[index] = logScore
		sum += score
	}
	for i := 0; i < n; i++ {
		scores[i] /= sum
	}
	inx, strict = findMax(scores)
	logInx, logStrict := findMax(logScores)

	// detect underflow -- the size
	// relation between scores and logScores
	// must be preserved or something is wrong
	if inx != logInx || strict != logStrict {
		err = errors.New("possible underflow detected")
	}
	c.seen++
	return scores, inx, strict, err
}

// WordFrequencies returns a matrix of word frequencies that currently
// exist in the classifier for each class state for the given input
// words. In other words, if you obtain the frequencies
//
//    freqs := c.WordFrequencies(/* [j]string */)
//
// then the expression freq[i][j] represents the frequency of the j-th
// word within the i-th class.
func (c *Classifier) WordFrequencies(words []string) (freqMatrix [][]float64) {
	n, l := len(c.Classes), len(words)
	freqMatrix = make([][]float64, n)
	for i, _ := range freqMatrix {
		arr := make([]float64, l)
		data := c.datas[c.Classes[i]]
		for j, _ := range arr {
			arr[j] = data.getWordProb(words[j])
		}
		freqMatrix[i] = arr
	}
	return
}

// Serialize this classifier to a file.
func (c *Classifier) WriteToFile(name string) (err error) {
	file, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	return c.WriteTo(file)
}

//Serialize this classifier to GOB and write to Writer
func (c *Classifier) WriteTo(w io.Writer) (err error) {
	enc := gob.NewEncoder(w)
	err = enc.Encode(&serializableClassifier{c.Classes, c.learned, c.seen, c.datas})
	return
}

// findMax finds the maximum of a set of scores; if the
// maximum is strict -- that is, it is the single unique
// maximum from the set -- then strict has return value
// true. Otherwise it is false.
func findMax(scores []float64) (inx int, strict bool) {
	inx = 0
	strict = true
	for i := 1; i < len(scores); i++ {
		if scores[inx] < scores[i] {
			inx = i
			strict = true
		} else if scores[inx] == scores[i] {
			strict = false
		}
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	log "github.com/Sirupsen/logrus"
	r "github.com/dancannon/gorethink"
	"github.com/stefanprodan/syros/models"
	"time"
)

type Repository struct {
	Config  *Config
	Session *r.Session
}

func NewRepository(config *Config) (*Repository, error) {

	session, err := r.Connect(r.ConnectOpts{
		Address:  config.RethinkDB,
		Database: config.Database,
	})
	if err != nil {
		return nil, err
	}

	repo := &Repository{
		Config:  config,
		Session: session,
	}
	return repo, nil
}

func (repo *Repository) Initialize() {
	var cursor *r.Cursor
	var err error
	var cnt int

	// init db
	cursor, err = r.DBList().Contains(repo.Config.Database).Run(repo.Session)
	if err != nil {
		log.Fatalf("RethinkDB database init query failed %v", err)
	}

	cursor.One(&cnt)
	cursor.Close()

	if cnt < 1 {
		log.Infof("RethinkDB no database found, creating %v", repo.Config.Database)
		_, err := r.DBCreate(repo.Config.Database).RunWrite(repo.Session)
		if err != nil {
			log.Fatalf("RethinkDB database creation failed %v", err)
		}
	}

	repo.CreateTable("hosts")
	repo.CreateIndex("hosts", "environment")
	repo.CreateIndex("hosts", "collected")
	repo.CreateTable("containers")
	repo.CreateIndex("containers", "host_id")
	repo.CreateIndex("containers", "environment")
	repo.CreateIndex("containers", "collected")
	repo.CreateTable("syros_services")
	repo.CreateIndex("syros_services", "environment")
	repo.CreateIndex("syros_services", "collected")
}

func (repo *Repository) CreateTable(table string) {
	rdb := r.DB(repo.Config.Database)
	cursor, err := rdb.TableList().Contains(table).Run(repo.Session)
	if err != nil {
		log.Fatalf("RethinkDB table init query failed %v", err)
	}
	var cnt int
	cursor.One(&cnt)
	cursor.Close()

	if cnt < 1 {
		log.Infof("RethinkDB no table found, creating %v", table)
		_, err := rdb.TableCreate(table).RunWrite(repo.Session)
		if err != nil {
			log.Fatalf("RethinkDB %v table creation failed %v", table, err)
		}
	}
}

func (repo *Repository) CreateIndex(table string, field string) {
	t := r.DB(repo.Config.Database).Table(table)
	cursor, err := t.IndexList().Contains(field).Run(repo.Session)
	if err != nil {
		log.Fatalf("RethinkDB index init query failed %v", err)
	}

	var cnt int
	cursor.One(&cnt)
	cursor.Close()

	if cnt < 1 {
		log.Infof("RethinkDB no index found on table %v, creating %v", table, field)
		err := t.IndexCreate(field).Exec(repo.Session)
		if err != nil {
			log.Fatalf("RethinkDB table %v index %v creation failed %v", table, field, err)
		}
		t.IndexWait().RunWrite(repo.Session)
		if err != nil {
			log.Fatalf("RethinkDB table %v index %v wait failed %v", table, field, err)
		}
	}
}

func (repo *Repository) HostUpsert(host models.DockerHost) {
	res, err := r.Table("hosts").Get(host.Id).Run(repo.Session)
	if err != nil {
		log.Errorf("Repository host upsert query after ID failed %v", err)
		return
	}

	if res.IsNil() {
		_, err := r.Table("hosts").Insert(host).RunWrite(repo.Session)
		if err != nil {
			log.Errorf("Repository host insert failed %v", err)
		}
	} else {
		_, err := r.Table("hosts").Get(host.Id).Update(host).Run(repo.Session)
		if err != nil {
			log.Errorf("Repository host update failed %v", err)
		}
	}
}

func (repo *Repository) ContainerUpsert(container models.DockerContainer) {
	res, err := r.Table("containers").Get(container.Id).Run(repo.Session)
	if err != nil {
		log.Errorf("Repository containers upsert query after ID failed %v", err)
		return
	}

	if res.IsNil() {
		_, err := r.Table("containers").Insert(container).RunWrite(repo.Session)
		if err != nil {
			log.Errorf("Repository containers insert failed %v", err)
		}
	} else {
		_, err := r.Table("containers").Get(container.Id).Update(container).Run(repo.Session)
		if err != nil {
			log.Errorf("Repository containers update failed %v", err)
		}
	}
}

func (repo *Repository) SyrosServiceUpsert(service models.SyrosService) {
	res, err := r.Table("syros_services").Get(service.Id).Run(repo.Session)
	if err != nil {
		log.Errorf("Repository syros_services upsert query after ID failed %v", err)
		return
	}

	if res.IsNil() {
		_, err := r.Table("syros_services").Insert(service).RunWrite(repo.Session)
		if err != nil {
			log.Errorf("Repository syros_services insert failed %v", err)
		}
	} else {
		_, err := r.Table("syros_services").Get(service.Id).Update(service).Run(repo.Session)
		if err != nil {
			log.Errorf("Repository syros_services update failed %v", err)
		}
	}
}

// Removes stale records that have not been updated for a while
func (repo *Repository) RunGarbageCollector(tables []string) {
	if repo.Config.DatabaseStale > 0 {
		go func() {
			for _, table := range tables {
				res, err := r.Table(table).
					Between(time.Now().Add(-12*time.Hour).UTC(),
						time.Now().Add(-time.Duration(repo.Config.DatabaseStale)*time.Minute).UTC(),
						r.BetweenOpts{Index: "collected"}).
					Delete().RunWrite(repo.Session)
				if err != nil {
					log.Errorf("Repository GC for table %v query failed %v", table, err)
				} else {
					if res.Deleted > 0 {
						log.Infof("Repository GC removed %v from %v", res.Deleted, table)
					}
				}
			}

			time.Sleep(60 * time.Second)
		}()
	}
}

func logRepositoryResponse(action string, response interface{}) {
	jBytes, _ := json.Marshal(response)
	log.Debugf("Repository %v result %s", action, string(jBytes))
}
<commit_msg>log repo GC<commit_after>package main

import (
	"encoding/json"
	log "github.com/Sirupsen/logrus"
	r "github.com/dancannon/gorethink"
	"github.com/stefanprodan/syros/models"
	"time"
)

type Repository struct {
	Config  *Config
	Session *r.Session
}

func NewRepository(config *Config) (*Repository, error) {

	session, err := r.Connect(r.ConnectOpts{
		Address:  config.RethinkDB,
		Database: config.Database,
	})
	if err != nil {
		return nil, err
	}

	repo := &Repository{
		Config:  config,
		Session: session,
	}
	return repo, nil
}

func (repo *Repository) Initialize() {
	var cursor *r.Cursor
	var err error
	var cnt int

	// init db
	cursor, err = r.DBList().Contains(repo.Config.Database).Run(repo.Session)
	if err != nil {
		log.Fatalf("RethinkDB database init query failed %v", err)
	}

	cursor.One(&cnt)
	cursor.Close()

	if cnt < 1 {
		log.Infof("RethinkDB no database found, creating %v", repo.Config.Database)
		_, err := r.DBCreate(repo.Config.Database).RunWrite(repo.Session)
		if err != nil {
			log.Fatalf("RethinkDB database creation failed %v", err)
		}
	}

	repo.CreateTable("hosts")
	repo.CreateIndex("hosts", "environment")
	repo.CreateIndex("hosts", "collected")
	repo.CreateTable("containers")
	repo.CreateIndex("containers", "host_id")
	repo.CreateIndex("containers", "environment")
	repo.CreateIndex("containers", "collected")
	repo.CreateTable("syros_services")
	repo.CreateIndex("syros_services", "environment")
	repo.CreateIndex("syros_services", "collected")
}

func (repo *Repository) CreateTable(table string) {
	rdb := r.DB(repo.Config.Database)
	cursor, err := rdb.TableList().Contains(table).Run(repo.Session)
	if err != nil {
		log.Fatalf("RethinkDB table init query failed %v", err)
	}
	var cnt int
	cursor.One(&cnt)
	cursor.Close()

	if cnt < 1 {
		log.Infof("RethinkDB no table found, creating %v", table)
		_, err := rdb.TableCreate(table).RunWrite(repo.Session)
		if err != nil {
			log.Fatalf("RethinkDB %v table creation failed %v", table, err)
		}
	}
}

func (repo *Repository) CreateIndex(table string, field string) {
	t := r.DB(repo.Config.Database).Table(table)
	cursor, err := t.IndexList().Contains(field).Run(repo.Session)
	if err != nil {
		log.Fatalf("RethinkDB index init query failed %v", err)
	}

	var cnt int
	cursor.One(&cnt)
	cursor.Close()

	if cnt < 1 {
		log.Infof("RethinkDB no index found on table %v, creating %v", table, field)
		err := t.IndexCreate(field).Exec(repo.Session)
		if err != nil {
			log.Fatalf("RethinkDB table %v index %v creation failed %v", table, field, err)
		}
		t.IndexWait().RunWrite(repo.Session)
		if err != nil {
			log.Fatalf("RethinkDB table %v index %v wait failed %v", table, field, err)
		}
	}
}

func (repo *Repository) HostUpsert(host models.DockerHost) {
	res, err := r.Table("hosts").Get(host.Id).Run(repo.Session)
	if err != nil {
		log.Errorf("Repository host upsert query after ID failed %v", err)
		return
	}

	if res.IsNil() {
		_, err := r.Table("hosts").Insert(host).RunWrite(repo.Session)
		if err != nil {
			log.Errorf("Repository host insert failed %v", err)
		}
	} else {
		_, err := r.Table("hosts").Get(host.Id).Update(host).Run(repo.Session)
		if err != nil {
			log.Errorf("Repository host update failed %v", err)
		}
	}
}

func (repo *Repository) ContainerUpsert(container models.DockerContainer) {
	res, err := r.Table("containers").Get(container.Id).Run(repo.Session)
	if err != nil {
		log.Errorf("Repository containers upsert query after ID failed %v", err)
		return
	}

	if res.IsNil() {
		_, err := r.Table("containers").Insert(container).RunWrite(repo.Session)
		if err != nil {
			log.Errorf("Repository containers insert failed %v", err)
		}
	} else {
		_, err := r.Table("containers").Get(container.Id).Update(container).Run(repo.Session)
		if err != nil {
			log.Errorf("Repository containers update failed %v", err)
		}
	}
}

func (repo *Repository) SyrosServiceUpsert(service models.SyrosService) {
	res, err := r.Table("syros_services").Get(service.Id).Run(repo.Session)
	if err != nil {
		log.Errorf("Repository syros_services upsert query after ID failed %v", err)
		return
	}

	if res.IsNil() {
		_, err := r.Table("syros_services").Insert(service).RunWrite(repo.Session)
		if err != nil {
			log.Errorf("Repository syros_services insert failed %v", err)
		}
	} else {
		_, err := r.Table("syros_services").Get(service.Id).Update(service).Run(repo.Session)
		if err != nil {
			log.Errorf("Repository syros_services update failed %v", err)
		}
	}
}

// Removes stale records that have not been updated for a while
func (repo *Repository) RunGarbageCollector(tables []string) {
	if repo.Config.DatabaseStale > 0 {
		go func(stale int) {
			log.Infof("Stating repository GC interval %v minutes", stale)
			for _, table := range tables {
				res, err := r.Table(table).
					Between(time.Now().Add(-12*time.Hour).UTC(),
						time.Now().Add(-time.Duration(stale)*time.Minute).UTC(),
						r.BetweenOpts{Index: "collected"}).
					Delete().RunWrite(repo.Session)
				if err != nil {
					log.Errorf("Repository GC for table %v query failed %v", table, err)
				} else {
					if res.Deleted > 0 {
						log.Infof("Repository GC removed %v from %v", res.Deleted, table)
					}
				}
			}

			time.Sleep(60 * time.Second)
		}(repo.Config.DatabaseStale)
	}
}

func logRepositoryResponse(action string, response interface{}) {
	jBytes, _ := json.Marshal(response)
	log.Debugf("Repository %v result %s", action, string(jBytes))
}
<|endoftext|>


<commit_before>package firego

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
)

// TransactionFn is used to run a transaction on a Firebase reference.
// See Firebase.Transaction for more information.
type TransactionFn func(currentSnapshot interface{}) (result interface{}, err error)

func getTransactionParams(headers http.Header, body []byte) (etag string, snapshot interface{}, err error) {
	etag = headers.Get("ETag")
	if len(etag) == 0 {
		return etag, snapshot, errors.New("no etag returned by Firebase")
	}

	if err := json.Unmarshal(body, &snapshot); err != nil {
		return etag, snapshot, fmt.Errorf("failed to unmarshal Firebase response. %s", err)
	}

	return etag, snapshot, nil
}

// Transaction runs a transaction on the data at this location. The TransactionFn parameter
// will be called, possibly multiple times, with the current data at this location.
// It is responsible for inspecting that data and specifying either the desired new data
// at the location or that the transaction should be aborted.
//
// Since the provided function may be called repeatedly for the same transaction, be extremely careful of
// any side effects that may be triggered by this method.
//
// Best practices for this method are to rely only on the data that is passed in.
func (fb *Firebase) Transaction(fn TransactionFn) error {
	// fetch etag and current value
	headers, body, err := fb.doRequest("GET", nil, withHeader("X-Firebase-ETag", "true"))
	if err != nil {
		return err
	}

	etag, snapshot, err := getTransactionParams(headers, body)
	if err != nil {
		return err
	}

	// set the error value to something non-nil so that
	// we step into the loop
	tErr := errors.New("")
	for i := 0; i < 25 && tErr != nil; i++ {
		// run transaction
		result, err := fn(snapshot)
		if err != nil {
			return nil
		}

		newBody, err := json.Marshal(result)
		if err != nil {
			return fmt.Errorf("failed to marshal transaction result. %s", err)
		}

		// attempt to update it
		headers, body, tErr = fb.doRequest("PUT", newBody, withHeader("if-match", etag))
		if tErr == nil {
			// we're good, break to loop
			break
		}

		// we failed to update, so grab the new snapshot/etag
		e, s, tErr := getTransactionParams(headers, body)
		if tErr != nil {
			return tErr
		}
		etag, snapshot = e, s
	}

	if tErr != nil {
		return fmt.Errorf("failed to run transaction. %s", tErr)
	}
	return nil
}
<commit_msg>fixing comment typo<commit_after>package firego

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
)

// TransactionFn is used to run a transaction on a Firebase reference.
// See Firebase.Transaction for more information.
type TransactionFn func(currentSnapshot interface{}) (result interface{}, err error)

func getTransactionParams(headers http.Header, body []byte) (etag string, snapshot interface{}, err error) {
	etag = headers.Get("ETag")
	if len(etag) == 0 {
		return etag, snapshot, errors.New("no etag returned by Firebase")
	}

	if err := json.Unmarshal(body, &snapshot); err != nil {
		return etag, snapshot, fmt.Errorf("failed to unmarshal Firebase response. %s", err)
	}

	return etag, snapshot, nil
}

// Transaction runs a transaction on the data at this location. The TransactionFn parameter
// will be called, possibly multiple times, with the current data at this location.
// It is responsible for inspecting that data and specifying either the desired new data
// at the location or that the transaction should be aborted.
//
// Since the provided function may be called repeatedly for the same transaction, be extremely careful of
// any side effects that may be triggered by this method.
//
// Best practices for this method are to rely only on the data that is passed in.
func (fb *Firebase) Transaction(fn TransactionFn) error {
	// fetch etag and current value
	headers, body, err := fb.doRequest("GET", nil, withHeader("X-Firebase-ETag", "true"))
	if err != nil {
		return err
	}

	etag, snapshot, err := getTransactionParams(headers, body)
	if err != nil {
		return err
	}

	// set the error value to something non-nil so that
	// we step into the loop
	tErr := errors.New("")
	for i := 0; i < 25 && tErr != nil; i++ {
		// run transaction
		result, err := fn(snapshot)
		if err != nil {
			return nil
		}

		newBody, err := json.Marshal(result)
		if err != nil {
			return fmt.Errorf("failed to marshal transaction result. %s", err)
		}

		// attempt to update it
		headers, body, tErr = fb.doRequest("PUT", newBody, withHeader("if-match", etag))
		if tErr == nil {
			// we're good, break the loop
			break
		}

		// we failed to update, so grab the new snapshot/etag
		e, s, tErr := getTransactionParams(headers, body)
		if tErr != nil {
			return tErr
		}
		etag, snapshot = e, s
	}

	if tErr != nil {
		return fmt.Errorf("failed to run transaction. %s", tErr)
	}
	return nil
}
<|endoftext|>


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

package alarm

import (
	"errors"
	"fmt"
	stdlog "log"
	"sync"
	"time"

	"github.com/robertkrimen/otto"
	"github.com/tsuru/tsuru-autoscale/action"
	"github.com/tsuru/tsuru-autoscale/datasource"
	"github.com/tsuru/tsuru-autoscale/db"
	"github.com/tsuru/tsuru-autoscale/log"
	"github.com/tsuru/tsuru-autoscale/tsuru"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

func StartAutoScale() {
	go runAutoScale()
}

func logger() *stdlog.Logger {
	return log.Logger()
}

// Alarm represents the configuration for the auto scale.
type Alarm struct {
	Name       string            `json:"name"`
	Actions    []string          `json:"actions"`
	Expression string            `json:"expression"`
	Enabled    bool              `json:"enabled"`
	Wait       time.Duration     `json:"wait"`
	DataSource string            `json:"datasource"`
	Instance   string            `json:"instance"`
	Envs       map[string]string `json:"envs"`
}

func NewAlarm(a *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		return nil
	}
	defer conn.Close()
	err = conn.Alarms().Insert(&a)
	if err != nil {
		return nil
	}
	return nil
}

func runAutoScaleOnce() {
	logger().Print("checking alarms")
	alarms := []Alarm{}
	conn, err := db.Conn()
	if err != nil {
		return
	}
	defer conn.Close()
	err = conn.Alarms().Find(bson.M{"enabled": true}).All(&alarms)
	if err != nil {
		return
	}
	var wg sync.WaitGroup
	for _, alarm := range alarms {
		wg.Add(1)
		go func(alarm Alarm) {
			defer wg.Done()
			logger().Printf("checking %s alarm", alarm.Name)
			err := scaleIfNeeded(&alarm)
			if err != nil {
				logger().Print(err.Error())
			}
		}(alarm)
	}
	wg.Wait()
}

func runAutoScale() {
	for {
		runAutoScaleOnce()
		time.Sleep(30 * time.Second)
	}
}

func scaleIfNeeded(alarm *Alarm) error {
	if alarm == nil {
		return errors.New("alarm: alarm is not configured.")
	}
	check, err := alarm.Check()
	if err != nil {
		logger().Printf("alarm %s check error: %s", alarm.Name, err.Error())
		return err
	}
	logger().Printf("alarm %s - %s - check: %t", alarm.Name, alarm.Expression, check)
	if check {
		if wait, err := shouldWait(alarm); err != nil {
			logger().Printf("waiting for alarm %s", alarm.Name)
			return err
		} else if wait {
			return nil
		}
		for _, alarmName := range alarm.Actions {
			a, err := action.FindByName(alarmName)
			if err != nil {
				logger().Printf("alarm %s not found - error: %s", alarmName, err)
			} else {
				logger().Printf("executing alarm %s action %s", alarm.Name, a.Name)
				instance, err := tsuru.GetInstanceByName(alarm.Instance)
				if err != nil {
					logger().Printf("Error trying to get instance by name, auto scale aborted: %s", err)
					return err
				}
				if len(instance.Apps) < 1 {
					msg := "Error trying to get app instance, auto scale aborted."
					logger().Printf(msg)
					err = errors.New(msg)
					return err
				}
				appName := instance.Apps[0]
				evt, err := NewEvent(alarm, a)
				if err != nil {
					logger().Printf("Error trying to insert auto scale event, auto scale aborted: %s", err)
				}
				aErr := a.Do(appName, alarm.Envs)
				if aErr != nil {
					logger().Printf("Error executing action %s in the alarm %s - error: %s", a.Name, alarm.Name, aErr)
				} else {
					logger().Printf("alarm %s action %s executed", alarm.Name, a.Name)
				}
				err = evt.update(aErr)
				if err != nil {
					logger().Printf("Error trying to update auto scale event: %s", err)
				}
			}
		}
		return nil
	}
	return nil
}

func shouldWait(alarm *Alarm) (bool, error) {
	now := time.Now().UTC()
	lastEvent, err := lastScaleEvent(alarm)
	if err != nil && err != mgo.ErrNotFound {
		logger().Printf("error on get last event for alarm %s - not waiting - err: %s", alarm.Name, err)
		return false, err
	}
	if err != mgo.ErrNotFound && lastEvent.EndTime.IsZero() {
		logger().Printf("last event not finished yet for alarm %s - waiting", alarm.Name)
		return true, nil
	}
	diff := now.Sub(lastEvent.EndTime)
	if diff > alarm.Wait {
		logger().Printf("diff %d > %d form alarm %s - not waiting", diff, alarm.Wait, alarm.Name)
		return false, nil
	}
	logger().Printf("diff %d < %d form alarm %s - waiting", diff, alarm.Wait, alarm.Name)
	return true, nil
}

func Enable(alarm *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Update(bson.M{"name": alarm.Name}, bson.M{"$set": bson.M{"enabled": true}})
}

func Disable(alarm *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Update(bson.M{"name": alarm.Name}, bson.M{"$set": bson.M{"enabled": false}})
}

func (a *Alarm) Check() (bool, error) {
	logger().Printf("getting data for alarm %s", a.Name)
	ds, err := datasource.Get(a.DataSource)
	if err != nil {
		logger().Printf("error getting data for alarm %s - error: %s", a.Name, err.Error())
		return false, err
	}
	instance, err := tsuru.GetInstanceByName(a.Instance)
	if err != nil {
		logger().Printf("Error trying to get instance by name, auto scale aborted: %s", err)
		return false, err
	}
	if len(instance.Apps) < 0 {
		msg := "Error trying to get app instance."
		logger().Printf(msg)
		err = errors.New(msg)
		return false, err
	}
	appName := instance.Apps[0]
	data, err := ds.Get(appName)
	if err != nil {
		logger().Printf("error getting data for alarm %s - error: %s", a.Name, err.Error())
		return false, err
	}
	logger().Printf("data for alarm %s - %s", a.Name, data)
	vm := otto.New()
	vm.Run(fmt.Sprintf("var data=%s;", data))
	vm.Run(fmt.Sprintf("var expression=%s", a.Expression))
	expression, err := vm.Get("expression")
	if err != nil {
		logger().Printf("error executing expresion for alarm %s - error: %s", a.Name, err.Error())
		return false, err
	}
	check, err := expression.ToBoolean()
	if err != nil {
		return false, err
	}
	return check, nil
}

// ListAlarmsByToken lists alarms by token.
func ListAlarmsByToken(token string) ([]Alarm, error) {
	i, err := tsuru.FindServiceInstance(token)
	if err != nil {
		logger().Printf("error find service instance by token %s - error: %s", token, err.Error())
		return nil, err
	}
	instances := []string{}
	for _, instance := range i {
		instances = append(instances, instance.Name)
	}
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var alarms []Alarm
	err = conn.Alarms().Find(bson.M{"instance": bson.M{"$in": instances}}).All(&alarms)
	if err != nil {
		logger().Printf("error find alarms by instance #%v", instances)
		return nil, err
	}
	return alarms, nil
}

// ListAlarmsByInstance lists alarms by instance.
func ListAlarmsByInstance(instanceName string) ([]Alarm, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var alarms []Alarm
	err = conn.Alarms().Find(bson.M{"instance": instanceName}).All(&alarms)
	if err != nil {
		logger().Printf("error find alarms by instance %q", instanceName)
		return nil, err
	}
	return alarms, nil
}

// FindAlarmByName find alarm by name.
func FindAlarmByName(name string) (*Alarm, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var alarm Alarm
	err = conn.Alarms().Find(bson.M{"name": name}).One(&alarm)
	if err != nil {
		return nil, err
	}
	return &alarm, nil
}

// RemoveAlarm removes an alarm.
func RemoveAlarm(a *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	return conn.Alarms().Remove(a)
}
<commit_msg>alarm/new: simplifying the new method<commit_after>// Copyright 2015 tsuru-autoscale authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package alarm

import (
	"errors"
	"fmt"
	stdlog "log"
	"sync"
	"time"

	"github.com/robertkrimen/otto"
	"github.com/tsuru/tsuru-autoscale/action"
	"github.com/tsuru/tsuru-autoscale/datasource"
	"github.com/tsuru/tsuru-autoscale/db"
	"github.com/tsuru/tsuru-autoscale/log"
	"github.com/tsuru/tsuru-autoscale/tsuru"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

func StartAutoScale() {
	go runAutoScale()
}

func logger() *stdlog.Logger {
	return log.Logger()
}

// Alarm represents the configuration for the auto scale.
type Alarm struct {
	Name       string            `json:"name"`
	Actions    []string          `json:"actions"`
	Expression string            `json:"expression"`
	Enabled    bool              `json:"enabled"`
	Wait       time.Duration     `json:"wait"`
	DataSource string            `json:"datasource"`
	Instance   string            `json:"instance"`
	Envs       map[string]string `json:"envs"`
}

func NewAlarm(a *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Insert(&a)
}

func runAutoScaleOnce() {
	logger().Print("checking alarms")
	alarms := []Alarm{}
	conn, err := db.Conn()
	if err != nil {
		return
	}
	defer conn.Close()
	err = conn.Alarms().Find(bson.M{"enabled": true}).All(&alarms)
	if err != nil {
		return
	}
	var wg sync.WaitGroup
	for _, alarm := range alarms {
		wg.Add(1)
		go func(alarm Alarm) {
			defer wg.Done()
			logger().Printf("checking %s alarm", alarm.Name)
			err := scaleIfNeeded(&alarm)
			if err != nil {
				logger().Print(err.Error())
			}
		}(alarm)
	}
	wg.Wait()
}

func runAutoScale() {
	for {
		runAutoScaleOnce()
		time.Sleep(30 * time.Second)
	}
}

func scaleIfNeeded(alarm *Alarm) error {
	if alarm == nil {
		return errors.New("alarm: alarm is not configured.")
	}
	check, err := alarm.Check()
	if err != nil {
		logger().Printf("alarm %s check error: %s", alarm.Name, err.Error())
		return err
	}
	logger().Printf("alarm %s - %s - check: %t", alarm.Name, alarm.Expression, check)
	if check {
		if wait, err := shouldWait(alarm); err != nil {
			logger().Printf("waiting for alarm %s", alarm.Name)
			return err
		} else if wait {
			return nil
		}
		for _, alarmName := range alarm.Actions {
			a, err := action.FindByName(alarmName)
			if err != nil {
				logger().Printf("alarm %s not found - error: %s", alarmName, err)
			} else {
				logger().Printf("executing alarm %s action %s", alarm.Name, a.Name)
				instance, err := tsuru.GetInstanceByName(alarm.Instance)
				if err != nil {
					logger().Printf("Error trying to get instance by name, auto scale aborted: %s", err)
					return err
				}
				if len(instance.Apps) < 1 {
					msg := "Error trying to get app instance, auto scale aborted."
					logger().Printf(msg)
					err = errors.New(msg)
					return err
				}
				appName := instance.Apps[0]
				evt, err := NewEvent(alarm, a)
				if err != nil {
					logger().Printf("Error trying to insert auto scale event, auto scale aborted: %s", err)
				}
				aErr := a.Do(appName, alarm.Envs)
				if aErr != nil {
					logger().Printf("Error executing action %s in the alarm %s - error: %s", a.Name, alarm.Name, aErr)
				} else {
					logger().Printf("alarm %s action %s executed", alarm.Name, a.Name)
				}
				err = evt.update(aErr)
				if err != nil {
					logger().Printf("Error trying to update auto scale event: %s", err)
				}
			}
		}
		return nil
	}
	return nil
}

func shouldWait(alarm *Alarm) (bool, error) {
	now := time.Now().UTC()
	lastEvent, err := lastScaleEvent(alarm)
	if err != nil && err != mgo.ErrNotFound {
		logger().Printf("error on get last event for alarm %s - not waiting - err: %s", alarm.Name, err)
		return false, err
	}
	if err != mgo.ErrNotFound && lastEvent.EndTime.IsZero() {
		logger().Printf("last event not finished yet for alarm %s - waiting", alarm.Name)
		return true, nil
	}
	diff := now.Sub(lastEvent.EndTime)
	if diff > alarm.Wait {
		logger().Printf("diff %d > %d form alarm %s - not waiting", diff, alarm.Wait, alarm.Name)
		return false, nil
	}
	logger().Printf("diff %d < %d form alarm %s - waiting", diff, alarm.Wait, alarm.Name)
	return true, nil
}

func Enable(alarm *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Update(bson.M{"name": alarm.Name}, bson.M{"$set": bson.M{"enabled": true}})
}

func Disable(alarm *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		return nil
	}
	defer conn.Close()
	return conn.Alarms().Update(bson.M{"name": alarm.Name}, bson.M{"$set": bson.M{"enabled": false}})
}

func (a *Alarm) Check() (bool, error) {
	logger().Printf("getting data for alarm %s", a.Name)
	ds, err := datasource.Get(a.DataSource)
	if err != nil {
		logger().Printf("error getting data for alarm %s - error: %s", a.Name, err.Error())
		return false, err
	}
	instance, err := tsuru.GetInstanceByName(a.Instance)
	if err != nil {
		logger().Printf("Error trying to get instance by name, auto scale aborted: %s", err)
		return false, err
	}
	if len(instance.Apps) < 0 {
		msg := "Error trying to get app instance."
		logger().Printf(msg)
		err = errors.New(msg)
		return false, err
	}
	appName := instance.Apps[0]
	data, err := ds.Get(appName)
	if err != nil {
		logger().Printf("error getting data for alarm %s - error: %s", a.Name, err.Error())
		return false, err
	}
	logger().Printf("data for alarm %s - %s", a.Name, data)
	vm := otto.New()
	vm.Run(fmt.Sprintf("var data=%s;", data))
	vm.Run(fmt.Sprintf("var expression=%s", a.Expression))
	expression, err := vm.Get("expression")
	if err != nil {
		logger().Printf("error executing expresion for alarm %s - error: %s", a.Name, err.Error())
		return false, err
	}
	check, err := expression.ToBoolean()
	if err != nil {
		return false, err
	}
	return check, nil
}

// ListAlarmsByToken lists alarms by token.
func ListAlarmsByToken(token string) ([]Alarm, error) {
	i, err := tsuru.FindServiceInstance(token)
	if err != nil {
		logger().Printf("error find service instance by token %s - error: %s", token, err.Error())
		return nil, err
	}
	instances := []string{}
	for _, instance := range i {
		instances = append(instances, instance.Name)
	}
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var alarms []Alarm
	err = conn.Alarms().Find(bson.M{"instance": bson.M{"$in": instances}}).All(&alarms)
	if err != nil {
		logger().Printf("error find alarms by instance #%v", instances)
		return nil, err
	}
	return alarms, nil
}

// ListAlarmsByInstance lists alarms by instance.
func ListAlarmsByInstance(instanceName string) ([]Alarm, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var alarms []Alarm
	err = conn.Alarms().Find(bson.M{"instance": instanceName}).All(&alarms)
	if err != nil {
		logger().Printf("error find alarms by instance %q", instanceName)
		return nil, err
	}
	return alarms, nil
}

// FindAlarmByName find alarm by name.
func FindAlarmByName(name string) (*Alarm, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var alarm Alarm
	err = conn.Alarms().Find(bson.M{"name": name}).One(&alarm)
	if err != nil {
		return nil, err
	}
	return &alarm, nil
}

// RemoveAlarm removes an alarm.
func RemoveAlarm(a *Alarm) error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	return conn.Alarms().Remove(a)
}
<|endoftext|>


<commit_before>package album

import (
	gorp "gopkg.in/gorp.v1"
)

type Store struct {
	Database *gorp.DbMap
}

func NewStore(database *gorp.DbMap) (Store, error) {
	store := Store{
		Database: database,
	}

	return store, nil
}

func (store Store) Add(album *Album) error {
	err := store.Database.Insert(album)

	return err
}

func (store Store) GetById(id int64) (Album, error) {
	var album Album

	err := store.Database.SelectOne(&album, "select * from albums where id=$1", id)

	return album, err
}

func (store Store) GetAll() ([]Album, error) {
	var albums []Album

	_, err := store.Database.Select(&albums, "select * from albums order by id")

	return albums, err
}

func (store Store) Delete(album *Album) (int64, error) {
	count, err := store.Database.Delete(album)

	return count, err
}

func (store Store) Update(album *Album) (int64, error) {
	count, err := store.Database.Update(album)

	return count, err
}

func (store Store) Size() (int64, error) {
	size, err := store.Database.SelectInt("select count(*) from albums")

	return size, err
}
<commit_msg>Implemented a function on the album stores to retrieve an album's photos<commit_after>package album

import (
	"github.com/karousel/core/photo"
	gorp "gopkg.in/gorp.v1"
)

type Store struct {
	Database *gorp.DbMap
}

func NewStore(database *gorp.DbMap) (Store, error) {
	store := Store{
		Database: database,
	}

	return store, nil
}

func (store Store) Add(album *Album) error {
	err := store.Database.Insert(album)

	return err
}

func (store Store) GetById(id int64) (Album, error) {
	var album Album

	err := store.Database.SelectOne(&album, "select * from albums where id=$1", id)

	return album, err
}

func (store Store) GetAll() ([]Album, error) {
	var albums []Album

	_, err := store.Database.Select(&albums, "select * from albums order by id")

	return albums, err
}

func (store Store) Delete(album *Album) (int64, error) {
	count, err := store.Database.Delete(album)

	return count, err
}

func (store Store) Update(album *Album) (int64, error) {
	count, err := store.Database.Update(album)

	return count, err
}

func (store Store) Size() (int64, error) {
	size, err := store.Database.SelectInt("select count(*) from albums")

	return size, err
}

func (store Store) GetPhotosForAlbum(album *Album) ([]photo.Photo, error) {
	var photos []photo.Photo

	_, err := store.Database.Select(&photos, "select * from photos where album_id=$1", album.Id)

	return photos, err
}
<|endoftext|>


<commit_before>package repository

import (
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/user"
	"time"

	"restic/backend"
	"restic/crypto"
	"restic/debug"
)

var (
	// ErrNoKeyFound is returned when no key for the repository could be decrypted.
	ErrNoKeyFound = errors.New("no key could be found")
)

// TODO: figure out scrypt values on the fly depending on the current
// hardware.
const (
	scryptN        = 65536
	scryptR        = 8
	scryptP        = 1
	scryptSaltsize = 64
)

// Key represents an encrypted master key for a repository.
type Key struct {
	Created  time.Time `json:"created"`
	Username string    `json:"username"`
	Hostname string    `json:"hostname"`

	KDF  string `json:"kdf"`
	N    int    `json:"N"`
	R    int    `json:"r"`
	P    int    `json:"p"`
	Salt []byte `json:"salt"`
	Data []byte `json:"data"`

	user   *crypto.Key
	master *crypto.Key

	name string
}

// createMasterKey creates a new master key in the given backend and encrypts
// it with the password.
func createMasterKey(s *Repository, password string) (*Key, error) {
	return AddKey(s, password, nil)
}

// OpenKey tries do decrypt the key specified by name with the given password.
func OpenKey(s *Repository, name string, password string) (*Key, error) {
	k, err := LoadKey(s, name)
	if err != nil {
		debug.Log("OpenKey", "LoadKey(%v) returned error %v", name[:12], err)
		return nil, err
	}

	// check KDF
	if k.KDF != "scrypt" {
		return nil, errors.New("only supported KDF is scrypt()")
	}

	// derive user key
	k.user, err = crypto.KDF(k.N, k.R, k.P, k.Salt, password)
	if err != nil {
		return nil, err
	}

	// decrypt master keys
	buf, err := crypto.Decrypt(k.user, []byte{}, k.Data)
	if err != nil {
		return nil, err
	}

	// restore json
	k.master = &crypto.Key{}
	err = json.Unmarshal(buf, k.master)
	if err != nil {
		debug.Log("OpenKey", "Unmarshal() returned error %v", err)
		return nil, err
	}
	k.name = name

	if !k.Valid() {
		return nil, errors.New("Invalid key for repository")
	}

	return k, nil
}

// SearchKey tries to decrypt all keys in the backend with the given password.
// If none could be found, ErrNoKeyFound is returned.
func SearchKey(s *Repository, password string) (*Key, error) {
	// try all keys in repo
	done := make(chan struct{})
	defer close(done)
	for name := range s.Backend().List(backend.Key, done) {
		debug.Log("SearchKey", "trying key %v", name[:12])
		key, err := OpenKey(s, name, password)
		if err != nil {
			debug.Log("SearchKey", "key %v returned error %v", name[:12], err)
			continue
		}

		debug.Log("SearchKey", "successfully opened key %v", name[:12])
		return key, nil
	}

	return nil, ErrNoKeyFound
}

// LoadKey loads a key from the backend.
func LoadKey(s *Repository, name string) (k *Key, err error) {
	h := backend.Handle{Type: backend.Key, Name: name}
	data, err := backend.LoadAll(s.be, h, nil)
	if err != nil {
		return nil, err
	}

	k = &Key{}
	err = json.Unmarshal(data, k)
	if err != nil {
		return nil, err
	}

	return k, nil
}

// AddKey adds a new key to an already existing repository.
func AddKey(s *Repository, password string, template *crypto.Key) (*Key, error) {
	// fill meta data about key
	newkey := &Key{
		Created: time.Now(),
		KDF:     "scrypt",
		N:       scryptN,
		R:       scryptR,
		P:       scryptP,
	}

	hn, err := os.Hostname()
	if err == nil {
		newkey.Hostname = hn
	}

	usr, err := user.Current()
	if err == nil {
		newkey.Username = usr.Username
	}

	// generate random salt
	newkey.Salt = make([]byte, scryptSaltsize)
	n, err := rand.Read(newkey.Salt)
	if n != scryptSaltsize || err != nil {
		panic("unable to read enough random bytes for salt")
	}

	// call KDF to derive user key
	newkey.user, err = crypto.KDF(newkey.N, newkey.R, newkey.P, newkey.Salt, password)
	if err != nil {
		return nil, err
	}

	if template == nil {
		// generate new random master keys
		newkey.master = crypto.NewRandomKey()
	} else {
		// copy master keys from old key
		newkey.master = template
	}

	// encrypt master keys (as json) with user key
	buf, err := json.Marshal(newkey.master)
	if err != nil {
		return nil, err
	}

	newkey.Data, err = crypto.Encrypt(newkey.user, nil, buf)

	// dump as json
	buf, err = json.Marshal(newkey)
	if err != nil {
		return nil, err
	}

	// store in repository and return
	h := backend.Handle{
		Type: backend.Key,
		Name: backend.Hash(buf).String(),
	}

	err = s.be.Save(h, buf)
	if err != nil {
		return nil, err
	}

	newkey.name = h.Name

	return newkey, nil
}

func (k *Key) String() string {
	if k == nil {
		return "<Key nil>"
	}
	return fmt.Sprintf("<Key of %s@%s, created on %s>", k.Username, k.Hostname, k.Created)
}

// Name returns an identifier for the key.
func (k Key) Name() string {
	return k.name
}

// Valid tests whether the mac and encryption keys are valid (i.e. not zero)
func (k *Key) Valid() bool {
	return k.user.Valid() && k.master.Valid()
}
<commit_msg>Change ErrNoKeyFound message<commit_after>package repository

import (
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/user"
	"time"

	"restic/backend"
	"restic/crypto"
	"restic/debug"
)

var (
	// ErrNoKeyFound is returned when no key for the repository could be decrypted.
	ErrNoKeyFound = errors.New("wrong password or no key found")
)

// TODO: figure out scrypt values on the fly depending on the current
// hardware.
const (
	scryptN        = 65536
	scryptR        = 8
	scryptP        = 1
	scryptSaltsize = 64
)

// Key represents an encrypted master key for a repository.
type Key struct {
	Created  time.Time `json:"created"`
	Username string    `json:"username"`
	Hostname string    `json:"hostname"`

	KDF  string `json:"kdf"`
	N    int    `json:"N"`
	R    int    `json:"r"`
	P    int    `json:"p"`
	Salt []byte `json:"salt"`
	Data []byte `json:"data"`

	user   *crypto.Key
	master *crypto.Key

	name string
}

// createMasterKey creates a new master key in the given backend and encrypts
// it with the password.
func createMasterKey(s *Repository, password string) (*Key, error) {
	return AddKey(s, password, nil)
}

// OpenKey tries do decrypt the key specified by name with the given password.
func OpenKey(s *Repository, name string, password string) (*Key, error) {
	k, err := LoadKey(s, name)
	if err != nil {
		debug.Log("OpenKey", "LoadKey(%v) returned error %v", name[:12], err)
		return nil, err
	}

	// check KDF
	if k.KDF != "scrypt" {
		return nil, errors.New("only supported KDF is scrypt()")
	}

	// derive user key
	k.user, err = crypto.KDF(k.N, k.R, k.P, k.Salt, password)
	if err != nil {
		return nil, err
	}

	// decrypt master keys
	buf, err := crypto.Decrypt(k.user, []byte{}, k.Data)
	if err != nil {
		return nil, err
	}

	// restore json
	k.master = &crypto.Key{}
	err = json.Unmarshal(buf, k.master)
	if err != nil {
		debug.Log("OpenKey", "Unmarshal() returned error %v", err)
		return nil, err
	}
	k.name = name

	if !k.Valid() {
		return nil, errors.New("Invalid key for repository")
	}

	return k, nil
}

// SearchKey tries to decrypt all keys in the backend with the given password.
// If none could be found, ErrNoKeyFound is returned.
func SearchKey(s *Repository, password string) (*Key, error) {
	// try all keys in repo
	done := make(chan struct{})
	defer close(done)
	for name := range s.Backend().List(backend.Key, done) {
		debug.Log("SearchKey", "trying key %v", name[:12])
		key, err := OpenKey(s, name, password)
		if err != nil {
			debug.Log("SearchKey", "key %v returned error %v", name[:12], err)
			continue
		}

		debug.Log("SearchKey", "successfully opened key %v", name[:12])
		return key, nil
	}

	return nil, ErrNoKeyFound
}

// LoadKey loads a key from the backend.
func LoadKey(s *Repository, name string) (k *Key, err error) {
	h := backend.Handle{Type: backend.Key, Name: name}
	data, err := backend.LoadAll(s.be, h, nil)
	if err != nil {
		return nil, err
	}

	k = &Key{}
	err = json.Unmarshal(data, k)
	if err != nil {
		return nil, err
	}

	return k, nil
}

// AddKey adds a new key to an already existing repository.
func AddKey(s *Repository, password string, template *crypto.Key) (*Key, error) {
	// fill meta data about key
	newkey := &Key{
		Created: time.Now(),
		KDF:     "scrypt",
		N:       scryptN,
		R:       scryptR,
		P:       scryptP,
	}

	hn, err := os.Hostname()
	if err == nil {
		newkey.Hostname = hn
	}

	usr, err := user.Current()
	if err == nil {
		newkey.Username = usr.Username
	}

	// generate random salt
	newkey.Salt = make([]byte, scryptSaltsize)
	n, err := rand.Read(newkey.Salt)
	if n != scryptSaltsize || err != nil {
		panic("unable to read enough random bytes for salt")
	}

	// call KDF to derive user key
	newkey.user, err = crypto.KDF(newkey.N, newkey.R, newkey.P, newkey.Salt, password)
	if err != nil {
		return nil, err
	}

	if template == nil {
		// generate new random master keys
		newkey.master = crypto.NewRandomKey()
	} else {
		// copy master keys from old key
		newkey.master = template
	}

	// encrypt master keys (as json) with user key
	buf, err := json.Marshal(newkey.master)
	if err != nil {
		return nil, err
	}

	newkey.Data, err = crypto.Encrypt(newkey.user, nil, buf)

	// dump as json
	buf, err = json.Marshal(newkey)
	if err != nil {
		return nil, err
	}

	// store in repository and return
	h := backend.Handle{
		Type: backend.Key,
		Name: backend.Hash(buf).String(),
	}

	err = s.be.Save(h, buf)
	if err != nil {
		return nil, err
	}

	newkey.name = h.Name

	return newkey, nil
}

func (k *Key) String() string {
	if k == nil {
		return "<Key nil>"
	}
	return fmt.Sprintf("<Key of %s@%s, created on %s>", k.Username, k.Hostname, k.Created)
}

// Name returns an identifier for the key.
func (k Key) Name() string {
	return k.name
}

// Valid tests whether the mac and encryption keys are valid (i.e. not zero)
func (k *Key) Valid() bool {
	return k.user.Valid() && k.master.Valid()
}
<|endoftext|>


<commit_before>// Package to communicate with bitfinex exchange

package bitfinex

import (
	"crypto/hmac"
	"crypto/sha512"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"time"
)

const (
	APIURL = "https://api.bitfinex.com/" // Bitfinex API URL
)

// Stores Bitfinex API credentials
type API struct {
	APIKey    string
	APISecret string
}

// Error from exchange
type ErrorMessage struct {
	Message string `json:"message"` // Returned only on error
}

// Order book data from the exchange
type Orderbook struct {
	Bids []OrderbookItems // Slice of bid data items
	Asks []OrderbookItems // Slice of ask data items
}

// Inner order book data from the exchange
type OrderbookItems struct {
	Price     float64 `json:"price,string"`     // Order price
	Amount    float64 `json:"amount,string"`    // Order volume
	Timestamp float64 `json:"timestamp,string"` // Exchange timestamp
}

type Order struct {
	ID              int     `json:"order_id"`                   // Order ID
	Symbol          string  `json:"symbol"`                     // The symbol name the order belongs to
	Exchange        string  `json:"exchange"`                   // Exchange name "bitfinex"
	Price           float64 `json:"price,string"`               // The price the order was issued at (can be null for market orders)
	ExecutionPrice  float64 `json:"avg_execution_price,string"` // The average price at which this order as been executed so far. 0 if the order has not been executed at all
	Side            string  `json:"side"`                       // Either "buy" or "sell"
	OrderType       string  `json:"type"`                       // Either "market" / "limit" / "stop" / "trailing-stop"
	Timestamp       float64 `json:"timestamp,string"`           // The timestamp the order was submitted
	IsLive          bool    `json:"is_live,bool"`               // Could the order still be filled?
	IsCancelled     bool    `json:"is_cancelled,bool"`          // Has the order been cancelled?
	WasForced       bool    `json:"was_forced,bool"`            // For margin only: true if it was forced by the system
	ExecutedAmount  float64 `json:"executed_amount,string"`     // How much of the order has been executed so far in its history?
	RemainingAmount float64 `json:"remaining_amount,string"`    // How much is still remaining to be submitted?
	OriginalAmount  float64 `json:"original_amount,string"`     // What was the order originally submitted for?
}

// Return a new Bitfinex API instance
func New(key, secret string) (api *API) {
	api = &API{
		APIKey:    key,
		APISecret: secret,
	}
	return api
}

// Get book data from exchange
func (api *API) GetBook(url string) Orderbook {
	var book Orderbook
	data, err := api.get("/v1/book/" + url)
	if err != nil {
		log.Fatal(err)
	}
	err = json.Unmarshal(data, &book)
	if err != nil {
		log.Fatal(err)
	}

	return book
}

// Post new order

// Request
// symbol (string): The name of the symbol (see `/symbols`).
// amount (decimal): Order size: how much to buy or sell.
// price (price): Price to buy or sell at. Must be positive. Use random number for market orders.
// exchange (string): "bitfinex".
// side (string): Either "buy" or "sell".
// type (string): Either "market" / "limit" / "stop" / "trailing-stop" / "fill-or-kill" / "exchange market" / "exchange limit" / "exchange stop" / "exchange trailing-stop" / "exchange fill-or-kill". (type starting by "exchange " are exchange orders, others are margin trading orders)
// is_hidden (bool) true if the order should be hidden. Default is false.

// Response
// order_id (int): A randomly generated ID for the order.
// and the information given by /order/status
func (api *API) NewOrder(symbol string, amount, price float64, exchange, side, otype string, hidden bool) (Order, error) {
	var order Order
	request := struct {
		URL      string  `json:"request"`
		Nonce    string  `json:"nonce"`
		Symbol   string  `json:"symbol"`
		Amount   float64 `json:"amount,string"`
		Price    float64 `json:"price,string"`
		Exchange string  `json:"exchange"`
		Side     string  `json:"side"`
		Type     string  `json:"type"`
		Hidden   bool    `json:"is_hidden,bool"`
	}{
		"/v1/order/new",
		strconv.FormatInt(time.Now().UnixNano(), 10),
		symbol,
		amount,
		price,
		exchange,
		side,
		otype,
		hidden,
	}

	body, err := api.post(request.URL, request)
	if err != nil {
		return order, err
	}

	err = json.Unmarshal(body, &order)
	if err != nil || order.ID == 0 { // Failed to unmarshal expected message
		// Attempt to unmarshal the error message
		errorMessage := ErrorMessage{}
		err = json.Unmarshal(body, &errorMessage)
		if err != nil { // Not expected message and not expected error, bailing...
			return order, err
		}

		return order, errors.New("API: " + errorMessage.Message)
	}

	return order, nil
}

// API GET
func (api *API) get(url string) ([]byte, error) {
	resp, err := http.Get(APIURL + url)
	if err != nil {
		return []byte{}, err
	}
	defer resp.Body.Close()

	return ioutil.ReadAll(resp.Body)
}

// API POST method based on github.com/eAndrius/bitfinex-go
func (api *API) post(url string, payload interface{}) ([]byte, error) {
	// X-BFX-PAYLOAD
	payloadJSON, err := json.Marshal(payload)
	if err != nil {
		return []byte{}, err
	}
	payloadBase64 := base64.StdEncoding.EncodeToString(payloadJSON)

	// X-BFX-SIGNATURE
	h := hmac.New(sha512.New384, []byte(api.APISecret))
	h.Write([]byte(payloadBase64))
	signature := hex.EncodeToString(h.Sum(nil))

	// POST
	client := &http.Client{}
	req, err := http.NewRequest("POST", APIURL+url, nil)
	if err != nil {
		return []byte{}, err
	}

	req.Header.Add("X-BFX-APIKEY", api.APIKey)
	req.Header.Add("X-BFX-PAYLOAD", payloadBase64)
	req.Header.Add("X-BFX-SIGNATURE", signature)

	resp, err := client.Do(req)
	if err != nil {
		return []byte{}, err
	}
	defer resp.Body.Close()

	return ioutil.ReadAll(resp.Body)
}
<commit_msg>refactor order post methods<commit_after>// Bitfinex trading API

package bitfinex

import (
	"crypto/hmac"
	"crypto/sha512"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"time"
)

const (
	APIURL = "https://api.bitfinex.com/" // Bitfinex API URL
)

// Stores Bitfinex API credentials
type API struct {
	APIKey    string
	APISecret string
}

// Error from exchange
type ErrorMessage struct {
	Message string `json:"message"` // Returned only on error
}

// Order book data from the exchange
type Orderbook struct {
	Bids []OrderbookItems // Slice of bid data items
	Asks []OrderbookItems // Slice of ask data items
}

// Inner order book data from the exchange
type OrderbookItems struct {
	Price     float64 `json:"price,string"`     // Order price
	Amount    float64 `json:"amount,string"`    // Order volume
	Timestamp float64 `json:"timestamp,string"` // Exchange timestamp
}

// Order data to/from the exchange
type Order struct {
	ID              int     `json:"order_id"`                   // Order ID
	Symbol          string  `json:"symbol"`                     // The symbol name the order belongs to
	Exchange        string  `json:"exchange"`                   // Exchange name "bitfinex"
	Price           float64 `json:"price,string"`               // The price the order was issued at (can be null for market orders)
	ExecutionPrice  float64 `json:"avg_execution_price,string"` // The average price at which this order as been executed so far. 0 if the order has not been executed at all
	Side            string  `json:"side"`                       // Either "buy" or "sell"
	OrderType       string  `json:"type"`                       // Either "market" / "limit" / "stop" / "trailing-stop"
	Timestamp       float64 `json:"timestamp,string"`           // The timestamp the order was submitted
	IsLive          bool    `json:"is_live,bool"`               // Could the order still be filled?
	IsCancelled     bool    `json:"is_cancelled,bool"`          // Has the order been cancelled?
	WasForced       bool    `json:"was_forced,bool"`            // For margin only: true if it was forced by the system
	ExecutedAmount  float64 `json:"executed_amount,string"`     // How much of the order has been executed so far in its history?
	RemainingAmount float64 `json:"remaining_amount,string"`    // How much is still remaining to be submitted?
	OriginalAmount  float64 `json:"original_amount,string"`     // What was the order originally submitted for?
}

// Slice of orders
type Orders []Order

// Return a new Bitfinex API instance
func New(key, secret string) (api *API) {
	api = &API{
		APIKey:    key,
		APISecret: secret,
	}
	return api
}

// Get book data from the exchange
func (api *API) GetBook(url string) Orderbook {
	var book Orderbook
	data, err := api.get("/v1/book/" + url)
	if err != nil {
		log.Fatal(err)
	}
	err = json.Unmarshal(data, &book)
	if err != nil {
		log.Fatal(err)
	}

	return book
}

// Post new order to the exchange
func (api *API) NewOrder(symbol string, amount, price float64, exchange, side, otype string, hidden bool) (Order, error) {
	request := struct {
		URL      string  `json:"request"`
		Nonce    string  `json:"nonce"`
		Symbol   string  `json:"symbol"`
		Amount   float64 `json:"amount,string"`
		Price    float64 `json:"price,string"`
		Exchange string  `json:"exchange"`
		Side     string  `json:"side"`
		Type     string  `json:"type"`
		Hidden   bool    `json:"is_hidden,bool"`
	}{
		"/v1/order/new",
		strconv.FormatInt(time.Now().UnixNano(), 10),
		symbol,
		amount,
		price,
		exchange,
		side,
		otype,
		hidden,
	}

	return api.postOrder(request.URL, request)
}

// Cancel existing order on the exchange
func (api *API) CancelOrder(id int) (Order, error) {
	request := struct {
		URL     string `json:"request"`
		Nonce   string `json:"nonce"`
		OrderID int    `json:"order_id"`
	}{
		"/v1/order/cancel",
		strconv.FormatInt(time.Now().UnixNano(), 10),
		id,
	}

	return api.postOrder(request.URL, request)
}

// func (api *API) ReplaceOrder() (Order, error) {
// }

// func (api *API) OrderStatus() (Order, error) {
// }

// func (api *API) ActiveOrders() (Orders, error) {
// }

// Post Order info, used in order-related API methods
func (api *API) postOrder(url string, request interface{}) (Order, error) {
	var order Order

	body, err := api.post(url, request)
	if err != nil {
		return order, err
	}

	err = json.Unmarshal(body, &order)
	if err != nil || order.ID == 0 {
		errorMessage := ErrorMessage{}
		err = json.Unmarshal(body, &errorMessage)
		if err != nil {
			return order, err
		}

		return order, errors.New("API: " + errorMessage.Message)
	}

	return order, nil
}

// API GET
func (api *API) get(url string) ([]byte, error) {
	resp, err := http.Get(APIURL + url)
	if err != nil {
		return []byte{}, err
	}
	defer resp.Body.Close()

	return ioutil.ReadAll(resp.Body)
}

// API POST, from github.com/eAndrius/bitfinex-go
func (api *API) post(url string, payload interface{}) ([]byte, error) {
	// X-BFX-PAYLOAD
	payloadJSON, err := json.Marshal(payload)
	if err != nil {
		return []byte{}, err
	}
	payloadBase64 := base64.StdEncoding.EncodeToString(payloadJSON)

	// X-BFX-SIGNATURE
	h := hmac.New(sha512.New384, []byte(api.APISecret))
	h.Write([]byte(payloadBase64))
	signature := hex.EncodeToString(h.Sum(nil))

	// POST
	client := &http.Client{}
	req, err := http.NewRequest("POST", APIURL+url, nil)
	if err != nil {
		return []byte{}, err
	}

	req.Header.Add("X-BFX-APIKEY", api.APIKey)
	req.Header.Add("X-BFX-PAYLOAD", payloadBase64)
	req.Header.Add("X-BFX-SIGNATURE", signature)

	resp, err := client.Do(req)
	if err != nil {
		return []byte{}, err
	}
	defer resp.Body.Close()

	return ioutil.ReadAll(resp.Body)
}
<|endoftext|>


<commit_before>/*
Copyright 2022 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    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 web

import (
	"bufio"
	"bytes"
	"context"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/gravitational/trace"
	"github.com/jonboulle/clockwork"
	"github.com/stretchr/testify/require"
)

func TestWriteUpgradeResponse(t *testing.T) {
	var buf bytes.Buffer
	require.NoError(t, writeUpgradeResponse(&buf, "custom"))

	resp, err := http.ReadResponse(bufio.NewReader(&buf), nil)
	require.NoError(t, err)
	require.Equal(t, resp.StatusCode, http.StatusSwitchingProtocols)
	require.Equal(t, "custom", resp.Header.Get("Upgrade"))
}

func TestHandlerConnectionUpgrade(t *testing.T) {
	t.Parallel()

	expectedPayload := "hello@"
	alpnHandler := func(_ context.Context, conn net.Conn) error {
		// Handles connection asynchronously to verify web handler waits until
		// connection is closed.
		go func() {
			defer conn.Close()
			n, err := conn.Write([]byte(expectedPayload))
			require.NoError(t, err)
			require.Equal(t, len(expectedPayload), n)
		}()
		return nil
	}

	// Cherry picked some attributes to create a Handler to test only the
	// connection upgrade portion.
	h := &Handler{
		cfg: Config{
			ALPNHandler: alpnHandler,
		},
		log:   newPackageLogger(),
		clock: clockwork.NewRealClock(),
	}

	t.Run("unsupported type", func(t *testing.T) {
		r, err := http.NewRequest("GET", "http://localhost/webapi/connectionupgrade", nil)
		require.NoError(t, err)
		r.Header.Add("Upgrade", "unsupported-protocol")

		_, err = h.connectionUpgrade(httptest.NewRecorder(), r, nil)
		require.True(t, trace.IsBadParameter(err))
	})

	t.Run("upgraded to ALPN", func(t *testing.T) {
		serverConn, clientConn := net.Pipe()
		defer serverConn.Close()
		defer clientConn.Close()

		r, err := http.NewRequest("GET", "http://localhost/webapi/connectionupgrade", nil)
		require.NoError(t, err)
		r.Header.Add("Upgrade", "alpn")

		go func() {
			// serverConn will be hijacked.
			w := newResponseWriterHijacker(nil, serverConn)
			_, err = h.connectionUpgrade(w, r, nil)
			require.NoError(t, err)
		}()

		// Verify clientConn receives http.StatusSwitchingProtocols.
		clientConnReader := bufio.NewReader(clientConn)
		response, err := http.ReadResponse(clientConnReader, r)
		require.NoError(t, err)
		require.Equal(t, http.StatusSwitchingProtocols, response.StatusCode)

		// Verify clientConn receives data sent by Config.ALPNHandler.
		receive, err := clientConnReader.ReadString(byte('@'))
		require.NoError(t, err)
		require.Equal(t, expectedPayload, receive)
	})
}

// responseWriterHijacker is a mock http.ResponseWriter that also serves a
// net.Conn for http.Hijacker.
type responseWriterHijacker struct {
	http.ResponseWriter
	conn net.Conn
}

func newResponseWriterHijacker(w http.ResponseWriter, conn net.Conn) *responseWriterHijacker {
	if w == nil {
		w = httptest.NewRecorder()
	}
	return &responseWriterHijacker{
		ResponseWriter: w,
		conn:           conn,
	}
}

func (h responseWriterHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return h.conn, nil, nil
}
<commit_msg>Fix flaky TestHandlerConnectionUpgrade (#16129)<commit_after>/*
Copyright 2022 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    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 web

import (
	"bufio"
	"bytes"
	"context"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/gravitational/trace"
	"github.com/jonboulle/clockwork"
	"github.com/stretchr/testify/require"
)

func TestWriteUpgradeResponse(t *testing.T) {
	var buf bytes.Buffer
	require.NoError(t, writeUpgradeResponse(&buf, "custom"))

	resp, err := http.ReadResponse(bufio.NewReader(&buf), nil)
	require.NoError(t, err)
	require.Equal(t, resp.StatusCode, http.StatusSwitchingProtocols)
	require.Equal(t, "custom", resp.Header.Get("Upgrade"))
}

func TestHandlerConnectionUpgrade(t *testing.T) {
	t.Parallel()

	expectedPayload := "hello@"
	alpnHandler := func(_ context.Context, conn net.Conn) error {
		// Handles connection asynchronously to verify web handler waits until
		// connection is closed.
		go func() {
			defer conn.Close()
			n, err := conn.Write([]byte(expectedPayload))
			require.NoError(t, err)
			require.Equal(t, len(expectedPayload), n)
		}()
		return nil
	}

	// Cherry picked some attributes to create a Handler to test only the
	// connection upgrade portion.
	h := &Handler{
		cfg: Config{
			ALPNHandler: alpnHandler,
		},
		log:   newPackageLogger(),
		clock: clockwork.NewRealClock(),
	}

	t.Run("unsupported type", func(t *testing.T) {
		r, err := http.NewRequest("GET", "http://localhost/webapi/connectionupgrade", nil)
		require.NoError(t, err)
		r.Header.Add("Upgrade", "unsupported-protocol")

		_, err = h.connectionUpgrade(httptest.NewRecorder(), r, nil)
		require.True(t, trace.IsBadParameter(err))
	})

	t.Run("upgraded to ALPN", func(t *testing.T) {
		serverConn, clientConn := net.Pipe()
		defer serverConn.Close()
		defer clientConn.Close()

		r, err := http.NewRequest("GET", "http://localhost/webapi/connectionupgrade", nil)
		require.NoError(t, err)
		r.Header.Add("Upgrade", "alpn")

		go func() {
			// serverConn will be hijacked.
			w := newResponseWriterHijacker(nil, serverConn)
			_, err := h.connectionUpgrade(w, r, nil)
			require.NoError(t, err)
		}()

		// Verify clientConn receives http.StatusSwitchingProtocols.
		clientConnReader := bufio.NewReader(clientConn)
		response, err := http.ReadResponse(clientConnReader, r)
		require.NoError(t, err)
		require.Equal(t, http.StatusSwitchingProtocols, response.StatusCode)

		// Verify clientConn receives data sent by Config.ALPNHandler.
		receive, err := clientConnReader.ReadString(byte('@'))
		require.NoError(t, err)
		require.Equal(t, expectedPayload, receive)
	})
}

// responseWriterHijacker is a mock http.ResponseWriter that also serves a
// net.Conn for http.Hijacker.
type responseWriterHijacker struct {
	http.ResponseWriter
	conn net.Conn
}

func newResponseWriterHijacker(w http.ResponseWriter, conn net.Conn) *responseWriterHijacker {
	if w == nil {
		w = httptest.NewRecorder()
	}
	return &responseWriterHijacker{
		ResponseWriter: w,
		conn:           conn,
	}
}

func (h responseWriterHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return h.conn, nil, nil
}
<|endoftext|>