<commit_before>package deploy

import (
	"errors"
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/git"
	hub "github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/github"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/google/go-github/github"
)

const (
	Name  = "deploy"
	Usage = "A command for creating GitHub deployments"
)

const DefaultRef = "master"

func init() {
	cli.AppHelpTemplate = `USAGE:
   # Deploy the master branch of remind101/acme-inc to staging
   {{.Name}} --env=staging --ref=master remind101/acme-inc

   # Deploy HEAD of the current branch to staging
   {{.Name}} --env=staging remind101/acme-inc

   # Deploy the current GitHub repo to staging
   {{.Name}} --env=staging
{{if .Flags}}
OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}
`
}

var flags = []cli.Flag{
	cli.StringFlag{
		Name:  "ref, branch, commit, tag",
		Value: "",
		Usage: "The git ref to deploy. Can be a git commit, branch or tag.",
	},
	cli.StringFlag{
		Name:  "env, e",
		Value: "",
		Usage: "The environment to deploy to.",
	},
	cli.BoolFlag{
		Name:  "force, f",
		Usage: "Ignore failed tests.",
	},
}

// NewApp returns a new cli.App for the deploy command.
func NewApp() *cli.App {
	app := cli.NewApp()
	app.Version = "0.0.1"
	app.Name = Name
	app.Usage = Usage
	app.Flags = flags
	app.Action = func(c *cli.Context) {
		if err := RunDeploy(c); err != nil {
			msg := err.Error()
			if err, ok := err.(*github.ErrorResponse); ok {
				msg = err.Message
			}

			fmt.Println(msg)
			os.Exit(-1)
		}
	}

	return app
}

// RunDeploy performs a deploy.
func RunDeploy(c *cli.Context) error {
	w := c.App.Writer

	h, err := hub.CurrentConfig().PromptForHost("github.com")
	if err != nil {
		return err
	}

	client, err := newGitHubClient(h)
	if err != nil {
		return err
	}

	nwo, err := Repo(c.Args())
	if err != nil {
		return err
	}

	owner, repo, err := SplitRepo(nwo)
	if err != nil {
		return fmt.Errorf("Invalid GitHub repo: %s", nwo)
	}

	r, err := newDeploymentRequest(c)
	if err != nil {
		return err
	}

	fmt.Fprintf(w, "Creating deployment request of %s@%s to %s... ", nwo, *r.Ref, *r.Environment)

	d, _, err := client.Repositories.CreateDeployment(owner, repo, r)
	if err != nil {
		return err
	}

	ch := make(chan *github.DeploymentStatus)

	go func() {
		for {
			statuses, _, err := client.Repositories.ListDeploymentStatuses(owner, repo, *d.ID, nil)
			if err != nil {
				continue
			}

			completed := CompletedStatus(statuses)
			if completed != nil {
				ch <- completed
				break
			}
		}
	}()

	status := <-ch

	var url string
	if status.TargetURL != nil {
		url = *status.TargetURL
	}

	state := "unkown"
	if status.State != nil {
		state = *status.State
	}

	fmt.Fprintf(w, "%s: %s\n", state, url)

	return nil
}

func newDeploymentRequest(c *cli.Context) (*github.DeploymentRequest, error) {
	ref := c.String("ref")
	if ref == "" {
		r, err := git.Ref("HEAD")
		if err == nil {
			ref = r
		} else {
			ref = DefaultRef
		}
	}

	env := c.String("env")
	if env == "" {
		return nil, fmt.Errorf("--env flag is required")
	}

	var contexts *[]string
	if c.Bool("force") {
		s := []string{}
		contexts = &s
	}

	return &github.DeploymentRequest{
		Ref:              github.String(ref),
		Task:             github.String("deploy"),
		AutoMerge:        github.Bool(false),
		Environment:      github.String(env),
		RequiredContexts: contexts,
		// TODO Description:
	}, nil
}

var completedStatuses = []string{"success", "error", "failure"}

// CompletedStatus takes a slice of github.DeploymentStatus and returns the
// first "completed" status. nil is returned if there are no completed
// deployment states.
func CompletedStatus(statuses []github.DeploymentStatus) *github.DeploymentStatus {
	for _, ds := range statuses {
		for _, s := range completedStatuses {
			if ds.State != nil && *ds.State == s {
				return &ds
			}
		}
	}

	return nil
}

// Repo will determine the correct GitHub repo to deploy to, based on a set of
// arguments.
func Repo(arguments []string) (string, error) {
	if len(arguments) != 0 {
		return arguments[0], nil
	}

	remotes, err := hub.Remotes()
	if err != nil {
		return "", err
	}

	repo := GitHubRepo(remotes)
	if repo == "" {
		return repo, errors.New("no GitHub repo found in .git/config")
	}

	return repo, nil
}

// A regular expression that can convert a URL.Path into a GitHub repo name.
var remoteRegex = regexp.MustCompile(`^/(.*)\.git$`)

// GitHubRepo, given a list of git remotes, will determine what the GitHub repo
// is.
func GitHubRepo(remotes []hub.Remote) string {
	// We only want to look at the `origin` remote.
	remote := findRemote("origin", remotes)
	if remote == nil {
		return ""
	}

	// Remotes that are not pointed at a GitHub repo are not valid.
	if remote.URL.Host != "github.com" {
		return ""
	}

	// Convert `/remind101/acme-inc.git` => `remind101/acme-inc`.
	return remoteRegex.ReplaceAllString(remote.URL.Path, "$1")
}

func findRemote(name string, remotes []hub.Remote) *hub.Remote {
	for _, r := range remotes {
		if r.Name == name {
			return &r
		}
	}

	return nil
}

var errInvalidRepo = errors.New("invalid repo")

// SplitRepo splits a repo string in the form remind101/acme-inc into it's owner
// and repo components.
func SplitRepo(nwo string) (owner string, repo string, err error) {
	parts := strings.Split(nwo, "/")

	if len(parts) != 2 {
		err = errInvalidRepo
		return
	}

	owner = parts[0]
	repo = parts[1]

	return
}
<commit_msg>Make commit status checks failure prettier.<commit_after>package deploy

import (
	"errors"
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/git"
	hub "github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/github"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/google/go-github/github"
)

const (
	Name  = "deploy"
	Usage = "A command for creating GitHub deployments"
)

const DefaultRef = "master"

func init() {
	cli.AppHelpTemplate = `USAGE:
   # Deploy the master branch of remind101/acme-inc to staging
   {{.Name}} --env=staging --ref=master remind101/acme-inc

   # Deploy HEAD of the current branch to staging
   {{.Name}} --env=staging remind101/acme-inc

   # Deploy the current GitHub repo to staging
   {{.Name}} --env=staging
{{if .Flags}}
OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}
`
}

var flags = []cli.Flag{
	cli.StringFlag{
		Name:  "ref, branch, commit, tag",
		Value: "",
		Usage: "The git ref to deploy. Can be a git commit, branch or tag.",
	},
	cli.StringFlag{
		Name:  "env, e",
		Value: "",
		Usage: "The environment to deploy to.",
	},
	cli.BoolFlag{
		Name:  "force, f",
		Usage: "Ignore failed tests.",
	},
}

// NewApp returns a new cli.App for the deploy command.
func NewApp() *cli.App {
	app := cli.NewApp()
	app.Version = "0.0.1"
	app.Name = Name
	app.Usage = Usage
	app.Flags = flags
	app.Action = func(c *cli.Context) {
		if err := RunDeploy(c); err != nil {
			msg := err.Error()
			if err, ok := err.(*github.ErrorResponse); ok {
				if strings.HasPrefix(err.Message, "Conflict: Commit status checks failed for") {
					msg = "Commit status checks failed. You can bypass commit status checks with the --force flag."
				} else {
					msg = err.Message
				}
			}

			fmt.Println(msg)
			os.Exit(-1)
		}
	}

	return app
}

// RunDeploy performs a deploy.
func RunDeploy(c *cli.Context) error {
	w := c.App.Writer

	h, err := hub.CurrentConfig().PromptForHost("github.com")
	if err != nil {
		return err
	}

	client, err := newGitHubClient(h)
	if err != nil {
		return err
	}

	nwo, err := Repo(c.Args())
	if err != nil {
		return err
	}

	owner, repo, err := SplitRepo(nwo)
	if err != nil {
		return fmt.Errorf("Invalid GitHub repo: %s", nwo)
	}

	r, err := newDeploymentRequest(c)
	if err != nil {
		return err
	}

	fmt.Fprintf(w, "Creating deployment request of %s@%s to %s... ", nwo, *r.Ref, *r.Environment)

	d, _, err := client.Repositories.CreateDeployment(owner, repo, r)
	if err != nil {
		return err
	}

	ch := make(chan *github.DeploymentStatus)

	go func() {
		for {
			statuses, _, err := client.Repositories.ListDeploymentStatuses(owner, repo, *d.ID, nil)
			if err != nil {
				continue
			}

			completed := CompletedStatus(statuses)
			if completed != nil {
				ch <- completed
				break
			}
		}
	}()

	status := <-ch

	var url string
	if status.TargetURL != nil {
		url = *status.TargetURL
	}

	state := "unkown"
	if status.State != nil {
		state = *status.State
	}

	fmt.Fprintf(w, "%s: %s\n", state, url)

	return nil
}

func newDeploymentRequest(c *cli.Context) (*github.DeploymentRequest, error) {
	ref := c.String("ref")
	if ref == "" {
		r, err := git.Ref("HEAD")
		if err == nil {
			ref = r
		} else {
			ref = DefaultRef
		}
	}

	env := c.String("env")
	if env == "" {
		return nil, fmt.Errorf("--env flag is required")
	}

	var contexts *[]string
	if c.Bool("force") {
		s := []string{}
		contexts = &s
	}

	return &github.DeploymentRequest{
		Ref:              github.String(ref),
		Task:             github.String("deploy"),
		AutoMerge:        github.Bool(false),
		Environment:      github.String(env),
		RequiredContexts: contexts,
		// TODO Description:
	}, nil
}

var completedStatuses = []string{"success", "error", "failure"}

// CompletedStatus takes a slice of github.DeploymentStatus and returns the
// first "completed" status. nil is returned if there are no completed
// deployment states.
func CompletedStatus(statuses []github.DeploymentStatus) *github.DeploymentStatus {
	for _, ds := range statuses {
		for _, s := range completedStatuses {
			if ds.State != nil && *ds.State == s {
				return &ds
			}
		}
	}

	return nil
}

// Repo will determine the correct GitHub repo to deploy to, based on a set of
// arguments.
func Repo(arguments []string) (string, error) {
	if len(arguments) != 0 {
		return arguments[0], nil
	}

	remotes, err := hub.Remotes()
	if err != nil {
		return "", err
	}

	repo := GitHubRepo(remotes)
	if repo == "" {
		return repo, errors.New("no GitHub repo found in .git/config")
	}

	return repo, nil
}

// A regular expression that can convert a URL.Path into a GitHub repo name.
var remoteRegex = regexp.MustCompile(`^/(.*)\.git$`)

// GitHubRepo, given a list of git remotes, will determine what the GitHub repo
// is.
func GitHubRepo(remotes []hub.Remote) string {
	// We only want to look at the `origin` remote.
	remote := findRemote("origin", remotes)
	if remote == nil {
		return ""
	}

	// Remotes that are not pointed at a GitHub repo are not valid.
	if remote.URL.Host != "github.com" {
		return ""
	}

	// Convert `/remind101/acme-inc.git` => `remind101/acme-inc`.
	return remoteRegex.ReplaceAllString(remote.URL.Path, "$1")
}

func findRemote(name string, remotes []hub.Remote) *hub.Remote {
	for _, r := range remotes {
		if r.Name == name {
			return &r
		}
	}

	return nil
}

var errInvalidRepo = errors.New("invalid repo")

// SplitRepo splits a repo string in the form remind101/acme-inc into it's owner
// and repo components.
func SplitRepo(nwo string) (owner string, repo string, err error) {
	parts := strings.Split(nwo, "/")

	if len(parts) != 2 {
		err = errInvalidRepo
		return
	}

	owner = parts[0]
	repo = parts[1]

	return
}
<|endoftext|>


<commit_before>package som

import (
	"fmt"
	"io"
	"math/rand"
	"runtime"
	"sync"
	"time"

	"github.com/gonum/matrix/mat64"
	"github.com/milosgajdos83/gosom/pkg/utils"
)

// CodebookInitFunc defines SOM codebook initialization function
type CodebookInitFunc func(*mat64.Dense, []int) (*mat64.Dense, error)

// CoordsInitFunc defines SOM grid coordinates initialization function
type CoordsInitFunc func(string, []int) (*mat64.Dense, error)

// NeighbFunc defines SOM neighbourhood function
type NeighbFunc func(float64, float64) float64

// Map is a Self Organizing Map (SOM)
type Map struct {
	// codebook is a matrix which contains SOM codebook vectors
	// codebook dimensions: SOM units x data features
	codebook *mat64.Dense
	// unitDist is a symmetric hollow matrix that maps distances between SOM units
	// unitDist dimesions: SOM units x SOM units
	unitDist *mat64.Dense
	// bmus stores indices of Best Match Units (BMU) for each input data
	bmus []int
}

// NewMap creates new SOM based on the provided configuration and input data
// NewMap allows you to pass in SOM codebook init function that is used to initialize
// SOM codebook vectors to initial values. If codebook InitFunc is nil, random initialization
// is used. NewMap returns error if the provided configuration is not valid or if the data matrix
// is nil or if the codebook matrix could not be initialized.
func NewMap(c *MapConfig, data *mat64.Dense) (*Map, error) {
	// if input data is empty throw error
	if data == nil {
		return nil, fmt.Errorf("Invalid input data: %v\n", data)
	}
	// validate the map configuration
	if err := validateMapConfig(c); err != nil {
		return nil, err
	}
	// compute the number of map units
	mUnits := utils.IntProduct(c.Dims)
	if mUnits <= 1 {
		return nil, fmt.Errorf("Incorrect map size dimensions: %v\n", c.Dims)
	}
	// initialize codebook
	codebook, err := c.InitFunc(data, c.Dims)
	if err != nil {
		return nil, err
	}
	// grid coordinates matrix
	gridCoords, err := GridCoords(c.UShape, c.Dims)
	if err != nil {
		return nil, err
	}
	// unit distance matrix
	unitDist, err := DistanceMx("euclidean", gridCoords)
	if err != nil {
		return nil, err
	}
	rows, _ := data.Dims()
	bmus := make([]int, rows)
	// return pointer to new map
	return &Map{
		codebook: codebook,
		unitDist: unitDist,
		bmus:     bmus,
	}, nil
}

// Codebook returns a matrix which contains SOM codebook vectors
func (m Map) Codebook() *mat64.Dense {
	return m.codebook
}

// UnitDist returns a matrix which contains Euclidean distances between SOM units
func (m Map) UnitDist() *mat64.Dense {
	return m.unitDist
}

// BMUs returns a slice which contains indices of Best Match Units (BMUs) of each input vector
func (m Map) BMUs() []int {
	return m.bmus
}

// MarshalTo serializes SOM codebook in a given format to writer w.
// At the moment only the native gonum binary format is supported.
// It returns the number of bytes written to w or fails with error.
func (m *Map) MarshalTo(format string, w io.Writer) (int, error) {
	switch format {
	case "gonum":
		return m.codebook.MarshalBinaryTo(w)
	}
	// marshal binary to file path
	return 0, fmt.Errorf("Unsupported format: %s\n", format)
}

// UMatrixOut generates SOM u-matrix in a given format and writes the output to w.
// At the moment only SVG format is supported. It fails with error if the write to w fails.
func (m Map) UMatrixOut(format, title string, w io.Writer) error {
	// TODO: needs some UMatrixSVG modifications first
	return nil
}

// Train runs a SOM training for a given data set and training configuration parameters.
// It modifies the map codebook vectors based on the chosen training algorithm.
// It returns error if the supplied training configuration is invalid or training fails
func (m *Map) Train(c *TrainConfig, data *mat64.Dense, iters int) error {
	// number of iterations must be a positive integer
	if iters <= 0 {
		return fmt.Errorf("Invalid number of iterations: %d\n", iters)
	}
	// nil data passed in
	if data == nil {
		return fmt.Errorf("Invalid data supplied: %v\n", data)
	}
	// validate the training configuration
	if err := validateTrainConfig(c); err != nil {
		return err
	}
	// run the training
	switch c.Method {
	case "seq":
		return m.seqTrain(c, data, iters)
	case "batch":
		return m.batchTrain(c, data, iters)
	}

	return nil
}

// seqTrain runs sequential SOM training algorithm on a given data set
func (m *Map) seqTrain(tc *TrainConfig, data *mat64.Dense, iters int) error {
	rows, _ := data.Dims()
	// create random number generator
	rSrc := rand.NewSource(time.Now().UnixNano())
	r := rand.New(rSrc)
	// retrieve Neighbourhood function
	neighbFn := Neighb[tc.NeighbFn]
	// perform iters number of learning iterations
	for i := 0; i < iters; i++ {
		// pick a random sample from dataset
		sample := data.RowView(r.Intn(rows))
		// no need to check for error here:
		// sample and codebook are not nil and have the same dimension
		bmu, _ := ClosestVec("euclidean", sample, m.codebook)
		// no need to check for errors:
		// LRate and Radius are checked by config validation
		lRate, _ := LRate(i, iters, tc.LDecay, tc.LRate)
		radius, _ := Radius(i, iters, tc.RDecay, tc.Radius)
		// pick the bmu unit distance row
		bmuDists := m.unitDist.RowView(bmu)
		// find units which are within the radius
		for i := 0; i < bmuDists.Len(); i++ {
			// bmu distance to i-th map unit
			dist := bmuDists.At(i, 0)
			// we are within BMU radius
			if dist < radius {
				// update particular codebook vector
				m.seqUpdateCbVec(i, sample, lRate, radius, dist, neighbFn)
			}
		}
	}

	return nil
}

// seqUpdateCbVec updates codebook vector on row cbIdx given the learning rate l,
// radius r, distance d and neihgbourhood function nFn
func (m *Map) seqUpdateCbVec(cbIdx int, vec *mat64.Vector, l, r, d float64, nFn NeighbFunc) {
	// pick codebook vector that should be updated
	cbVec := m.codebook.RowView(cbIdx)
	// update codebook vector according to the algorithm
	diff := mat64.NewVector(cbVec.Len(), nil)
	diff.AddScaledVec(vec, -1.0, cbVec)
	mul := l
	// nFn returns 1 for d == 0; skipping this case will save us some CPU time
	if d > 0.0 {
		mul *= nFn(d, r)
	}
	cbVec.AddScaledVec(cbVec, mul, diff)
}

// batchConfig holds batch training configuration
type batchConfig struct {
	// tc is SOM training configuration
	tc *TrainConfig
	// iters is a number of batch iterations
	iters int
}

// batchResult holds result of batch algorithm for a particular data input
// It holds scaled data vector, neighbourhood of its BMU and index of codebook vector
type batchResult struct {
	// vec is a scaled data vector
	vec *mat64.Vector
	// nghb is vec BMU neighbourhood
	nghb float64
	// idx is an index of codebook vector to update
	idx int
}

// batchTrain runs batch SOM training on a given data set
func (m *Map) batchTrain(tc *TrainConfig, data *mat64.Dense, iters int) error {
	cbRows, _ := m.codebook.Dims()
	rows, _ := data.Dims()
	// batchConfig holds training config and number of iterations
	bc := &batchConfig{
		tc:    tc,
		iters: iters,
	}
	// number of worker goroutines
	workers := runtime.NumCPU()
	// evenly distribute batch work between workers
	workerBatch := rows / workers
	// train for a number of iterations
	for i := 0; i < iters; i++ {
		// reset from index and input count
		from := 0
		count := workerBatch
		// create batch results channel
		results := make(chan *batchResult, workers*4)
		wg := &sync.WaitGroup{}
		// start worker goroutines
		for j := 0; j < workers; j++ {
			// from is data matrix row pointer
			from += j * workerBatch
			// last worker will work through the batch reminder
			if j == workers-1 {
				count += rows % workers
			}
			// if we go over the number of rows adjust bSamples
			if from+count > rows {
				count = rows - from
			}
			wg.Add(1)
			go m.processBatch(results, wg, bc, data, from, count, i)
		}
		// wait for workers to finish and close the result channel
		go func() {
			wg.Wait()
			close(results)
		}()
		// collect batch results from all workers
		cbVecs := make([]*mat64.Vector, cbRows)
		nghbs := make([]float64, cbRows)
		for result := range results {
			if cbVecs[result.idx] != nil {
				cbVecs[result.idx].AddVec(cbVecs[result.idx], result.vec)
			} else {
				cbVecs[result.idx] = result.vec
			}
			nghbs[result.idx] += result.nghb
		}
		// update codebook vectors
		for k := 0; k < cbRows; k++ {
			if cbVecs[k] != nil {
				cbVecs[k].ScaleVec(1.0/nghbs[k], cbVecs[k])
				m.codebook.SetRow(k, cbVecs[k].RawVector().Data)
			}
		}
	}

	return nil
}

// processRow processes data rows and sends tehm down the results channel
func (m Map) processBatch(res chan<- *batchResult, wg *sync.WaitGroup,
	bc *batchConfig, data *mat64.Dense, from, count, iter int) {
	// retrieve Neighbourhood function
	neighbFn := Neighb[bc.tc.NeighbFn]
	for i := from; i < count+from; i++ {
		row := data.RowView(i)
		// find codebook BMU for this data row
		bmu, _ := ClosestVec("euclidean", row, m.codebook)
		// calculate radius for this iteration
		radius, _ := Radius(iter, bc.iters, bc.tc.RDecay, bc.tc.Radius)
		// pick the BMU's distance row
		bmuDists := m.unitDist.RowView(bmu)
		for j := 0; j < bmuDists.Len(); j++ {
			// bmu distance to i-th map unit
			dist := bmuDists.At(j, 0)
			// when in BMU radius, scale and add to all neighbourhood vecs
			if dist < radius {
				// calculate neighbourhood function
				nghb := neighbFn(dist, radius)
				// allocate new vector and copy row data to it
				vec := new(mat64.Vector)
				vec.CloneVec(row)
				vec.ScaleVec(nghb, vec)
				// send batchResult down results channel
				res <- &batchResult{vec: vec, nghb: nghb, idx: j}
			}
		}
	}
	wg.Done()
}
<commit_msg>fixed a bug in batch algorithm<commit_after>package som

import (
	"fmt"
	"io"
	"math/rand"
	"runtime"
	"sync"
	"time"

	"github.com/gonum/matrix/mat64"
	"github.com/milosgajdos83/gosom/pkg/utils"
)

// CodebookInitFunc defines SOM codebook initialization function
type CodebookInitFunc func(*mat64.Dense, []int) (*mat64.Dense, error)

// CoordsInitFunc defines SOM grid coordinates initialization function
type CoordsInitFunc func(string, []int) (*mat64.Dense, error)

// NeighbFunc defines SOM neighbourhood function
type NeighbFunc func(float64, float64) float64

// Map is a Self Organizing Map (SOM)
type Map struct {
	// codebook is a matrix which contains SOM codebook vectors
	// codebook dimensions: SOM units x data features
	codebook *mat64.Dense
	// unitDist is a symmetric hollow matrix that maps distances between SOM units
	// unitDist dimesions: SOM units x SOM units
	unitDist *mat64.Dense
	// bmus stores indices of Best Match Units (BMU) for each input data
	bmus []int
}

// NewMap creates new SOM based on the provided configuration and input data
// NewMap allows you to pass in SOM codebook init function that is used to initialize
// SOM codebook vectors to initial values. If codebook InitFunc is nil, random initialization
// is used. NewMap returns error if the provided configuration is not valid or if the data matrix
// is nil or if the codebook matrix could not be initialized.
func NewMap(c *MapConfig, data *mat64.Dense) (*Map, error) {
	// if input data is empty throw error
	if data == nil {
		return nil, fmt.Errorf("Invalid input data: %v\n", data)
	}
	// validate the map configuration
	if err := validateMapConfig(c); err != nil {
		return nil, err
	}
	// compute the number of map units
	mUnits := utils.IntProduct(c.Dims)
	if mUnits <= 1 {
		return nil, fmt.Errorf("Incorrect map size dimensions: %v\n", c.Dims)
	}
	// initialize codebook
	codebook, err := c.InitFunc(data, c.Dims)
	if err != nil {
		return nil, err
	}
	// grid coordinates matrix
	gridCoords, err := GridCoords(c.UShape, c.Dims)
	if err != nil {
		return nil, err
	}
	// unit distance matrix
	unitDist, err := DistanceMx("euclidean", gridCoords)
	if err != nil {
		return nil, err
	}
	rows, _ := data.Dims()
	bmus := make([]int, rows)
	// return pointer to new map
	return &Map{
		codebook: codebook,
		unitDist: unitDist,
		bmus:     bmus,
	}, nil
}

// Codebook returns a matrix which contains SOM codebook vectors
func (m Map) Codebook() *mat64.Dense {
	return m.codebook
}

// UnitDist returns a matrix which contains Euclidean distances between SOM units
func (m Map) UnitDist() *mat64.Dense {
	return m.unitDist
}

// BMUs returns a slice which contains indices of Best Match Units (BMUs) of each input vector
func (m Map) BMUs() []int {
	return m.bmus
}

// MarshalTo serializes SOM codebook in a given format to writer w.
// At the moment only the native gonum binary format is supported.
// It returns the number of bytes written to w or fails with error.
func (m *Map) MarshalTo(format string, w io.Writer) (int, error) {
	switch format {
	case "gonum":
		return m.codebook.MarshalBinaryTo(w)
	}
	// marshal binary to file path
	return 0, fmt.Errorf("Unsupported format: %s\n", format)
}

// UMatrixOut generates SOM u-matrix in a given format and writes the output to w.
// At the moment only SVG format is supported. It fails with error if the write to w fails.
func (m Map) UMatrixOut(format, title string, w io.Writer) error {
	// TODO: needs some UMatrixSVG modifications first
	return nil
}

// Train runs a SOM training for a given data set and training configuration parameters.
// It modifies the map codebook vectors based on the chosen training algorithm.
// It returns error if the supplied training configuration is invalid or training fails
func (m *Map) Train(c *TrainConfig, data *mat64.Dense, iters int) error {
	// number of iterations must be a positive integer
	if iters <= 0 {
		return fmt.Errorf("Invalid number of iterations: %d\n", iters)
	}
	// nil data passed in
	if data == nil {
		return fmt.Errorf("Invalid data supplied: %v\n", data)
	}
	// validate the training configuration
	if err := validateTrainConfig(c); err != nil {
		return err
	}
	// run the training
	switch c.Method {
	case "seq":
		return m.seqTrain(c, data, iters)
	case "batch":
		return m.batchTrain(c, data, iters)
	}

	return nil
}

// seqTrain runs sequential SOM training algorithm on a given data set
func (m *Map) seqTrain(tc *TrainConfig, data *mat64.Dense, iters int) error {
	rows, _ := data.Dims()
	// create random number generator
	rSrc := rand.NewSource(time.Now().UnixNano())
	r := rand.New(rSrc)
	// retrieve Neighbourhood function
	neighbFn := Neighb[tc.NeighbFn]
	// perform iters number of learning iterations
	for i := 0; i < iters; i++ {
		// pick a random sample from dataset
		sample := data.RowView(r.Intn(rows))
		// no need to check for error here:
		// sample and codebook are not nil and have the same dimension
		bmu, _ := ClosestVec("euclidean", sample, m.codebook)
		// no need to check for errors:
		// LRate and Radius are checked by config validation
		lRate, _ := LRate(i, iters, tc.LDecay, tc.LRate)
		radius, _ := Radius(i, iters, tc.RDecay, tc.Radius)
		// pick the bmu unit distance row
		bmuDists := m.unitDist.RowView(bmu)
		// find units which are within the radius
		for i := 0; i < bmuDists.Len(); i++ {
			// bmu distance to i-th map unit
			dist := bmuDists.At(i, 0)
			// we are within BMU radius
			if dist < radius {
				// update particular codebook vector
				m.seqUpdateCbVec(i, sample, lRate, radius, dist, neighbFn)
			}
		}
	}

	return nil
}

// seqUpdateCbVec updates codebook vector on row cbIdx given the learning rate l,
// radius r, distance d and neihgbourhood function nFn
func (m *Map) seqUpdateCbVec(cbIdx int, vec *mat64.Vector, l, r, d float64, nFn NeighbFunc) {
	// pick codebook vector that should be updated
	cbVec := m.codebook.RowView(cbIdx)
	// update codebook vector according to the algorithm
	diff := mat64.NewVector(cbVec.Len(), nil)
	diff.AddScaledVec(vec, -1.0, cbVec)
	mul := l
	// nFn returns 1 for d == 0; skipping this case will save us some CPU time
	if d > 0.0 {
		mul *= nFn(d, r)
	}
	cbVec.AddScaledVec(cbVec, mul, diff)
}

// batchConfig holds batch training configuration
type batchConfig struct {
	// tc is SOM training configuration
	tc *TrainConfig
	// iters is a number of batch iterations
	iters int
}

// batchResult holds result of batch algorithm for a particular data input
// It holds scaled data vector, neighbourhood of its BMU and index of codebook vector
type batchResult struct {
	// vec is a scaled data vector
	vec *mat64.Vector
	// nghb is vec BMU neighbourhood
	nghb float64
	// idx is an index of codebook vector to update
	idx int
}

// batchTrain runs batch SOM training on a given data set
func (m *Map) batchTrain(tc *TrainConfig, data *mat64.Dense, iters int) error {
	cbRows, _ := m.codebook.Dims()
	rows, _ := data.Dims()
	// batchConfig holds training config and number of iterations
	bc := &batchConfig{
		tc:    tc,
		iters: iters,
	}
	// number of worker goroutines
	workers := runtime.NumCPU()
	// evenly distribute batch work between workers
	workerBatch := rows / workers
	// train for a number of iterations
	for i := 0; i < iters; i++ {
		// reset from index and input count
		from := 0
		count := workerBatch
		// create batch results channel
		results := make(chan *batchResult, workers*4)
		wg := &sync.WaitGroup{}
		// start worker goroutines
		for j := 0; j < workers; j++ {
			// from is data matrix row pointer
			from = j * workerBatch
			// last worker will work through the batch reminder
			if j == workers-1 {
				count += rows % workers
			}
			// if we go over the number of rows adjust bSamples
			if from+count > rows {
				count = rows - from
			}
			wg.Add(1)
			go m.processBatch(results, wg, bc, data, from, count, i)
		}
		// wait for workers to finish and close the result channel
		go func() {
			wg.Wait()
			close(results)
		}()
		// collect batch results from all workers
		cbVecs := make([]*mat64.Vector, cbRows)
		nghbs := make([]float64, cbRows)
		for result := range results {
			if cbVecs[result.idx] != nil {
				cbVecs[result.idx].AddVec(cbVecs[result.idx], result.vec)
			} else {
				cbVecs[result.idx] = result.vec
			}
			nghbs[result.idx] += result.nghb
		}
		// update codebook vectors
		for k := 0; k < cbRows; k++ {
			if cbVecs[k] != nil {
				cbVecs[k].ScaleVec(1.0/nghbs[k], cbVecs[k])
				m.codebook.SetRow(k, cbVecs[k].RawVector().Data)
			}
		}
	}

	return nil
}

// processRow processes data rows and sends tehm down the results channel
func (m Map) processBatch(res chan<- *batchResult, wg *sync.WaitGroup,
	bc *batchConfig, data *mat64.Dense, from, count, iter int) {
	// retrieve Neighbourhood function
	neighbFn := Neighb[bc.tc.NeighbFn]
	for i := from; i < count+from; i++ {
		row := data.RowView(i)
		// find codebook BMU for this data row
		bmu, _ := ClosestVec("euclidean", row, m.codebook)
		// calculate radius for this iteration
		radius, _ := Radius(iter, bc.iters, bc.tc.RDecay, bc.tc.Radius)
		// pick the BMU's distance row
		bmuDists := m.unitDist.RowView(bmu)
		for j := 0; j < bmuDists.Len(); j++ {
			// bmu distance to i-th map unit
			dist := bmuDists.At(j, 0)
			// when in BMU radius, scale and add to all neighbourhood vecs
			if dist < radius {
				// calculate neighbourhood function
				nghb := neighbFn(dist, radius)
				// allocate new vector and copy row data to it
				vec := new(mat64.Vector)
				vec.CloneVec(row)
				vec.ScaleVec(nghb, vec)
				// send batchResult down results channel
				res <- &batchResult{vec: vec, nghb: nghb, idx: j}
			}
		}
	}
	wg.Done()
}
<|endoftext|>


<commit_before>package controller

import (
	"errors"
	"fmt"
	"socialapi/workers/emailnotifier/models"
	"socialapi/workers/helper"
	notificationmodels "socialapi/workers/notification/models"
	"time"

	"github.com/koding/logging"
	"github.com/koding/rabbitmq"
	"github.com/koding/worker"
	"github.com/streadway/amqp"
)

var emailConfig = map[string]string{
	notificationmodels.NotificationContent_TYPE_COMMENT: "comment",
	notificationmodels.NotificationContent_TYPE_LIKE:    "likeActivities",
	notificationmodels.NotificationContent_TYPE_FOLLOW:  "followActions",
	notificationmodels.NotificationContent_TYPE_JOIN:    "groupJoined",
	notificationmodels.NotificationContent_TYPE_LEAVE:   "groupLeft",
	notificationmodels.NotificationContent_TYPE_MENTION: "mention",
}

const (
	DAY         = 24 * time.Hour
	TIMEFORMAT  = "20060102"
	CACHEPREFIX = "dailymail"
)

type Action func(*Controller, []byte) error

type Controller struct {
	routes   map[string]Action
	log      logging.Logger
	rmqConn  *amqp.Connection
	settings *models.EmailSettings
}

func (n *Controller) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	n.log.Error("an error occured: %s", err)
	delivery.Ack(false)

	return false
}

func (n *Controller) HandleEvent(event string, data []byte) error {
	n.log.Debug("New Event Received %s", event)
	handler, ok := n.routes[event]
	if !ok {
		return worker.HandlerNotFoundErr
	}

	return handler(n, data)
}

func New(rmq *rabbitmq.RabbitMQ, log logging.Logger, es *EmailSettings) (*Controller, error) {
	rmqConn, err := rmq.Connect("NewEmailNotifierWorkerController")
	if err != nil {
		return nil, err
	}

	nwc := &Controller{
		log:      log,
		rmqConn:  rmqConn.Conn(),
		settings: es,
	}

	routes := map[string]Action{
		"notification.notification_created": (*Controller).SendInstantEmail,
		"notification.notification_updated": (*Controller).SendInstantEmail,
	}

	nwc.routes = routes

	return nwc, nil
}

func (n *Controller) initDailyEmailCron() {

	cronJob = cron.New()
	cronJob.AddFunc(SCHEDULE, n.sendDailyMails)
	cronJob.Start()
}

func (n *Controller) SendInstantEmail(data []byte) error {
	channel, err := n.rmqConn.Channel()
	if err != nil {
		return errors.New("channel connection error")
	}
	defer channel.Close()

	notification := notificationmodels.NewNotification()
	if err := notification.MapMessage(data); err != nil {
		return err
	}

	// fetch latest activity for checking actor
	activity, nc, err := notification.FetchLastActivity()
	if err != nil {
		return err
	}

	if !validNotification(activity, notification) {
		return nil
	}

	uc, err := models.FetchUserContact(notification.AccountId)
	if err != nil {
		return fmt.Errorf("an error occurred while fetching user contact: %s", err)
	}

	if !n.checkMailSettings(uc, activity, nc) {
		return nil
	}

	mc := models.NewMailerContainer()
	mc.AccountId = notification.AccountId
	mc.Activity = activity
	mc.Content = nc

	if err := mc.PrepareContainer(); err != nil {
		return err
	}

	mc.CreatedAt = notification.ActivatedAt

	tp := models.NewTemplateParser()
	tp.UserContact = uc
	body, err := tp.RenderInstantTemplate(mc)
	if err != nil {
		return fmt.Errorf("an error occurred while preparing notification email: %s", err)
	}

	tg := &models.TokenGenerator{
		UserContact:      uc,
		NotificationType: emailConfig[nc.TypeConstant],
	}

	if err := tg.CreateToken(); err != nil {
		return err
	}

	mailer := models.NewMailer()
	mailer.EmailSettings = n.settings
	mailer.UserContact = uc
	mailer.Body = body
	mailer.Subject = prepareSubject(mc)

	if err := mailer.SendMail(); err != nil {
		return err
	}

	n.log.Info("%s notified by email", uc.Username)

	return nil
}

func prepareSubject(mc *models.MailerContainer) string {
	t, err := mc.Content.GetContentType()
	if err != nil {
		return ""
	}

	return t.GetDefinition()
}

func validNotification(a *notificationmodels.NotificationActivity, n *notificationmodels.Notification) bool {
	// do not notify actor for her own action
	if a.ActorId == n.AccountId {
		return false
	}

	// do not notify user when notification is not yet activated
	return !n.ActivatedAt.IsZero()
}

func (n *EmailNotifierWorkerController) checkMailSettings(uc *models.UserContact,
	a *notificationmodels.NotificationActivity, nc *notificationmodels.NotificationContent) bool {
	// notifications are disabled
	if val := uc.EmailSettings["global"]; !val {
		return false
	}

	notificationEnabled := uc.EmailSettings[emailConfig[nc.TypeConstant]]
	// daily notifications are enabled
	if val := uc.EmailSettings["daily"]; val {
		if notificationEnabled {
			go n.saveDailyMail(uc.AccountId, a.Id)
		}

		return false
	}

	// get config
	return notificationEnabled
}

func (n *EmailNotifierWorkerController) saveDailyMail(accountId, activityId int64) {
	redisConn := helper.MustGetRedisConn()
	key := prepareSetterCacheKey(accountId)
	if _, err := redisConn.AddSetMembers(key, activityId); err != nil {
		n.log.Error("daily mail error: %s", err)
		return
	}

	if err := redisConn.Expire(key, DAY); err != nil {
		n.log.Error("daily mail error: %s", err)
	}
}

func containsObject(nc *models.NotificationContent) bool {
	return nc.TypeConstant == models.NotificationContent_TYPE_LIKE ||
		nc.TypeConstant == models.NotificationContent_TYPE_MENTION ||
		nc.TypeConstant == models.NotificationContent_TYPE_COMMENT
}

func fetchContentBody(nc *models.NotificationContent, cm *socialmodels.ChannelMessage) string {

	switch nc.TypeConstant {
	case models.NotificationContent_TYPE_LIKE:
		return cm.Body
	case models.NotificationContent_TYPE_MENTION:
		return cm.Body
	case models.NotificationContent_TYPE_COMMENT:
		return fetchLastReplyBody(cm.Id)
	}

	return ""
}

func fetchLastReplyBody(targetId int64) string {
	mr := socialmodels.NewMessageReply()
	mr.MessageId = targetId
	query := socialmodels.NewQuery()
	query.Limit = 1
	messages, err := mr.List(query)
	if err != nil {
		return ""
	}

	if len(messages) == 0 {
		return ""
	}

	return messages[0].Body
}

func fetchRepliedMessage(replyId int64) *socialmodels.ChannelMessage {
	mr := socialmodels.NewMessageReply()
	mr.ReplyId = replyId

	parent, err := mr.FetchRepliedMessage()
	if err != nil {
		parent = socialmodels.NewChannelMessage()
	}

	return parent
}

func (n *Controller) SendMail(uc *UserContact, body, subject string) error {
	es := n.settings
	sg := sendgrid.NewSendGridClient(es.Username, es.Password)
	fullname := fmt.Sprintf("%s %s", uc.FirstName, uc.LastName)

	message := sendgrid.NewMail()
	message.AddTo(uc.Email)
	message.AddToName(fullname)
	message.SetSubject(subject)
	message.SetHTML(body)
	message.SetFrom(es.FromMail)
	message.SetFromName(es.FromName)

	if err := sg.Send(message); err != nil {
		return fmt.Errorf("an error occurred while sending notification email to %s", uc.Username)
	}
	n.log.Info("%s notified by email", uc.Username)

	return nil
}

func prepareSetterCacheKey(accountId int64) string {
	return fmt.Sprintf("%s:%d:%s", CACHEPREFIX, accountId, time.Now().Format(TIMEFORMAT))
}
<commit_msg>Notification: each daily notification email recipient is added to a set in redis<commit_after>package controller

import (
	"errors"
	"fmt"
	"socialapi/config"
	"socialapi/workers/emailnotifier/models"
	"socialapi/workers/helper"
	notificationmodels "socialapi/workers/notification/models"
	"time"

	"github.com/koding/logging"
	"github.com/koding/rabbitmq"
	"github.com/koding/worker"
	"github.com/streadway/amqp"
)

var emailConfig = map[string]string{
	notificationmodels.NotificationContent_TYPE_COMMENT: "comment",
	notificationmodels.NotificationContent_TYPE_LIKE:    "likeActivities",
	notificationmodels.NotificationContent_TYPE_FOLLOW:  "followActions",
	notificationmodels.NotificationContent_TYPE_JOIN:    "groupJoined",
	notificationmodels.NotificationContent_TYPE_LEAVE:   "groupLeft",
	notificationmodels.NotificationContent_TYPE_MENTION: "mention",
}

const (
	DAY           = 24 * time.Hour
	TIMEFORMAT    = "20060102"
	CACHEPREFIX   = "dailymail"
	RECIPIENTSKEY = "recipients"
)

type Action func(*Controller, []byte) error

type Controller struct {
	routes   map[string]Action
	log      logging.Logger
	rmqConn  *amqp.Connection
	settings *models.EmailSettings
}

func (n *Controller) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	n.log.Error("an error occured: %s", err)
	delivery.Ack(false)

	return false
}

func (n *Controller) HandleEvent(event string, data []byte) error {
	n.log.Debug("New Event Received %s", event)
	handler, ok := n.routes[event]
	if !ok {
		return worker.HandlerNotFoundErr
	}

	return handler(n, data)
}

func New(rmq *rabbitmq.RabbitMQ, log logging.Logger, es *EmailSettings) (*Controller, error) {
	rmqConn, err := rmq.Connect("NewEmailNotifierWorkerController")
	if err != nil {
		return nil, err
	}

	nwc := &Controller{
		log:      log,
		rmqConn:  rmqConn.Conn(),
		settings: es,
	}

	routes := map[string]Action{
		"notification.notification_created": (*Controller).SendInstantEmail,
		"notification.notification_updated": (*Controller).SendInstantEmail,
	}

	nwc.routes = routes

	return nwc, nil
}

func (n *Controller) initDailyEmailCron() {

	cronJob = cron.New()
	cronJob.AddFunc(SCHEDULE, n.sendDailyMails)
	cronJob.Start()
}

func (n *Controller) SendInstantEmail(data []byte) error {
	channel, err := n.rmqConn.Channel()
	if err != nil {
		return errors.New("channel connection error")
	}
	defer channel.Close()

	notification := notificationmodels.NewNotification()
	if err := notification.MapMessage(data); err != nil {
		return err
	}

	// fetch latest activity for checking actor
	activity, nc, err := notification.FetchLastActivity()
	if err != nil {
		return err
	}

	if !validNotification(activity, notification) {
		return nil
	}

	uc, err := models.FetchUserContact(notification.AccountId)
	if err != nil {
		return fmt.Errorf("an error occurred while fetching user contact: %s", err)
	}

	if !n.checkMailSettings(uc, activity, nc) {
		return nil
	}

	mc := models.NewMailerContainer()
	mc.AccountId = notification.AccountId
	mc.Activity = activity
	mc.Content = nc

	if err := mc.PrepareContainer(); err != nil {
		return err
	}

	mc.CreatedAt = notification.ActivatedAt

	tp := models.NewTemplateParser()
	tp.UserContact = uc
	body, err := tp.RenderInstantTemplate(mc)
	if err != nil {
		return fmt.Errorf("an error occurred while preparing notification email: %s", err)
	}

	tg := &models.TokenGenerator{
		UserContact:      uc,
		NotificationType: emailConfig[nc.TypeConstant],
	}

	if err := tg.CreateToken(); err != nil {
		return err
	}

	mailer := models.NewMailer()
	mailer.EmailSettings = n.settings
	mailer.UserContact = uc
	mailer.Body = body
	mailer.Subject = prepareSubject(mc)

	if err := mailer.SendMail(); err != nil {
		return err
	}

	n.log.Info("%s notified by email", uc.Username)

	return nil
}

func prepareSubject(mc *models.MailerContainer) string {
	t, err := mc.Content.GetContentType()
	if err != nil {
		return ""
	}

	return t.GetDefinition()
}

func validNotification(a *notificationmodels.NotificationActivity, n *notificationmodels.Notification) bool {
	// do not notify actor for her own action
	if a.ActorId == n.AccountId {
		return false
	}

	// do not notify user when notification is not yet activated
	return !n.ActivatedAt.IsZero()
}

func (n *EmailNotifierWorkerController) checkMailSettings(uc *models.UserContact,
	a *notificationmodels.NotificationActivity, nc *notificationmodels.NotificationContent) bool {
	// notifications are disabled
	if val := uc.EmailSettings["global"]; !val {
		return false
	}

	notificationEnabled := uc.EmailSettings[emailConfig[nc.TypeConstant]]
	// daily notifications are enabled
	if val := uc.EmailSettings["daily"]; val {
		if notificationEnabled {
			go n.saveDailyMail(uc.AccountId, a.Id)
		}

		return false
	}

	// get config
	return notificationEnabled
}

func (n *EmailNotifierWorkerController) saveDailyMail(accountId, activityId int64) {
	if err := saveRecipient(accountId); err != nil {
		n.log.Error("daily mail error: %s", err)
	}

	if err := saveActivity(accountId, activityId); err != nil {
		n.log.Error("daily mail error: %s", err)
	}
}

func saveRecipient(accountId int64) error {
	redisConn := helper.MustGetRedisConn()
	key := prepareRecipientsCacheKey()
	if _, err := redisConn.AddSetMembers(key, accountId); err != nil {
		return err
	}

	if err := redisConn.Expire(key, DAY); err != nil {
		return fmt.Errorf("Could not set ttl of recipients: %s", err)
	}

	return nil
}

func saveActivity(accountId, activityId int64) error {
	redisConn := helper.MustGetRedisConn()
	key := prepareSetterCacheKey(accountId)
	if _, err := redisConn.AddSetMembers(key, activityId); err != nil {
		return err
	}

	if err := redisConn.Expire(key, DAY); err != nil {
		return fmt.Errorf("Could not set ttl of activity: %s", err)
	}

	return nil
}

func prepareRecipientsCacheKey() string {
	return fmt.Sprintf("%s:%s:%s:%s",
		config.Get().Environment,
		CACHEPREFIX,
		RECIPIENTSKEY,
		time.Now().Format(TIMEFORMAT))
}

func containsObject(nc *models.NotificationContent) bool {
	return nc.TypeConstant == models.NotificationContent_TYPE_LIKE ||
		nc.TypeConstant == models.NotificationContent_TYPE_MENTION ||
		nc.TypeConstant == models.NotificationContent_TYPE_COMMENT
}

func fetchContentBody(nc *models.NotificationContent, cm *socialmodels.ChannelMessage) string {

	switch nc.TypeConstant {
	case models.NotificationContent_TYPE_LIKE:
		return cm.Body
	case models.NotificationContent_TYPE_MENTION:
		return cm.Body
	case models.NotificationContent_TYPE_COMMENT:
		return fetchLastReplyBody(cm.Id)
	}

	return ""
}

func fetchLastReplyBody(targetId int64) string {
	mr := socialmodels.NewMessageReply()
	mr.MessageId = targetId
	query := socialmodels.NewQuery()
	query.Limit = 1
	messages, err := mr.List(query)
	if err != nil {
		return ""
	}

	if len(messages) == 0 {
		return ""
	}

	return messages[0].Body
}

func fetchRepliedMessage(replyId int64) *socialmodels.ChannelMessage {
	mr := socialmodels.NewMessageReply()
	mr.ReplyId = replyId

	parent, err := mr.FetchRepliedMessage()
	if err != nil {
		parent = socialmodels.NewChannelMessage()
	}

	return parent
}

func (n *Controller) SendMail(uc *UserContact, body, subject string) error {
	es := n.settings
	sg := sendgrid.NewSendGridClient(es.Username, es.Password)
	fullname := fmt.Sprintf("%s %s", uc.FirstName, uc.LastName)

	message := sendgrid.NewMail()
	message.AddTo(uc.Email)
	message.AddToName(fullname)
	message.SetSubject(subject)
	message.SetHTML(body)
	message.SetFrom(es.FromMail)
	message.SetFromName(es.FromName)

	if err := sg.Send(message); err != nil {
		return fmt.Errorf("an error occurred while sending notification email to %s", uc.Username)
	}
	n.log.Info("%s notified by email", uc.Username)

	return nil
}

func prepareSetterCacheKey(accountId int64) string {
	return fmt.Sprintf("%s:%s:%d:%s",
		config.Get().Environment,
		CACHEPREFIX,
		accountId,
		time.Now().Format(TIMEFORMAT))
}
<|endoftext|>


<commit_before>package nessie

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
)

// Nessus implements most of the communication with Nessus.
type Nessus struct {
	// client is the HTTP client to use to issue requests to nessus.
	client *http.Client
	// authCookie is the login token returned by nessus upon successful login.
	authCookie string
	apiURL     string
}

// NewNessus will return a new Nessus initialized with a client matching the security parameters.
// if caCertPath is empty, the host certificate roots will be used to check for the validity of the nessus server API certificate.
func NewNessus(apiURL, caCertPath string, ignoreSSLCertsErrors bool) (*Nessus, error) {
	var roots *x509.CertPool
	if len(caCertPath) != 0 {
		roots = x509.NewCertPool()
		rootPEM, err := ioutil.ReadFile(caCertPath)
		if err != nil {
			return nil, err
		}
		ok := roots.AppendCertsFromPEM(rootPEM)
		if !ok {
			return nil, fmt.Errorf("could not append certs from PEM %s", caCertPath)
		}
	}
	return &Nessus{
		apiURL: apiURL,
		client: &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{
					InsecureSkipVerify: ignoreSSLCertsErrors,
					RootCAs:            roots,
				},
			},
		},
	}, nil
}

func (n *Nessus) doRequest(method string, resource string, data url.Values, wantStatus []int) (resp *http.Response, err error) {
	u, err := url.ParseRequestURI(n.apiURL)
	if err != nil {
		return nil, err
	}
	u.Path = resource
	urlStr := fmt.Sprintf("%v", u)

	req, err := http.NewRequest(method, urlStr, bytes.NewBufferString(data.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded;charset=utf-8")
	req.Header.Add("Accept", "application/json")
	if n.authCookie != "" {
		req.Header.Add("X-Cookie", fmt.Sprintf("token=%s", n.authCookie))
	}

	resp, err = n.client.Do(req)
	if err != nil {
		return nil, err
	}
	var statusFound bool
	for _, status := range wantStatus {
		if resp.StatusCode == status {
			statusFound = true
			break
		}
	}
	if !statusFound {
		return nil, fmt.Errorf("Unexpected status code during login, got %d wanted %s", resp.StatusCode, wantStatus)
	}
	return resp, nil
}

type loginResp struct {
	Token string `json:"token"`
}

// ServerProperties is the structure returned by the ServerProperties() method.
type ServerProperties struct {
	Token           string `json:"token"`
	NessusType      string `json:"nessus_type"`
	NessusUIVersion string `json:"nessus_ui_version"`
	ServerVersion   string `json:"server_version"`
	Feed            string `json:"feed"`
	Enterprise      bool   `json:"enterprise"`
	LoadedPluginSet string `json:"loaded_plugin_set"`
	ServerUUID      string `json:"server_uuid"`
	Expiration      int64  `json:"expiration"`
	Notifications   []struct {
		Type string `json:"type"`
		Msg  string `json:"message"`
	} `json:"notifications"`
	ExpirationTime int64 `json:"expiration_time"`
	Capabilities   struct {
		MultiScanner      bool `json:"multi_scanner"`
		ReportEmailConfig bool `json:"report_email_config"`
	} `json:"capabilities"`
	PluginSet       string `json:"plugin_set"`
	IdleTImeout     int64  `json:"idle_timeout"`
	ScannerBoottime int64  `json:"scanner_boottime"`
	LoginBanner     bool   `json:"login_banner"`
}

// ServerStatus is the stucture returned  by the ServerStatus() method.
type ServerStatus struct {
	Status             string `json:"status"`
	Progress           int64  `json:"progress"`
	MustDestroySession bool
}

// Login will log into nessus with the username and passwords given from the command line flags.
func (n *Nessus) Login(username, password string) error {
	log.Printf("Login into %s\n", n.apiURL)
	data := url.Values{}
	data.Set("username", username)
	data.Set("password", password)

	resp, err := n.doRequest("POST", "/session", data, []int{http.StatusOK})
	if err != nil {
		return err
	}
	reply := &loginResp{}
	if err = json.NewDecoder(resp.Body).Decode(&reply); err != nil {
		return err
	}
	n.authCookie = reply.Token
	return nil
}

// Logout will invalidate the current session token.
func (n *Nessus) Logout() error {
	if n.authCookie == "" {
		log.Println("Not logged in, nothing to do to logout...")
		return nil
	}
	log.Println("Logout...")

	if _, err := n.doRequest("DELETE", "/session", nil, []int{http.StatusOK}); err != nil {
		return err
	}
	n.authCookie = ""
	return nil
}

// ServerProperties will return the current state of the nessus instance.
func (n *Nessus) ServerProperties() (*ServerProperties, error) {
	log.Println("Server properties...")

	resp, err := n.doRequest("GET", "/server/properties", nil, []int{http.StatusOK})
	if err != nil {
		return nil, err
	}
	reply := &ServerProperties{}
	if err = json.NewDecoder(resp.Body).Decode(&reply); err != nil {
		return nil, err
	}
	return reply, nil
}

// ServerStatus will return the current status of the nessus instance.
func (n *Nessus) ServerStatus() (*ServerStatus, error) {
	log.Println("Server status...")

	resp, err := n.doRequest("GET", "/server/status", nil, []int{http.StatusOK, http.StatusServiceUnavailable})
	if err != nil {
		return nil, err
	}
	reply := &ServerStatus{}
	if err = json.NewDecoder(resp.Body).Decode(&reply); err != nil {
		return nil, err
	}
	if resp.StatusCode == http.StatusServiceUnavailable {
		reply.MustDestroySession = true
	}
	return reply, nil
}
<commit_msg>fixed build<commit_after>package nessie

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
)

// Nessus implements most of the communication with Nessus.
type Nessus struct {
	// client is the HTTP client to use to issue requests to nessus.
	client *http.Client
	// authCookie is the login token returned by nessus upon successful login.
	authCookie string
	apiURL     string
}

// NewNessus will return a new Nessus initialized with a client matching the security parameters.
// if caCertPath is empty, the host certificate roots will be used to check for the validity of the nessus server API certificate.
func NewNessus(apiURL, caCertPath string, ignoreSSLCertsErrors bool) (*Nessus, error) {
	var roots *x509.CertPool
	if len(caCertPath) != 0 {
		roots = x509.NewCertPool()
		rootPEM, err := ioutil.ReadFile(caCertPath)
		if err != nil {
			return nil, err
		}
		ok := roots.AppendCertsFromPEM(rootPEM)
		if !ok {
			return nil, fmt.Errorf("could not append certs from PEM %s", caCertPath)
		}
	}
	return &Nessus{
		apiURL: apiURL,
		client: &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{
					InsecureSkipVerify: ignoreSSLCertsErrors,
					RootCAs:            roots,
				},
			},
		},
	}, nil
}

func (n *Nessus) doRequest(method string, resource string, data url.Values, wantStatus []int) (resp *http.Response, err error) {
	u, err := url.ParseRequestURI(n.apiURL)
	if err != nil {
		return nil, err
	}
	u.Path = resource
	urlStr := fmt.Sprintf("%v", u)

	req, err := http.NewRequest(method, urlStr, bytes.NewBufferString(data.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded;charset=utf-8")
	req.Header.Add("Accept", "application/json")
	if n.authCookie != "" {
		req.Header.Add("X-Cookie", fmt.Sprintf("token=%s", n.authCookie))
	}

	resp, err = n.client.Do(req)
	if err != nil {
		return nil, err
	}
	var statusFound bool
	for _, status := range wantStatus {
		if resp.StatusCode == status {
			statusFound = true
			break
		}
	}
	if !statusFound {
		return nil, fmt.Errorf("Unexpected status code during login, got %d wanted %s", resp.StatusCode, wantStatus)
	}
	return resp, nil
}

// Login will log into nessus with the username and passwords given from the command line flags.
func (n *Nessus) Login(username, password string) error {
	log.Printf("Login into %s\n", n.apiURL)
	data := url.Values{}
	data.Set("username", username)
	data.Set("password", password)

	resp, err := n.doRequest("POST", "/session", data, []int{http.StatusOK})
	if err != nil {
		return err
	}
	reply := &loginResp{}
	if err = json.NewDecoder(resp.Body).Decode(&reply); err != nil {
		return err
	}
	n.authCookie = reply.Token
	return nil
}

// Logout will invalidate the current session token.
func (n *Nessus) Logout() error {
	if n.authCookie == "" {
		log.Println("Not logged in, nothing to do to logout...")
		return nil
	}
	log.Println("Logout...")

	if _, err := n.doRequest("DELETE", "/session", nil, []int{http.StatusOK}); err != nil {
		return err
	}
	n.authCookie = ""
	return nil
}

// ServerProperties will return the current state of the nessus instance.
func (n *Nessus) ServerProperties() (*ServerProperties, error) {
	log.Println("Server properties...")

	resp, err := n.doRequest("GET", "/server/properties", nil, []int{http.StatusOK})
	if err != nil {
		return nil, err
	}
	reply := &ServerProperties{}
	if err = json.NewDecoder(resp.Body).Decode(&reply); err != nil {
		return nil, err
	}
	return reply, nil
}

// ServerStatus will return the current status of the nessus instance.
func (n *Nessus) ServerStatus() (*ServerStatus, error) {
	log.Println("Server status...")

	resp, err := n.doRequest("GET", "/server/status", nil, []int{http.StatusOK, http.StatusServiceUnavailable})
	if err != nil {
		return nil, err
	}
	reply := &ServerStatus{}
	if err = json.NewDecoder(resp.Body).Decode(&reply); err != nil {
		return nil, err
	}
	if resp.StatusCode == http.StatusServiceUnavailable {
		reply.MustDestroySession = true
	}
	return reply, nil
}
<|endoftext|>


<commit_before>package vault

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/vault/logical"
)

const (
	// expirationSubPath is the sub-path used for the expiration manager
	// view. This is nested under the system view.
	expirationSubPath = "expire/"

	// maxRevokeAttempts limits how many revoke attempts are made
	maxRevokeAttempts = 6

	// revokeRetryBase is a baseline retry time
	revokeRetryBase = 10 * time.Second

	// minRevokeDelay is used to prevent an instant revoke on restore
	minRevokeDelay = 5 * time.Second
)

// ExpirationManager is used by the Core to manage leases. Secrets
// can provide a lease, meaning that they can be renewed or revoked.
// If a secret is not renewed in timely manner, it may be expired, and
// the ExpirationManager will handle doing automatic revocation.
type ExpirationManager struct {
	router *Router
	view   *BarrierView
	logger *log.Logger

	pending     map[string]*time.Timer
	pendingLock sync.Mutex
}

// NewExpirationManager creates a new ExpirationManager that is backed
// using a given view, and uses the provided router for revocation.
func NewExpirationManager(router *Router, view *BarrierView, logger *log.Logger) *ExpirationManager {
	if logger == nil {
		logger = log.New(os.Stderr, "", log.LstdFlags)
	}
	exp := &ExpirationManager{
		router:  router,
		view:    view,
		logger:  logger,
		pending: make(map[string]*time.Timer),
	}
	return exp
}

// setupExpiration is invoked after we've loaded the mount table to
// initialize the expiration manager
func (c *Core) setupExpiration() error {
	// Create a sub-view
	view := c.systemView.SubView(expirationSubPath)

	// Create the manager
	mgr := NewExpirationManager(c.router, view, c.logger)
	c.expiration = mgr

	// Restore the existing state
	if err := c.expiration.Restore(); err != nil {
		return fmt.Errorf("expiration state restore failed: %v", err)
	}
	return nil
}

// stopExpiration is used to stop the expiration manager before
// sealing the Vault.
func (c *Core) stopExpiration() error {
	if err := c.expiration.Stop(); err != nil {
		return err
	}
	c.expiration = nil
	return nil
}

// Restore is used to recover the lease states when starting.
// This is used after starting the vault.
func (m *ExpirationManager) Restore() error {
	m.pendingLock.Lock()
	defer m.pendingLock.Unlock()

	// Accumulate existing leases
	existing, err := CollectKeys(m.view)
	if err != nil {
		return fmt.Errorf("failed to scan for leases: %v", err)
	}

	// Restore each key
	for _, vaultID := range existing {
		// Load the entry
		le, err := m.loadEntry(vaultID)
		if err != nil {
			return err
		}

		// If there is no entry, nothing to restore
		if le == nil {
			continue
		}

		// Determine the remaining time to expiration
		expires := le.ExpireTime.Sub(time.Now().UTC())
		if expires <= 0 {
			expires = minRevokeDelay
		}

		// Setup revocation timer
		m.pending[le.VaultID] = time.AfterFunc(expires, func() {
			m.expireID(le.VaultID)
		})
	}
	m.logger.Printf("[INFO] expire: restored %d leases", len(m.pending))
	return nil
}

// Stop is used to prevent further automatic revocations.
// This must be called before sealing the view.
func (m *ExpirationManager) Stop() error {
	// Stop all the pending expiration timers
	m.pendingLock.Lock()
	for _, timer := range m.pending {
		timer.Stop()
	}
	m.pending = make(map[string]*time.Timer)
	m.pendingLock.Unlock()
	return nil
}

// Revoke is used to revoke a secret named by the given vaultID
func (m *ExpirationManager) Revoke(vaultID string) error {
	// Load the entry
	le, err := m.loadEntry(vaultID)
	if err != nil {
		return err
	}

	// If there is no entry, nothing to revoke
	if le == nil {
		return nil
	}

	// Revoke the entry
	if err := m.revokeEntry(le); err != nil {
		return err
	}

	// Delete the entry
	if err := m.deleteEntry(vaultID); err != nil {
		return err
	}

	// Clear the expiration handler
	m.pendingLock.Lock()
	if timer, ok := m.pending[vaultID]; ok {
		timer.Stop()
		delete(m.pending, vaultID)
	}
	m.pendingLock.Unlock()
	return nil
}

// RevokePrefix is used to revoke all secrets with a given prefix.
// The prefix maps to that of the mount table to make this simpler
// to reason about.
func (m *ExpirationManager) RevokePrefix(prefix string) error {
	// Ensure there is a trailing slash
	if !strings.HasSuffix(prefix, "/") {
		prefix = prefix + "/"
	}

	// Accumulate existing leases
	sub := m.view.SubView(prefix)
	existing, err := CollectKeys(sub)
	if err != nil {
		return fmt.Errorf("failed to scan for leases: %v", err)
	}

	// Revoke all the keys
	for idx, suffix := range existing {
		vaultID := prefix + suffix
		if err := m.Revoke(vaultID); err != nil {
			return fmt.Errorf("failed to revoke '%s' (%d / %d): %v",
				vaultID, idx+1, len(existing), err)
		}
	}
	return nil
}

// Renew is used to renew a secret using the given vaultID
// and a renew interval. The increment may be ignored.
func (m *ExpirationManager) Renew(vaultID string, increment time.Duration) (*logical.Response, error) {
	// Load the entry
	le, err := m.loadEntry(vaultID)
	if err != nil {
		return nil, err
	}

	// If there is no entry, cannot review
	if le == nil {
		return nil, fmt.Errorf("lease not found")
	}

	// Determine if the lease is expired
	if le.ExpireTime.Before(time.Now().UTC()) {
		return nil, fmt.Errorf("lease expired")
	}

	// Attempt to renew the entry
	resp, err := m.renewEntry(le, increment)
	if err != nil {
		return nil, err
	}

	// Fast-path if there is no lease
	if resp == nil || resp.Secret == nil || resp.Secret.Lease == 0 {
		return resp, nil
	}

	// Validate the lease
	if err := resp.Secret.Validate(); err != nil {
		return nil, err
	}

	// Attach the VaultID
	resp.Secret.VaultID = vaultID

	// Update the lease entry
	leaseTotal := resp.Secret.Lease + resp.Secret.LeaseGracePeriod
	le.Data = resp.Data
	le.Secret = resp.Secret
	le.ExpireTime = time.Now().UTC().Add(leaseTotal)
	if err := m.persistEntry(le); err != nil {
		return nil, err
	}

	// Update the expiration time
	m.pendingLock.Lock()
	if timer, ok := m.pending[vaultID]; ok {
		timer.Reset(leaseTotal)
	}
	m.pendingLock.Unlock()

	// Return the response
	return resp, nil
}

// Register is used to take a request and response with an associated
// lease. The secret gets assigned a vaultId and the management of
// of lease is assumed by the expiration manager.
func (m *ExpirationManager) Register(req *logical.Request, resp *logical.Response) (string, error) {
	// Ignore if there is no leased secret
	if resp == nil || resp.Secret == nil || resp.Secret.Lease == 0 {
		return "", nil
	}

	// Validate the secret
	if err := resp.Secret.Validate(); err != nil {
		return "", err
	}

	// Create a lease entry
	now := time.Now().UTC()
	leaseTotal := resp.Secret.Lease + resp.Secret.LeaseGracePeriod
	le := leaseEntry{
		VaultID:    path.Join(req.Path, generateUUID()),
		Path:       req.Path,
		Data:       resp.Data,
		Secret:     resp.Secret,
		IssueTime:  now,
		ExpireTime: now.Add(leaseTotal),
	}

	// Encode the entry
	if err := m.persistEntry(&le); err != nil {
		return "", err
	}

	// Setup revocation timer
	m.pendingLock.Lock()
	m.pending[le.VaultID] = time.AfterFunc(leaseTotal, func() {
		m.expireID(le.VaultID)
	})
	m.pendingLock.Unlock()

	// Done
	return le.VaultID, nil
}

// expireID is invoked when a given ID is expired
func (m *ExpirationManager) expireID(vaultID string) {
	// Clear from the pending expiration
	m.pendingLock.Lock()
	delete(m.pending, vaultID)
	m.pendingLock.Unlock()

	for attempt := uint(0); attempt < maxRevokeAttempts; attempt++ {
		err := m.Revoke(vaultID)
		if err == nil {
			m.logger.Printf("[INFO] expire: revoked '%s'", vaultID)
			return
		}
		m.logger.Printf("[ERR] expire: failed to revoke '%s': %v", vaultID, err)
		time.Sleep((1 << attempt) * revokeRetryBase)
	}
	m.logger.Printf("[ERR] expire: maximum revoke attempts for '%s' reached", vaultID)
}

// revokeEntry is used to attempt revocation of an internal entry
func (m *ExpirationManager) revokeEntry(le *leaseEntry) error {
	_, err := m.router.Route(logical.RevokeRequest(
		le.Path, le.Secret, le.Data))
	if err != nil {
		return fmt.Errorf("failed to revoke entry: %v", err)
	}
	return nil
}

// renewEntry is used to attempt renew of an internal entry
func (m *ExpirationManager) renewEntry(le *leaseEntry, increment time.Duration) (*logical.Response, error) {
	secret := *le.Secret
	secret.LeaseIncrement = increment
	secret.VaultID = ""

	resp, err := m.router.Route(logical.RenewRequest(
		le.Path, &secret, le.Data))
	if err != nil {
		return nil, fmt.Errorf("failed to renew entry: %v", err)
	}
	return resp, nil
}

// loadEntry is used to read a lease entry
func (m *ExpirationManager) loadEntry(vaultID string) (*leaseEntry, error) {
	out, err := m.view.Get(vaultID)
	if err != nil {
		return nil, fmt.Errorf("failed to read lease entry: %v", err)
	}
	if out == nil {
		return nil, nil
	}
	le, err := decodeLeaseEntry(out.Value)
	if err != nil {
		return nil, fmt.Errorf("failed to decode lease entry: %v", err)
	}
	return le, nil
}

// persistEntry is used to persist a lease entry
func (m *ExpirationManager) persistEntry(le *leaseEntry) error {
	// Encode the entry
	buf, err := le.encode()
	if err != nil {
		return fmt.Errorf("failed to encode lease entry: %v", err)
	}

	// Write out to the view
	ent := logical.StorageEntry{
		Key:   le.VaultID,
		Value: buf,
	}
	if err := m.view.Put(&ent); err != nil {
		return fmt.Errorf("failed to persist lease entry: %v", err)
	}
	return nil
}

// deleteEntry is used to delete a lease entry
func (m *ExpirationManager) deleteEntry(vaultID string) error {
	if err := m.view.Delete(vaultID); err != nil {
		return fmt.Errorf("failed to delete lease entry: %v", err)
	}
	return nil
}

// leaseEntry is used to structure the values the expiration
// manager stores. This is used to handle renew and revocation.
type leaseEntry struct {
	VaultID    string                 `json:"vault_id"`
	Path       string                 `json:"path"`
	Data       map[string]interface{} `json:"data"`
	Secret     *logical.Secret        `json:"secret"`
	IssueTime  time.Time              `json:"issue_time"`
	ExpireTime time.Time              `json:"expire_time"`
}

// encode is used to JSON encode the lease entry
func (l *leaseEntry) encode() ([]byte, error) {
	return json.Marshal(l)
}

// decodeLeaseEntry is used to reverse encode and return a new entry
func decodeLeaseEntry(buf []byte) (*leaseEntry, error) {
	out := new(leaseEntry)
	return out, json.Unmarshal(buf, out)
}
<commit_msg>vault: only log expiration notice if useful<commit_after>package vault

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/vault/logical"
)

const (
	// expirationSubPath is the sub-path used for the expiration manager
	// view. This is nested under the system view.
	expirationSubPath = "expire/"

	// maxRevokeAttempts limits how many revoke attempts are made
	maxRevokeAttempts = 6

	// revokeRetryBase is a baseline retry time
	revokeRetryBase = 10 * time.Second

	// minRevokeDelay is used to prevent an instant revoke on restore
	minRevokeDelay = 5 * time.Second
)

// ExpirationManager is used by the Core to manage leases. Secrets
// can provide a lease, meaning that they can be renewed or revoked.
// If a secret is not renewed in timely manner, it may be expired, and
// the ExpirationManager will handle doing automatic revocation.
type ExpirationManager struct {
	router *Router
	view   *BarrierView
	logger *log.Logger

	pending     map[string]*time.Timer
	pendingLock sync.Mutex
}

// NewExpirationManager creates a new ExpirationManager that is backed
// using a given view, and uses the provided router for revocation.
func NewExpirationManager(router *Router, view *BarrierView, logger *log.Logger) *ExpirationManager {
	if logger == nil {
		logger = log.New(os.Stderr, "", log.LstdFlags)
	}
	exp := &ExpirationManager{
		router:  router,
		view:    view,
		logger:  logger,
		pending: make(map[string]*time.Timer),
	}
	return exp
}

// setupExpiration is invoked after we've loaded the mount table to
// initialize the expiration manager
func (c *Core) setupExpiration() error {
	// Create a sub-view
	view := c.systemView.SubView(expirationSubPath)

	// Create the manager
	mgr := NewExpirationManager(c.router, view, c.logger)
	c.expiration = mgr

	// Restore the existing state
	if err := c.expiration.Restore(); err != nil {
		return fmt.Errorf("expiration state restore failed: %v", err)
	}
	return nil
}

// stopExpiration is used to stop the expiration manager before
// sealing the Vault.
func (c *Core) stopExpiration() error {
	if err := c.expiration.Stop(); err != nil {
		return err
	}
	c.expiration = nil
	return nil
}

// Restore is used to recover the lease states when starting.
// This is used after starting the vault.
func (m *ExpirationManager) Restore() error {
	m.pendingLock.Lock()
	defer m.pendingLock.Unlock()

	// Accumulate existing leases
	existing, err := CollectKeys(m.view)
	if err != nil {
		return fmt.Errorf("failed to scan for leases: %v", err)
	}

	// Restore each key
	for _, vaultID := range existing {
		// Load the entry
		le, err := m.loadEntry(vaultID)
		if err != nil {
			return err
		}

		// If there is no entry, nothing to restore
		if le == nil {
			continue
		}

		// Determine the remaining time to expiration
		expires := le.ExpireTime.Sub(time.Now().UTC())
		if expires <= 0 {
			expires = minRevokeDelay
		}

		// Setup revocation timer
		m.pending[le.VaultID] = time.AfterFunc(expires, func() {
			m.expireID(le.VaultID)
		})
	}
	if len(m.pending) > 0 {
		m.logger.Printf("[INFO] expire: restored %d leases", len(m.pending))
	}
	return nil
}

// Stop is used to prevent further automatic revocations.
// This must be called before sealing the view.
func (m *ExpirationManager) Stop() error {
	// Stop all the pending expiration timers
	m.pendingLock.Lock()
	for _, timer := range m.pending {
		timer.Stop()
	}
	m.pending = make(map[string]*time.Timer)
	m.pendingLock.Unlock()
	return nil
}

// Revoke is used to revoke a secret named by the given vaultID
func (m *ExpirationManager) Revoke(vaultID string) error {
	// Load the entry
	le, err := m.loadEntry(vaultID)
	if err != nil {
		return err
	}

	// If there is no entry, nothing to revoke
	if le == nil {
		return nil
	}

	// Revoke the entry
	if err := m.revokeEntry(le); err != nil {
		return err
	}

	// Delete the entry
	if err := m.deleteEntry(vaultID); err != nil {
		return err
	}

	// Clear the expiration handler
	m.pendingLock.Lock()
	if timer, ok := m.pending[vaultID]; ok {
		timer.Stop()
		delete(m.pending, vaultID)
	}
	m.pendingLock.Unlock()
	return nil
}

// RevokePrefix is used to revoke all secrets with a given prefix.
// The prefix maps to that of the mount table to make this simpler
// to reason about.
func (m *ExpirationManager) RevokePrefix(prefix string) error {
	// Ensure there is a trailing slash
	if !strings.HasSuffix(prefix, "/") {
		prefix = prefix + "/"
	}

	// Accumulate existing leases
	sub := m.view.SubView(prefix)
	existing, err := CollectKeys(sub)
	if err != nil {
		return fmt.Errorf("failed to scan for leases: %v", err)
	}

	// Revoke all the keys
	for idx, suffix := range existing {
		vaultID := prefix + suffix
		if err := m.Revoke(vaultID); err != nil {
			return fmt.Errorf("failed to revoke '%s' (%d / %d): %v",
				vaultID, idx+1, len(existing), err)
		}
	}
	return nil
}

// Renew is used to renew a secret using the given vaultID
// and a renew interval. The increment may be ignored.
func (m *ExpirationManager) Renew(vaultID string, increment time.Duration) (*logical.Response, error) {
	// Load the entry
	le, err := m.loadEntry(vaultID)
	if err != nil {
		return nil, err
	}

	// If there is no entry, cannot review
	if le == nil {
		return nil, fmt.Errorf("lease not found")
	}

	// Determine if the lease is expired
	if le.ExpireTime.Before(time.Now().UTC()) {
		return nil, fmt.Errorf("lease expired")
	}

	// Attempt to renew the entry
	resp, err := m.renewEntry(le, increment)
	if err != nil {
		return nil, err
	}

	// Fast-path if there is no lease
	if resp == nil || resp.Secret == nil || resp.Secret.Lease == 0 {
		return resp, nil
	}

	// Validate the lease
	if err := resp.Secret.Validate(); err != nil {
		return nil, err
	}

	// Attach the VaultID
	resp.Secret.VaultID = vaultID

	// Update the lease entry
	leaseTotal := resp.Secret.Lease + resp.Secret.LeaseGracePeriod
	le.Data = resp.Data
	le.Secret = resp.Secret
	le.ExpireTime = time.Now().UTC().Add(leaseTotal)
	if err := m.persistEntry(le); err != nil {
		return nil, err
	}

	// Update the expiration time
	m.pendingLock.Lock()
	if timer, ok := m.pending[vaultID]; ok {
		timer.Reset(leaseTotal)
	}
	m.pendingLock.Unlock()

	// Return the response
	return resp, nil
}

// Register is used to take a request and response with an associated
// lease. The secret gets assigned a vaultId and the management of
// of lease is assumed by the expiration manager.
func (m *ExpirationManager) Register(req *logical.Request, resp *logical.Response) (string, error) {
	// Ignore if there is no leased secret
	if resp == nil || resp.Secret == nil || resp.Secret.Lease == 0 {
		return "", nil
	}

	// Validate the secret
	if err := resp.Secret.Validate(); err != nil {
		return "", err
	}

	// Create a lease entry
	now := time.Now().UTC()
	leaseTotal := resp.Secret.Lease + resp.Secret.LeaseGracePeriod
	le := leaseEntry{
		VaultID:    path.Join(req.Path, generateUUID()),
		Path:       req.Path,
		Data:       resp.Data,
		Secret:     resp.Secret,
		IssueTime:  now,
		ExpireTime: now.Add(leaseTotal),
	}

	// Encode the entry
	if err := m.persistEntry(&le); err != nil {
		return "", err
	}

	// Setup revocation timer
	m.pendingLock.Lock()
	m.pending[le.VaultID] = time.AfterFunc(leaseTotal, func() {
		m.expireID(le.VaultID)
	})
	m.pendingLock.Unlock()

	// Done
	return le.VaultID, nil
}

// expireID is invoked when a given ID is expired
func (m *ExpirationManager) expireID(vaultID string) {
	// Clear from the pending expiration
	m.pendingLock.Lock()
	delete(m.pending, vaultID)
	m.pendingLock.Unlock()

	for attempt := uint(0); attempt < maxRevokeAttempts; attempt++ {
		err := m.Revoke(vaultID)
		if err == nil {
			m.logger.Printf("[INFO] expire: revoked '%s'", vaultID)
			return
		}
		m.logger.Printf("[ERR] expire: failed to revoke '%s': %v", vaultID, err)
		time.Sleep((1 << attempt) * revokeRetryBase)
	}
	m.logger.Printf("[ERR] expire: maximum revoke attempts for '%s' reached", vaultID)
}

// revokeEntry is used to attempt revocation of an internal entry
func (m *ExpirationManager) revokeEntry(le *leaseEntry) error {
	_, err := m.router.Route(logical.RevokeRequest(
		le.Path, le.Secret, le.Data))
	if err != nil {
		return fmt.Errorf("failed to revoke entry: %v", err)
	}
	return nil
}

// renewEntry is used to attempt renew of an internal entry
func (m *ExpirationManager) renewEntry(le *leaseEntry, increment time.Duration) (*logical.Response, error) {
	secret := *le.Secret
	secret.LeaseIncrement = increment
	secret.VaultID = ""

	resp, err := m.router.Route(logical.RenewRequest(
		le.Path, &secret, le.Data))
	if err != nil {
		return nil, fmt.Errorf("failed to renew entry: %v", err)
	}
	return resp, nil
}

// loadEntry is used to read a lease entry
func (m *ExpirationManager) loadEntry(vaultID string) (*leaseEntry, error) {
	out, err := m.view.Get(vaultID)
	if err != nil {
		return nil, fmt.Errorf("failed to read lease entry: %v", err)
	}
	if out == nil {
		return nil, nil
	}
	le, err := decodeLeaseEntry(out.Value)
	if err != nil {
		return nil, fmt.Errorf("failed to decode lease entry: %v", err)
	}
	return le, nil
}

// persistEntry is used to persist a lease entry
func (m *ExpirationManager) persistEntry(le *leaseEntry) error {
	// Encode the entry
	buf, err := le.encode()
	if err != nil {
		return fmt.Errorf("failed to encode lease entry: %v", err)
	}

	// Write out to the view
	ent := logical.StorageEntry{
		Key:   le.VaultID,
		Value: buf,
	}
	if err := m.view.Put(&ent); err != nil {
		return fmt.Errorf("failed to persist lease entry: %v", err)
	}
	return nil
}

// deleteEntry is used to delete a lease entry
func (m *ExpirationManager) deleteEntry(vaultID string) error {
	if err := m.view.Delete(vaultID); err != nil {
		return fmt.Errorf("failed to delete lease entry: %v", err)
	}
	return nil
}

// leaseEntry is used to structure the values the expiration
// manager stores. This is used to handle renew and revocation.
type leaseEntry struct {
	VaultID    string                 `json:"vault_id"`
	Path       string                 `json:"path"`
	Data       map[string]interface{} `json:"data"`
	Secret     *logical.Secret        `json:"secret"`
	IssueTime  time.Time              `json:"issue_time"`
	ExpireTime time.Time              `json:"expire_time"`
}

// encode is used to JSON encode the lease entry
func (l *leaseEntry) encode() ([]byte, error) {
	return json.Marshal(l)
}

// decodeLeaseEntry is used to reverse encode and return a new entry
func decodeLeaseEntry(buf []byte) (*leaseEntry, error) {
	out := new(leaseEntry)
	return out, json.Unmarshal(buf, out)
}
<|endoftext|>


<commit_before>// Copyright 2020 The gVisor Authors.
//
// 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 udp_recv_mcast_bcast_test

import (
	"flag"
	"net"
	"testing"

	"golang.org/x/sys/unix"
	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/test/packetimpact/testbench"
)

func init() {
	testbench.RegisterFlags(flag.CommandLine)
}

func TestUDPRecvMulticastBroadcast(t *testing.T) {
	dut := testbench.NewDUT(t)
	defer dut.TearDown()
	boundFD, remotePort := dut.CreateBoundSocket(t, unix.SOCK_DGRAM, unix.IPPROTO_UDP, net.IPv4(0, 0, 0, 0))
	defer dut.Close(t, boundFD)
	conn := testbench.NewUDPIPv4(t, testbench.UDP{DstPort: &remotePort}, testbench.UDP{SrcPort: &remotePort})
	defer conn.Close(t)

	for _, bcastAddr := range []net.IP{
		broadcastAddr(net.ParseIP(testbench.RemoteIPv4), net.CIDRMask(testbench.IPv4PrefixLength, 32)),
		net.IPv4(255, 255, 255, 255),
		net.IPv4(224, 0, 0, 1),
	} {
		payload := testbench.GenerateRandomPayload(t, 1<<10)
		conn.SendIP(
			t,
			testbench.IPv4{DstAddr: testbench.Address(tcpip.Address(bcastAddr.To4()))},
			testbench.UDP{},
			&testbench.Payload{Bytes: payload},
		)
		t.Logf("Receiving packet sent to address: %s", bcastAddr)
		if got, want := string(dut.Recv(t, boundFD, int32(len(payload)), 0)), string(payload); got != want {
			t.Errorf("received payload does not match sent payload got: %s, want: %s", got, want)
		}
	}
}

func broadcastAddr(ip net.IP, mask net.IPMask) net.IP {
	ip4 := ip.To4()
	for i := range ip4 {
		ip4[i] |= ^mask[i]
	}
	return ip4
}
<commit_msg>More test cases on receiving UDP mcast/bcast<commit_after>// Copyright 2020 The gVisor Authors.
//
// 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 udp_recv_mcast_bcast_test

import (
	"context"
	"flag"
	"fmt"
	"net"
	"syscall"
	"testing"

	"golang.org/x/sys/unix"
	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/test/packetimpact/testbench"
)

func init() {
	testbench.RegisterFlags(flag.CommandLine)
}

func TestUDPRecvMcastBcast(t *testing.T) {
	subnetBcastAddr := broadcastAddr(net.ParseIP(testbench.RemoteIPv4), net.CIDRMask(testbench.IPv4PrefixLength, 32))

	for _, v := range []struct {
		bound, to net.IP
	}{
		{bound: net.IPv4(0, 0, 0, 0), to: subnetBcastAddr},
		{bound: net.IPv4(0, 0, 0, 0), to: net.IPv4bcast},
		{bound: net.IPv4(0, 0, 0, 0), to: net.IPv4allsys},

		{bound: subnetBcastAddr, to: subnetBcastAddr},
		{bound: subnetBcastAddr, to: net.IPv4bcast},

		{bound: net.IPv4bcast, to: net.IPv4bcast},
		{bound: net.IPv4allsys, to: net.IPv4allsys},
	} {
		t.Run(fmt.Sprintf("bound=%s,to=%s", v.bound, v.to), func(t *testing.T) {
			dut := testbench.NewDUT(t)
			defer dut.TearDown()
			boundFD, remotePort := dut.CreateBoundSocket(t, unix.SOCK_DGRAM, unix.IPPROTO_UDP, v.bound)
			defer dut.Close(t, boundFD)
			conn := testbench.NewUDPIPv4(t, testbench.UDP{DstPort: &remotePort}, testbench.UDP{SrcPort: &remotePort})
			defer conn.Close(t)

			payload := testbench.GenerateRandomPayload(t, 1<<10)
			conn.SendIP(
				t,
				testbench.IPv4{DstAddr: testbench.Address(tcpip.Address(v.to.To4()))},
				testbench.UDP{},
				&testbench.Payload{Bytes: payload},
			)
			if got, want := string(dut.Recv(t, boundFD, int32(len(payload)), 0)), string(payload); got != want {
				t.Errorf("received payload does not match sent payload got: %s, want: %s", got, want)
			}
		})
	}
}

func TestUDPDoesntRecvMcastBcastOnUnicastAddr(t *testing.T) {
	dut := testbench.NewDUT(t)
	defer dut.TearDown()
	boundFD, remotePort := dut.CreateBoundSocket(t, unix.SOCK_DGRAM, unix.IPPROTO_UDP, net.ParseIP(testbench.RemoteIPv4))
	dut.SetSockOptTimeval(t, boundFD, unix.SOL_SOCKET, unix.SO_RCVTIMEO, &unix.Timeval{Sec: 1, Usec: 0})
	defer dut.Close(t, boundFD)
	conn := testbench.NewUDPIPv4(t, testbench.UDP{DstPort: &remotePort}, testbench.UDP{SrcPort: &remotePort})
	defer conn.Close(t)

	for _, to := range []net.IP{
		broadcastAddr(net.ParseIP(testbench.RemoteIPv4), net.CIDRMask(testbench.IPv4PrefixLength, 32)),
		net.IPv4(255, 255, 255, 255),
		net.IPv4(224, 0, 0, 1),
	} {
		t.Run(fmt.Sprint("to=%s", to), func(t *testing.T) {
			payload := testbench.GenerateRandomPayload(t, 1<<10)
			conn.SendIP(
				t,
				testbench.IPv4{DstAddr: testbench.Address(tcpip.Address(to.To4()))},
				testbench.UDP{},
				&testbench.Payload{Bytes: payload},
			)
			ret, payload, errno := dut.RecvWithErrno(context.Background(), t, boundFD, 100, 0)
			if errno != syscall.EAGAIN || errno != syscall.EWOULDBLOCK {
				t.Errorf("Recv got unexpected result, ret=%d, payload=%q, errno=%s", ret, payload, errno)
			}
		})
	}
}

func broadcastAddr(ip net.IP, mask net.IPMask) net.IP {
	ip4 := ip.To4()
	for i := range ip4 {
		ip4[i] |= ^mask[i]
	}
	return ip4
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"

	"github.com/go-check/check"
)

func (s *DockerSuite) TestExecResizeApiHeightWidthNoInt(c *check.C) {
	testRequires(c, DaemonIsLinux)
	out, _ := dockerCmd(c, "run", "-d", "busybox", "top")
	cleanedContainerID := strings.TrimSpace(out)

	endpoint := "/exec/" + cleanedContainerID + "/resize?h=foo&w=bar"
	status, _, err := sockRequest("POST", endpoint, nil)
	c.Assert(status, check.Equals, http.StatusInternalServerError)
	c.Assert(err, check.IsNil)
}

// Part of #14845
func (s *DockerSuite) TestExecResizeImmediatelyAfterExecStart(c *check.C) {
	testRequires(c, DaemonIsLinux)
	testRequires(c, NativeExecDriver)

	name := "exec_resize_test"
	dockerCmd(c, "run", "-d", "-i", "-t", "--name", name, "--restart", "always", "busybox", "/bin/sh")

	testExecResize := func() error {
		data := map[string]interface{}{
			"AttachStdin": true,
			"Cmd":         []string{"/bin/sh"},
		}
		uri := fmt.Sprintf("/containers/%s/exec", name)
		status, body, err := sockRequest("POST", uri, data)
		if err != nil {
			return err
		}
		if status != http.StatusCreated {
			return fmt.Errorf("POST %s is expected to return %d, got %d", uri, http.StatusCreated, status)
		}

		out := map[string]string{}
		err = json.Unmarshal(body, &out)
		if err != nil {
			return fmt.Errorf("ExecCreate returned invalid json. Error: %q", err.Error())
		}

		execID := out["Id"]
		if len(execID) < 1 {
			return fmt.Errorf("ExecCreate got invalid execID")
		}

		payload := bytes.NewBufferString(`{"Tty":true}`)
		conn, _, err := sockRequestHijack("POST", fmt.Sprintf("/exec/%s/start", execID), payload, "application/json")
		if err != nil {
			return fmt.Errorf("Failed to start the exec: %q", err.Error())
		}
		defer conn.Close()

		_, rc, err := sockRequestRaw("POST", fmt.Sprintf("/exec/%s/resize?h=24&w=80", execID), nil, "text/plain")
		// It's probably a panic of the daemon if io.ErrUnexpectedEOF is returned.
		if err == io.ErrUnexpectedEOF {
			return fmt.Errorf("The daemon might have crashed.")
		}

		if err == nil {
			rc.Close()
		}

		// We only interested in the io.ErrUnexpectedEOF error, so we return nil otherwise.
		return nil
	}

	// The panic happens when daemon.ContainerExecStart is called but the
	// container.Exec is not called.
	// Because the panic is not 100% reproducible, we send the requests concurrently
	// to increase the probability that the problem is triggered.
	var (
		n  = 10
		ch = make(chan error, n)
		wg sync.WaitGroup
	)
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := testExecResize(); err != nil {
				ch <- err
			}
		}()
	}

	wg.Wait()
	select {
	case err := <-ch:
		c.Fatal(err.Error())
	default:
	}
}
<commit_msg>cleaned up integration-cli/docker_api_exec_resize_test.go<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"

	"github.com/docker/docker/pkg/integration/checker"
	"github.com/go-check/check"
)

func (s *DockerSuite) TestExecResizeApiHeightWidthNoInt(c *check.C) {
	testRequires(c, DaemonIsLinux)
	out, _ := dockerCmd(c, "run", "-d", "busybox", "top")
	cleanedContainerID := strings.TrimSpace(out)

	endpoint := "/exec/" + cleanedContainerID + "/resize?h=foo&w=bar"
	status, _, err := sockRequest("POST", endpoint, nil)
	c.Assert(err, checker.IsNil)
	c.Assert(status, checker.Equals, http.StatusInternalServerError)
}

// Part of #14845
func (s *DockerSuite) TestExecResizeImmediatelyAfterExecStart(c *check.C) {
	testRequires(c, DaemonIsLinux)
	testRequires(c, NativeExecDriver)

	name := "exec_resize_test"
	dockerCmd(c, "run", "-d", "-i", "-t", "--name", name, "--restart", "always", "busybox", "/bin/sh")

	testExecResize := func() error {
		data := map[string]interface{}{
			"AttachStdin": true,
			"Cmd":         []string{"/bin/sh"},
		}
		uri := fmt.Sprintf("/containers/%s/exec", name)
		status, body, err := sockRequest("POST", uri, data)
		if err != nil {
			return err
		}
		if status != http.StatusCreated {
			return fmt.Errorf("POST %s is expected to return %d, got %d", uri, http.StatusCreated, status)
		}

		out := map[string]string{}
		err = json.Unmarshal(body, &out)
		if err != nil {
			return fmt.Errorf("ExecCreate returned invalid json. Error: %q", err.Error())
		}

		execID := out["Id"]
		if len(execID) < 1 {
			return fmt.Errorf("ExecCreate got invalid execID")
		}

		payload := bytes.NewBufferString(`{"Tty":true}`)
		conn, _, err := sockRequestHijack("POST", fmt.Sprintf("/exec/%s/start", execID), payload, "application/json")
		if err != nil {
			return fmt.Errorf("Failed to start the exec: %q", err.Error())
		}
		defer conn.Close()

		_, rc, err := sockRequestRaw("POST", fmt.Sprintf("/exec/%s/resize?h=24&w=80", execID), nil, "text/plain")
		// It's probably a panic of the daemon if io.ErrUnexpectedEOF is returned.
		if err == io.ErrUnexpectedEOF {
			return fmt.Errorf("The daemon might have crashed.")
		}

		if err == nil {
			rc.Close()
		}

		// We only interested in the io.ErrUnexpectedEOF error, so we return nil otherwise.
		return nil
	}

	// The panic happens when daemon.ContainerExecStart is called but the
	// container.Exec is not called.
	// Because the panic is not 100% reproducible, we send the requests concurrently
	// to increase the probability that the problem is triggered.
	var (
		n  = 10
		ch = make(chan error, n)
		wg sync.WaitGroup
	)
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := testExecResize(); err != nil {
				ch <- err
			}
		}()
	}

	wg.Wait()
	select {
	case err := <-ch:
		c.Fatal(err.Error())
	default:
	}
}
<|endoftext|>