<commit_before>package dockertest

import (
	"fmt"
	"net"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

// A Container is a container inside docker
type Container struct {
	Name string
	Args []string
	Addr string
	cmd  *exec.Cmd
}

// Shutdown ends the container
func (c *Container) Shutdown() {
	c.cmd.Process.Signal(syscall.SIGINT)
	//  Wait till the process exits.
	c.cmd.Wait()
}

// RunContainer runs a given docker container and returns a port on which the
// container can be reached
func RunContainer(container string, port string, waitFunc func(addr string) error, args ...string) (*Container, error) {
	free := freePort()
	host := getHost()
	addr := fmt.Sprintf("%s:%d", host, free)
	argsFull := append(args, "run", "-p", fmt.Sprintf("%d:%s", free, port), container)
	cmd := exec.Command("docker", argsFull...)
	// run this in the background

	err := cmd.Start()
	if err != nil {
		return nil, fmt.Errorf("could not run container, %s", err)
	}
	for {
		err := waitFunc(addr)
		if err == nil {
			break
		}

		time.Sleep(time.Millisecond * 150)
	}

	return &Container{
		Name: container,
		Addr: addr,
		Args: args,
		cmd:  cmd,
	}, nil
}

func getHost() string {
	out, err := exec.Command("docker-machine", "ip", os.Getenv("DOCKER_MACHINE_NAME")).Output()
	if err == nil {
		return strings.TrimSpace(string(out[:]))
	}
	return "localhost"
}

func freePort() int {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		panic(err)
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		panic(err)
	}
	defer l.Close()

	return l.Addr().(*net.TCPAddr).Port
}
<commit_msg>make sure run is the first docker arg<commit_after>package dockertest

import (
	"fmt"
	"net"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

// A Container is a container inside docker
type Container struct {
	Name string
	Args []string
	Addr string
	cmd  *exec.Cmd
}

// Shutdown ends the container
func (c *Container) Shutdown() {
	c.cmd.Process.Signal(syscall.SIGINT)
	//  Wait till the process exits.
	c.cmd.Wait()
}

// RunContainer runs a given docker container and returns a port on which the
// container can be reached
func RunContainer(container string, port string, waitFunc func(addr string) error, args ...string) (*Container, error) {
	free := freePort()
	host := getHost()
	addr := fmt.Sprintf("%s:%d", host, free)
	argsFull := append([]string{"run"}, args...)
	argsFull = append(argsFull, "-p", fmt.Sprintf("%d:%s", free, port), container)
	cmd := exec.Command("docker", argsFull...)
	// run this in the background

	err := cmd.Start()
	if err != nil {
		return nil, fmt.Errorf("could not run container, %s", err)
	}
	for {
		err := waitFunc(addr)
		if err == nil {
			break
		}

		time.Sleep(time.Millisecond * 150)
	}

	return &Container{
		Name: container,
		Addr: addr,
		Args: args,
		cmd:  cmd,
	}, nil
}

func getHost() string {
	out, err := exec.Command("docker-machine", "ip", os.Getenv("DOCKER_MACHINE_NAME")).Output()
	if err == nil {
		return strings.TrimSpace(string(out[:]))
	}
	return "localhost"
}

func freePort() int {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		panic(err)
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		panic(err)
	}
	defer l.Close()

	return l.Addr().(*net.TCPAddr).Port
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"io/ioutil"
	"net/http"
	"regexp"
	"strings"

	"github.com/pkg/errors"

	log "github.com/Sirupsen/logrus"
	"github.com/google/go-github/github"
	"github.com/lelenanam/downsize"
	"github.com/lelenanam/screenshot"
)

//StartCutiePullReq is the number of pull request
//where new template (with a picture of a cute animal) was added
const StartCutiePullReq = 20514

//Repo is repository with docker cuties
const Repo = "moby"

//Owner of repository
const Owner = "moby"

//SearchIssuesLimit is github limit for Search.Issues results
const SearchIssuesLimit = 1000

//PerPage number of pull requests per page
const PerPage = 50

//TwitterUser is account for docker cuties in twitter
const TwitterUser = "DockerCuties"

//ProjectOwner is a twitter account for notifications
const ProjectOwner = "lelenanam"

//TwitterUploadLimit is limit for media upload in bytes
const TwitterUploadLimit = 3145728

// GetURLFromPull parses the body of the pull request and return last image URL if found
// try to find flickr first
func GetURLFromPull(pull *github.Issue) string {
	str := *pull.Body
	res := ""

	// Example:
	// [![kitteh](https://c2.staticflickr.com/4/3147/2567501805_17ee8fd947_z.jpg)](https://flic.kr/p/4UT7Qv)
	flickre := regexp.MustCompile(`\[!\[.*\]\((.*)\)\]\(.*\)`)
	flickResult := flickre.FindAllStringSubmatch(str, -1)

	if len(flickResult) > 0 {
		lastres := flickResult[len(flickResult)-1]
		if len(lastres) > 1 {
			res := lastres[len(lastres)-1]
			res = strings.SplitN(res, " ", 2)[0]
			return res
		}
	}

	// Example:
	// ![image](https://cloud.githubusercontent.com/assets/2367858/23283487/02bb756e-f9db-11e6-9aa8-5f3e1bb80df3.png  "Swans")
	imagere := regexp.MustCompile(`!\[.*\]\((.*)\)`)
	imageResult := imagere.FindAllStringSubmatch(str, -1)

	if len(imageResult) > 0 {
		lastres := imageResult[len(imageResult)-1]
		if len(lastres) > 1 {
			res := lastres[len(lastres)-1]
			res = strings.SplitN(res, " ", 2)[0]
			return res
		}
	}
	return res
}

// GetImageFromURL downloads an image from url and returns image img, its size, format and error
// for animated gif format returns data in gifByte slice
func GetImageFromURL(url string) (img image.Image, format string, size int, gifByte []byte, err error) {
	log.WithFields(log.Fields{"URL": url}).Debug("Download")
	res, err := http.Get(url)
	if err != nil {
		log.WithFields(log.Fields{"URL": url}).WithError(err).Error("Cannot download")
		return nil, "", 0, nil, err
	}
	log.WithFields(log.Fields{"Content Length": res.ContentLength, "Status Code": res.StatusCode}).Debug("Got")

	defer func() {
		if err := res.Body.Close(); err != nil {
			log.WithFields(log.Fields{"URL": url}).WithError(err).Error("Cannot close body")
		}
	}()

	var reader io.Reader = res.Body

	if res.Header["Content-Type"][0] == "image/gif" {
		gifByte, err = ioutil.ReadAll(res.Body)
		if err != nil {
			log.WithFields(log.Fields{"URL": url}).WithError(err).Error("Cannot read body")
			return nil, "gif", 0, nil, err
		}
		reader = bytes.NewReader(gifByte)
	}

	img, format, err = image.Decode(reader)
	if err != nil {
		log.WithFields(log.Fields{"URL": url}).WithError(err).Error("Cannot decode image")
		return nil, "", 0, nil, err
	}

	log.WithFields(log.Fields{"format": format, "size": int(res.ContentLength)}).Debug("Image decoded")
	return img, format, int(res.ContentLength), gifByte, nil
}

// GetStringFromImage returns string of cutie image img or error
func GetStringFromImage(img image.Image, format string, size int, gifByte []byte) (string, error) {
	b := bytes.NewBuffer(nil)
	encoder := base64.NewEncoder(base64.StdEncoding, b)

	// Need to resize image
	if size >= TwitterUploadLimit || size < 0 {
		log.WithFields(log.Fields{"Twitter upload limit": TwitterUploadLimit}).Debug("Downsize image")
		opts := &downsize.Options{Size: TwitterUploadLimit, Format: format}
		err := downsize.Encode(encoder, img, opts)
		if err != nil {
			log.WithFields(log.Fields{"format": format}).WithError(err).Error("Cannot downsize image")
			return "", err
		}
	} else {
		if format == "gif" {
			gifReader := bytes.NewReader(gifByte)
			_, err := io.Copy(encoder, gifReader)
			if err != nil {
				log.WithFields(log.Fields{"format": format}).WithError(err).Error("Cannot copy gif image")
				return "", err
			}
		} else {
			// no need to resize, just encode
			if err := ImageEncode(encoder, img, format); err != nil {
				log.WithFields(log.Fields{"format": format}).WithError(err).Error("Cannot encode image")
				return "", err
			}
		}
	}

	if err := encoder.Close(); err != nil {
		log.WithError(err).Error("Cannot close encoder")
		return "", err
	}

	if b.Len() == 0 {
		log.Warn("Empty image data")
		return "", nil
	}
	return b.String(), nil
}

var errIsScreenshot = errors.New("picture is screenshot")
var errImageNotFound = errors.New("picture not found")

// GetCutieFromPull returns string of cutie image from pull request pull
func GetCutieFromPull(pull *github.Issue) (string, error) {
	url := GetURLFromPull(pull)
	if url == "" {
		return "", errImageNotFound
	}
	img, format, size, gifByte, err := GetImageFromURL(url)
	if err != nil {
		return "", errors.Wrap(err, "cannot get image from URL")
	}
	if screenshot.Detect(img) {
		return "", errIsScreenshot
	}
	str, err := GetStringFromImage(img, format, size, gifByte)
	if err != nil {
		return "", errors.Wrap(err, "cannot get string for image")
	}
	return str, nil
}

// ImageEncode encodes image m with format to writer w
func ImageEncode(w io.Writer, m image.Image, format string) error {
	switch format {
	case "jpeg":
		return jpeg.Encode(w, m, &jpeg.Options{Quality: 95})
	case "png":
		return png.Encode(w, m)
	case "gif":
		return gif.Encode(w, m, nil)
	default:
		return fmt.Errorf("Unknown image format %q", format)
	}
}
<commit_msg>Add User-Agent header to GET request<commit_after>package main

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"io/ioutil"
	"net/http"
	"regexp"
	"strings"

	"github.com/pkg/errors"

	log "github.com/Sirupsen/logrus"
	"github.com/google/go-github/github"
	"github.com/lelenanam/downsize"
	"github.com/lelenanam/screenshot"
)

//StartCutiePullReq is the number of pull request
//where new template (with a picture of a cute animal) was added
const StartCutiePullReq = 20514

//Repo is repository with docker cuties
const Repo = "moby"

//Owner of repository
const Owner = "moby"

//SearchIssuesLimit is github limit for Search.Issues results
const SearchIssuesLimit = 1000

//PerPage number of pull requests per page
const PerPage = 50

//TwitterUser is account for docker cuties in twitter
const TwitterUser = "DockerCuties"

//ProjectOwner is a twitter account for notifications
const ProjectOwner = "lelenanam"

//TwitterUploadLimit is limit for media upload in bytes
const TwitterUploadLimit = 3145728

// GetURLFromPull parses the body of the pull request and return last image URL if found
// try to find flickr first
func GetURLFromPull(pull *github.Issue) string {
	str := *pull.Body
	res := ""

	// Example:
	// [![kitteh](https://c2.staticflickr.com/4/3147/2567501805_17ee8fd947_z.jpg)](https://flic.kr/p/4UT7Qv)
	flickre := regexp.MustCompile(`\[!\[.*\]\((.*)\)\]\(.*\)`)
	flickResult := flickre.FindAllStringSubmatch(str, -1)

	if len(flickResult) > 0 {
		lastres := flickResult[len(flickResult)-1]
		if len(lastres) > 1 {
			res := lastres[len(lastres)-1]
			res = strings.SplitN(res, " ", 2)[0]
			return res
		}
	}

	// Example:
	// ![image](https://cloud.githubusercontent.com/assets/2367858/23283487/02bb756e-f9db-11e6-9aa8-5f3e1bb80df3.png  "Swans")
	imagere := regexp.MustCompile(`!\[.*\]\((.*)\)`)
	imageResult := imagere.FindAllStringSubmatch(str, -1)

	if len(imageResult) > 0 {
		lastres := imageResult[len(imageResult)-1]
		if len(lastres) > 1 {
			res := lastres[len(lastres)-1]
			res = strings.SplitN(res, " ", 2)[0]
			return res
		}
	}
	return res
}

// GetImageFromURL downloads an image from url and returns image img, its size, format and error
// for animated gif format returns data in gifByte slice
func GetImageFromURL(url string) (img image.Image, format string, size int, gifByte []byte, err error) {
	log.WithFields(log.Fields{"URL": url}).Debug("Download")
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.WithFields(log.Fields{"URL": url}).WithError(err).Error("Cannot create new GET request")
		return nil, "", 0, nil, err
	}
	req.Header.Add("User-Agent", "Mozilla/5.0 Firefox/26.0")
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		log.WithFields(log.Fields{"URL": url}).WithError(err).Error("Cannot download")
		return nil, "", 0, nil, err
	}
	log.WithFields(log.Fields{"Content Length": res.ContentLength, "Status Code": res.StatusCode}).Debug("Got")

	defer func() {
		if err := res.Body.Close(); err != nil {
			log.WithFields(log.Fields{"URL": url}).WithError(err).Error("Cannot close body")
		}
	}()

	var reader io.Reader = res.Body

	if res.Header["Content-Type"][0] == "image/gif" {
		gifByte, err = ioutil.ReadAll(res.Body)
		if err != nil {
			log.WithFields(log.Fields{"URL": url}).WithError(err).Error("Cannot read body")
			return nil, "gif", 0, nil, err
		}
		reader = bytes.NewReader(gifByte)
	}

	img, format, err = image.Decode(reader)
	if err != nil {
		log.WithFields(log.Fields{"URL": url}).WithError(err).Error("Cannot decode image")
		return nil, "", 0, nil, err
	}

	log.WithFields(log.Fields{"format": format, "size": int(res.ContentLength)}).Debug("Image decoded")
	return img, format, int(res.ContentLength), gifByte, nil
}

// GetStringFromImage returns string of cutie image img or error
func GetStringFromImage(img image.Image, format string, size int, gifByte []byte) (string, error) {
	b := bytes.NewBuffer(nil)
	encoder := base64.NewEncoder(base64.StdEncoding, b)

	// Need to resize image
	if size >= TwitterUploadLimit || size < 0 {
		log.WithFields(log.Fields{"Twitter upload limit": TwitterUploadLimit}).Debug("Downsize image")
		opts := &downsize.Options{Size: TwitterUploadLimit, Format: format}
		err := downsize.Encode(encoder, img, opts)
		if err != nil {
			log.WithFields(log.Fields{"format": format}).WithError(err).Error("Cannot downsize image")
			return "", err
		}
	} else {
		if format == "gif" {
			gifReader := bytes.NewReader(gifByte)
			_, err := io.Copy(encoder, gifReader)
			if err != nil {
				log.WithFields(log.Fields{"format": format}).WithError(err).Error("Cannot copy gif image")
				return "", err
			}
		} else {
			// no need to resize, just encode
			if err := ImageEncode(encoder, img, format); err != nil {
				log.WithFields(log.Fields{"format": format}).WithError(err).Error("Cannot encode image")
				return "", err
			}
		}
	}

	if err := encoder.Close(); err != nil {
		log.WithError(err).Error("Cannot close encoder")
		return "", err
	}

	if b.Len() == 0 {
		log.Warn("Empty image data")
		return "", nil
	}
	return b.String(), nil
}

var errIsScreenshot = errors.New("picture is screenshot")
var errImageNotFound = errors.New("picture not found")

// GetCutieFromPull returns string of cutie image from pull request pull
func GetCutieFromPull(pull *github.Issue) (string, error) {
	url := GetURLFromPull(pull)
	if url == "" {
		return "", errImageNotFound
	}
	img, format, size, gifByte, err := GetImageFromURL(url)
	if err != nil {
		return "", errors.Wrap(err, "cannot get image from URL")
	}
	if screenshot.Detect(img) {
		return "", errIsScreenshot
	}
	str, err := GetStringFromImage(img, format, size, gifByte)
	if err != nil {
		return "", errors.Wrap(err, "cannot get string for image")
	}
	return str, nil
}

// ImageEncode encodes image m with format to writer w
func ImageEncode(w io.Writer, m image.Image, format string) error {
	switch format {
	case "jpeg":
		return jpeg.Encode(w, m, &jpeg.Options{Quality: 95})
	case "png":
		return png.Encode(w, m)
	case "gif":
		return gif.Encode(w, m, nil)
	default:
		return fmt.Errorf("Unknown image format %q", format)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package api_test

import (
	"bytes"
	"encoding/base64"
	"io/ioutil"
	"net/http"

	gc "gopkg.in/check.v1"

	"github.com/juju/juju/api"
	jujutesting "github.com/juju/juju/juju/testing"
	"github.com/juju/juju/provider/dummy"
)

type fakeHTTPClient struct {
	calls []string

	response *http.Response
	err      error

	reqArg *http.Request
}

func (d *fakeHTTPClient) Do(req *http.Request) (*http.Response, error) {
	d.calls = append(d.calls, "Do")
	d.reqArg = req
	if d.err != nil {
		return nil, d.err
	}
	return d.response, nil
}

type httpSuite struct {
	jujutesting.JujuConnSuite
	fake *fakeHTTPClient
}

var _ = gc.Suite(&httpSuite{})

func (s *httpSuite) SetUpTest(c *gc.C) {
	s.JujuConnSuite.SetUpTest(c)

	resp := http.Response{
		StatusCode: http.StatusOK,
		Header:     make(http.Header),
		Body:       ioutil.NopCloser(&bytes.Buffer{}),
	}

	fake := fakeHTTPClient{
		response: &resp,
	}
	s.fake = &fake

	s.PatchValue(api.NewHTTPClient,
		func(*api.State) api.HTTPClient {
			return s.fake
		},
	)
}

func (s *httpSuite) checkRequest(c *gc.C, req *http.Request, method, path string) {
	// Only check API-related request fields.

	c.Check(req.Method, gc.Equals, method)

	url := `https://localhost:\d+/environment/[-0-9a-f]+/` + path
	c.Check(req.URL.String(), gc.Matches, url)

	c.Assert(req.Header, gc.HasLen, 1)
	username := dummy.AdminUserTag().String()
	password := jujutesting.AdminSecret
	auth := base64.StdEncoding.EncodeToString([]byte(username + ":" + password))
	c.Check(req.Header.Get("Authorization"), gc.Equals, "Basic "+auth)
}

func (s *httpSuite) TestNewHTTPRequestSuccess(c *gc.C) {
	req, err := s.APIState.NewHTTPRequest("GET", "somefacade")
	c.Assert(err, gc.IsNil)

	s.checkRequest(c, req, "GET", "somefacade")
}

func (s *httpSuite) TestNewHTTPClientCorrectTransport(c *gc.C) {
	doer := s.APIState.NewHTTPClient()

	c.Assert(doer, gc.FitsTypeOf, (*http.Client)(nil))
	httpClient := doer.(*http.Client)

	c.Assert(httpClient.Transport, gc.NotNil)
	c.Assert(httpClient.Transport, gc.FitsTypeOf, (*http.Transport)(nil))
	config := httpClient.Transport.(*http.Transport).TLSClientConfig

	c.Check(config.RootCAs, gc.NotNil)
}

func (s *httpSuite) TestNewHTTPClientValidatesCert(c *gc.C) {
	req, err := s.APIState.NewHTTPRequest("GET", "somefacade")
	doer := s.APIState.NewHTTPClient()
	resp, err := doer.Do(req)
	c.Assert(err, gc.IsNil)

	c.Check(resp.StatusCode, gc.Equals, http.StatusNotFound)
}

func (s *httpSuite) TestSendHTTPRequestSuccess(c *gc.C) {
	req, err := s.APIState.NewHTTPRequest("GET", "somefacade")
	c.Assert(err, gc.IsNil)
	resp, err := s.APIState.SendHTTPRequest(req)
	c.Assert(err, gc.IsNil)

	c.Check(s.fake.calls, gc.DeepEquals, []string{"Do"})
	c.Check(s.fake.reqArg, gc.Equals, req)
	c.Check(resp, gc.Equals, s.fake.response)
}
<commit_msg>Simplify fake error handling.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package api_test

import (
	"bytes"
	"encoding/base64"
	"io/ioutil"
	"net/http"

	gc "gopkg.in/check.v1"

	"github.com/juju/juju/api"
	jujutesting "github.com/juju/juju/juju/testing"
	"github.com/juju/juju/provider/dummy"
)

type fakeHTTPClient struct {
	calls []string

	response *http.Response
	err      error

	reqArg *http.Request
}

func (d *fakeHTTPClient) Do(req *http.Request) (*http.Response, error) {
	d.calls = append(d.calls, "Do")
	d.reqArg = req
	return d.response, d.err
}

type httpSuite struct {
	jujutesting.JujuConnSuite
	fake *fakeHTTPClient
}

var _ = gc.Suite(&httpSuite{})

func (s *httpSuite) SetUpTest(c *gc.C) {
	s.JujuConnSuite.SetUpTest(c)

	resp := http.Response{
		StatusCode: http.StatusOK,
		Header:     make(http.Header),
		Body:       ioutil.NopCloser(&bytes.Buffer{}),
	}

	fake := fakeHTTPClient{
		response: &resp,
	}
	s.fake = &fake

	s.PatchValue(api.NewHTTPClient,
		func(*api.State) api.HTTPClient {
			return s.fake
		},
	)
}

func (s *httpSuite) checkRequest(c *gc.C, req *http.Request, method, path string) {
	// Only check API-related request fields.

	c.Check(req.Method, gc.Equals, method)

	url := `https://localhost:\d+/environment/[-0-9a-f]+/` + path
	c.Check(req.URL.String(), gc.Matches, url)

	c.Assert(req.Header, gc.HasLen, 1)
	username := dummy.AdminUserTag().String()
	password := jujutesting.AdminSecret
	auth := base64.StdEncoding.EncodeToString([]byte(username + ":" + password))
	c.Check(req.Header.Get("Authorization"), gc.Equals, "Basic "+auth)
}

func (s *httpSuite) TestNewHTTPRequestSuccess(c *gc.C) {
	req, err := s.APIState.NewHTTPRequest("GET", "somefacade")
	c.Assert(err, gc.IsNil)

	s.checkRequest(c, req, "GET", "somefacade")
}

func (s *httpSuite) TestNewHTTPClientCorrectTransport(c *gc.C) {
	doer := s.APIState.NewHTTPClient()

	c.Assert(doer, gc.FitsTypeOf, (*http.Client)(nil))
	httpClient := doer.(*http.Client)

	c.Assert(httpClient.Transport, gc.NotNil)
	c.Assert(httpClient.Transport, gc.FitsTypeOf, (*http.Transport)(nil))
	config := httpClient.Transport.(*http.Transport).TLSClientConfig

	c.Check(config.RootCAs, gc.NotNil)
}

func (s *httpSuite) TestNewHTTPClientValidatesCert(c *gc.C) {
	req, err := s.APIState.NewHTTPRequest("GET", "somefacade")
	doer := s.APIState.NewHTTPClient()
	resp, err := doer.Do(req)
	c.Assert(err, gc.IsNil)

	c.Check(resp.StatusCode, gc.Equals, http.StatusNotFound)
}

func (s *httpSuite) TestSendHTTPRequestSuccess(c *gc.C) {
	req, err := s.APIState.NewHTTPRequest("GET", "somefacade")
	c.Assert(err, gc.IsNil)
	resp, err := s.APIState.SendHTTPRequest(req)
	c.Assert(err, gc.IsNil)

	c.Check(s.fake.calls, gc.DeepEquals, []string{"Do"})
	c.Check(s.fake.reqArg, gc.Equals, req)
	c.Check(resp, gc.Equals, s.fake.response)
}
<|endoftext|>


<commit_before>package action_test

import (
	"errors"

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

	. "github.com/maximilien/bosh-softlayer-cpi/action"

	fakedisk "github.com/maximilien/bosh-softlayer-cpi/softlayer/disk/fakes"
)

var _ = Describe("CreateDisk", func() {
	var (
		diskCreator *fakedisk.FakeCreator
		action      CreateDisk
	)

	BeforeEach(func() {
		diskCreator = &fakedisk.FakeCreator{}
		action = NewCreateDisk(diskCreator)
	})

	Describe("Run", func() {
		It("returns id for created disk for specific size", func() {
			diskCreator.CreateDisk = fakedisk.NewFakeDisk(1234)

			id, err := action.Run(20, VMCID(1234))
			Expect(err).ToNot(HaveOccurred())
			Expect(id).To(Equal(DiskCID(1234)))

			Expect(diskCreator.CreateSize).To(Equal(20))
		})

		It("returns error if creating disk fails", func() {
			diskCreator.CreateErr = errors.New("fake-create-err")

			id, err := action.Run(20, VMCID(1234))
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("fake-create-err"))
			Expect(id).To(Equal(DiskCID(1234)))
		})
	})
})
<commit_msg>in fail mode the Run() method for create_disk action should return disk of ID 0<commit_after>package action_test

import (
	"errors"

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

	. "github.com/maximilien/bosh-softlayer-cpi/action"

	fakedisk "github.com/maximilien/bosh-softlayer-cpi/softlayer/disk/fakes"
)

var _ = Describe("CreateDisk", func() {
	var (
		diskCreator *fakedisk.FakeCreator
		action      CreateDisk
	)

	BeforeEach(func() {
		diskCreator = &fakedisk.FakeCreator{}
		action = NewCreateDisk(diskCreator)
	})

	Describe("Run", func() {
		It("returns id for created disk for specific size", func() {
			diskCreator.CreateDisk = fakedisk.NewFakeDisk(1234)

			id, err := action.Run(20, VMCID(1234))
			Expect(err).ToNot(HaveOccurred())
			Expect(id).To(Equal(DiskCID(1234)))

			Expect(diskCreator.CreateSize).To(Equal(20))
		})

		It("returns error if creating disk fails", func() {
			diskCreator.CreateErr = errors.New("fake-create-err")

			id, err := action.Run(20, VMCID(1234))
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("fake-create-err"))
			Expect(id).To(Equal(DiskCID(0)))
		})
	})
})
<|endoftext|>


<commit_before>package api

import (
	"log"
	"sync"
	"testing"
	"time"
)

func TestLock_LockUnlock(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	lock, err := c.LockKey("test/lock")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Initial unlock should fail
	err = lock.Unlock()
	if err != ErrLockNotHeld {
		t.Fatalf("err: %v", err)
	}

	// Should work
	leaderCh, err := lock.Lock(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if leaderCh == nil {
		t.Fatalf("not leader")
	}

	// Double lock should fail
	_, err = lock.Lock(nil)
	if err != ErrLockHeld {
		t.Fatalf("err: %v", err)
	}

	// Should be leader
	select {
	case <-leaderCh:
		t.Fatalf("should be leader")
	default:
	}

	// Initial unlock should work
	err = lock.Unlock()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Double unlock should fail
	err = lock.Unlock()
	if err != ErrLockNotHeld {
		t.Fatalf("err: %v", err)
	}

	// Should loose leadership
	select {
	case <-leaderCh:
	case <-time.After(time.Second):
		t.Fatalf("should not be leader")
	}
}

func TestLock_ForceInvalidate(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	lock, err := c.LockKey("test/lock")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should work
	leaderCh, err := lock.Lock(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if leaderCh == nil {
		t.Fatalf("not leader")
	}
	defer lock.Unlock()

	go func() {
		// Nuke the session, simulator an operator invalidation
		// or a health check failure
		session := c.Session()
		session.Destroy(lock.lockSession, nil)
	}()

	// Should loose leadership
	select {
	case <-leaderCh:
	case <-time.After(time.Second):
		t.Fatalf("should not be leader")
	}
}

func TestLock_DeleteKey(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	// This uncovered some issues around special-case handling of low index
	// numbers where it would work with a low number but fail for higher
	// ones, so we loop this a bit to sweep the index up out of that
	// territory.
	for i := 0; i < 10; i++ {
		func() {
			lock, err := c.LockKey("test/lock")
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			// Should work
			leaderCh, err := lock.Lock(nil)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if leaderCh == nil {
				t.Fatalf("not leader")
			}
			defer lock.Unlock()

			go func() {
				// Nuke the key, simulate an operator intervention
				kv := c.KV()
				kv.Delete("test/lock", nil)
			}()

			// Should loose leadership
			select {
			case <-leaderCh:
			case <-time.After(time.Second):
				t.Fatalf("should not be leader")
			}
		}()
	}
}

func TestLock_Contend(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	wg := &sync.WaitGroup{}
	acquired := make([]bool, 3)
	for idx := range acquired {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			lock, err := c.LockKey("test/lock")
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			// Should work eventually, will contend
			leaderCh, err := lock.Lock(nil)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if leaderCh == nil {
				t.Fatalf("not leader")
			}
			defer lock.Unlock()
			log.Printf("Contender %d acquired", idx)

			// Set acquired and then leave
			acquired[idx] = true
		}(idx)
	}

	// Wait for termination
	doneCh := make(chan struct{})
	go func() {
		wg.Wait()
		close(doneCh)
	}()

	// Wait for everybody to get a turn
	select {
	case <-doneCh:
	case <-time.After(3 * DefaultLockRetryTime):
		t.Fatalf("timeout")
	}

	for idx, did := range acquired {
		if !did {
			t.Fatalf("contender %d never acquired", idx)
		}
	}
}

func TestLock_Destroy(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	lock, err := c.LockKey("test/lock")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should work
	leaderCh, err := lock.Lock(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if leaderCh == nil {
		t.Fatalf("not leader")
	}

	// Destroy should fail
	if err := lock.Destroy(); err != ErrLockHeld {
		t.Fatalf("err: %v", err)
	}

	// Should be able to release
	err = lock.Unlock()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Acquire with a different lock
	l2, err := c.LockKey("test/lock")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should work
	leaderCh, err = l2.Lock(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if leaderCh == nil {
		t.Fatalf("not leader")
	}

	// Destroy should still fail
	if err := lock.Destroy(); err != ErrLockInUse {
		t.Fatalf("err: %v", err)
	}

	// Should release
	err = l2.Unlock()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Destroy should work
	err = lock.Destroy()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Double destroy should work
	err = l2.Destroy()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
}

func TestLock_Conflict(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	sema, err := c.SemaphorePrefix("test/lock/", 2)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should work
	lockCh, err := sema.Acquire(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if lockCh == nil {
		t.Fatalf("not hold")
	}
	defer sema.Release()

	lock, err := c.LockKey("test/lock/.lock")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should conflict with semaphore
	_, err = lock.Lock(nil)
	if err != ErrLockConflict {
		t.Fatalf("err: %v", err)
	}

	// Should conflict with semaphore
	err = lock.Destroy()
	if err != ErrLockConflict {
		t.Fatalf("err: %v", err)
	}
}

func TestLock_ReclaimLock(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	session, _, err := c.Session().Create(&SessionEntry{}, nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	lock, err := c.LockOpts(&LockOptions{Key: "test/lock", Session: session})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should work
	leaderCh, err := lock.Lock(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if leaderCh == nil {
		t.Fatalf("not leader")
	}
	defer lock.Unlock()

	l2, err := c.LockOpts(&LockOptions{Key: "test/lock", Session: session})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	reclaimed := make(chan (<-chan struct{}), 1)
	go func() {
		l2Ch, err := l2.Lock(nil)
		if err != nil {
			t.Fatalf("not locked: %v", err)
		}
		reclaimed <- l2Ch
	}()

	// Should reclaim the lock
	var leader2Ch <-chan struct{}

	select {
	case leader2Ch = <-reclaimed:
	case <-time.After(time.Second):
		t.Fatalf("should have locked")
	}

	// unlock should work
	err = l2.Unlock()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	//Both locks should see the unlock
	select {
	case <-leader2Ch:
	case <-time.After(time.Second):
		t.Fatalf("should not be leader")
	}

	select {
	case <-leaderCh:
	case <-time.After(time.Second):
		t.Fatalf("should not be leader")
	}
}
<commit_msg>Fixes a typo.<commit_after>package api

import (
	"log"
	"sync"
	"testing"
	"time"
)

func TestLock_LockUnlock(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	lock, err := c.LockKey("test/lock")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Initial unlock should fail
	err = lock.Unlock()
	if err != ErrLockNotHeld {
		t.Fatalf("err: %v", err)
	}

	// Should work
	leaderCh, err := lock.Lock(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if leaderCh == nil {
		t.Fatalf("not leader")
	}

	// Double lock should fail
	_, err = lock.Lock(nil)
	if err != ErrLockHeld {
		t.Fatalf("err: %v", err)
	}

	// Should be leader
	select {
	case <-leaderCh:
		t.Fatalf("should be leader")
	default:
	}

	// Initial unlock should work
	err = lock.Unlock()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Double unlock should fail
	err = lock.Unlock()
	if err != ErrLockNotHeld {
		t.Fatalf("err: %v", err)
	}

	// Should lose leadership
	select {
	case <-leaderCh:
	case <-time.After(time.Second):
		t.Fatalf("should not be leader")
	}
}

func TestLock_ForceInvalidate(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	lock, err := c.LockKey("test/lock")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should work
	leaderCh, err := lock.Lock(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if leaderCh == nil {
		t.Fatalf("not leader")
	}
	defer lock.Unlock()

	go func() {
		// Nuke the session, simulator an operator invalidation
		// or a health check failure
		session := c.Session()
		session.Destroy(lock.lockSession, nil)
	}()

	// Should loose leadership
	select {
	case <-leaderCh:
	case <-time.After(time.Second):
		t.Fatalf("should not be leader")
	}
}

func TestLock_DeleteKey(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	// This uncovered some issues around special-case handling of low index
	// numbers where it would work with a low number but fail for higher
	// ones, so we loop this a bit to sweep the index up out of that
	// territory.
	for i := 0; i < 10; i++ {
		func() {
			lock, err := c.LockKey("test/lock")
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			// Should work
			leaderCh, err := lock.Lock(nil)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if leaderCh == nil {
				t.Fatalf("not leader")
			}
			defer lock.Unlock()

			go func() {
				// Nuke the key, simulate an operator intervention
				kv := c.KV()
				kv.Delete("test/lock", nil)
			}()

			// Should loose leadership
			select {
			case <-leaderCh:
			case <-time.After(time.Second):
				t.Fatalf("should not be leader")
			}
		}()
	}
}

func TestLock_Contend(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	wg := &sync.WaitGroup{}
	acquired := make([]bool, 3)
	for idx := range acquired {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			lock, err := c.LockKey("test/lock")
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			// Should work eventually, will contend
			leaderCh, err := lock.Lock(nil)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if leaderCh == nil {
				t.Fatalf("not leader")
			}
			defer lock.Unlock()
			log.Printf("Contender %d acquired", idx)

			// Set acquired and then leave
			acquired[idx] = true
		}(idx)
	}

	// Wait for termination
	doneCh := make(chan struct{})
	go func() {
		wg.Wait()
		close(doneCh)
	}()

	// Wait for everybody to get a turn
	select {
	case <-doneCh:
	case <-time.After(3 * DefaultLockRetryTime):
		t.Fatalf("timeout")
	}

	for idx, did := range acquired {
		if !did {
			t.Fatalf("contender %d never acquired", idx)
		}
	}
}

func TestLock_Destroy(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	lock, err := c.LockKey("test/lock")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should work
	leaderCh, err := lock.Lock(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if leaderCh == nil {
		t.Fatalf("not leader")
	}

	// Destroy should fail
	if err := lock.Destroy(); err != ErrLockHeld {
		t.Fatalf("err: %v", err)
	}

	// Should be able to release
	err = lock.Unlock()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Acquire with a different lock
	l2, err := c.LockKey("test/lock")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should work
	leaderCh, err = l2.Lock(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if leaderCh == nil {
		t.Fatalf("not leader")
	}

	// Destroy should still fail
	if err := lock.Destroy(); err != ErrLockInUse {
		t.Fatalf("err: %v", err)
	}

	// Should release
	err = l2.Unlock()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Destroy should work
	err = lock.Destroy()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Double destroy should work
	err = l2.Destroy()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
}

func TestLock_Conflict(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	sema, err := c.SemaphorePrefix("test/lock/", 2)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should work
	lockCh, err := sema.Acquire(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if lockCh == nil {
		t.Fatalf("not hold")
	}
	defer sema.Release()

	lock, err := c.LockKey("test/lock/.lock")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should conflict with semaphore
	_, err = lock.Lock(nil)
	if err != ErrLockConflict {
		t.Fatalf("err: %v", err)
	}

	// Should conflict with semaphore
	err = lock.Destroy()
	if err != ErrLockConflict {
		t.Fatalf("err: %v", err)
	}
}

func TestLock_ReclaimLock(t *testing.T) {
	t.Parallel()
	c, s := makeClient(t)
	defer s.Stop()

	session, _, err := c.Session().Create(&SessionEntry{}, nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	lock, err := c.LockOpts(&LockOptions{Key: "test/lock", Session: session})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should work
	leaderCh, err := lock.Lock(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if leaderCh == nil {
		t.Fatalf("not leader")
	}
	defer lock.Unlock()

	l2, err := c.LockOpts(&LockOptions{Key: "test/lock", Session: session})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	reclaimed := make(chan (<-chan struct{}), 1)
	go func() {
		l2Ch, err := l2.Lock(nil)
		if err != nil {
			t.Fatalf("not locked: %v", err)
		}
		reclaimed <- l2Ch
	}()

	// Should reclaim the lock
	var leader2Ch <-chan struct{}

	select {
	case leader2Ch = <-reclaimed:
	case <-time.After(time.Second):
		t.Fatalf("should have locked")
	}

	// unlock should work
	err = l2.Unlock()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	//Both locks should see the unlock
	select {
	case <-leader2Ch:
	case <-time.After(time.Second):
		t.Fatalf("should not be leader")
	}

	select {
	case <-leaderCh:
	case <-time.After(time.Second):
		t.Fatalf("should not be leader")
	}
}
<|endoftext|>


<commit_before>// Copyright © 2017 The Kubicorn 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 resources

import (
	"encoding/base64"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/kris-nova/kubicorn/apis/cluster"
	"github.com/kris-nova/kubicorn/cloud"
	"github.com/kris-nova/kubicorn/cutil/compare"
	"github.com/kris-nova/kubicorn/cutil/defaults"
	"github.com/kris-nova/kubicorn/cutil/logger"
	"github.com/kris-nova/kubicorn/cutil/script"
)

var _ cloud.Resource = &Lc{}

type Lc struct {
	Shared
	InstanceType     string
	Image            string
	SpotPrice        string
	ServerPool       *cluster.ServerPool
	BootstrapScripts []string
}

const (
	MasterIPAttempts               = 40
	MasterIPSleepSecondsPerAttempt = 3
)

func (r *Lc) Actual(immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Actual")
	newResource := &Lc{
		Shared: Shared{
			Name: r.Name,
			Tags: make(map[string]string),
		},
	}
	if r.ServerPool.Identifier != "" {
		lcInput := &autoscaling.DescribeLaunchConfigurationsInput{
			LaunchConfigurationNames: []*string{&r.ServerPool.Identifier},
		}
		lcOutput, err := Sdk.ASG.DescribeLaunchConfigurations(lcInput)
		if err != nil {
			return nil, nil, err
		}
		llc := len(lcOutput.LaunchConfigurations)
		if llc != 1 {
			return nil, nil, fmt.Errorf("Found [%d] Launch Configurations for ID [%s]", llc, r.ServerPool.Identifier)
		}
		lc := lcOutput.LaunchConfigurations[0]
		newResource.Image = *lc.ImageId
		if lc.SpotPrice != nil {
			newResource.SpotPrice = *lc.SpotPrice
		}
		newResource.Identifier = *lc.LaunchConfigurationName
		newResource.Tags = map[string]string{
			"Name":              r.Name,
			"KubernetesCluster": immutable.Name,
		}
	} else {
		newResource.Image = r.ServerPool.Image
		newResource.InstanceType = r.ServerPool.Size
		if r.ServerPool.Type == cluster.ServerPoolTypeNode {
			newResource.SpotPrice = r.ServerPool.AwsConfiguration.SpotPrice
		}
	}
	newResource.BootstrapScripts = r.ServerPool.BootstrapScripts

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Expected(immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("asg.Expected")
	newResource := &Lc{
		Shared: Shared{
			Tags: map[string]string{
				"Name":              r.Name,
				"KubernetesCluster": immutable.Name,
			},
			Identifier: r.ServerPool.Identifier,
			Name:       r.Name,
		},
		InstanceType:     r.ServerPool.Size,
		Image:            r.ServerPool.Image,
		BootstrapScripts: r.ServerPool.BootstrapScripts,
	}
	if r.ServerPool.Type == cluster.ServerPoolTypeNode {
		newResource.SpotPrice = r.ServerPool.AwsConfiguration.SpotPrice
	}
	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Apply(actual, expected cloud.Resource, immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Apply")
	applyResource := expected.(*Lc)
	isEqual, err := compare.IsEqual(actual.(*Lc), expected.(*Lc))
	if err != nil {
		return nil, nil, err
	}
	if isEqual {
		return immutable, applyResource, nil
	}
	var sgs []*string
	found := false
	for _, serverPool := range immutable.ServerPools {
		if serverPool.Name == expected.(*Lc).Name {
			for _, firewall := range serverPool.Firewalls {
				sgs = append(sgs, &firewall.Identifier)
			}
			found = true
		}
	}
	if !found {
		return nil, nil, fmt.Errorf("Unable to lookup serverpool for Launch Configuration %s", r.Name)
	}

	// --- Hack in here for master IP
	privip := ""
	pubip := ""
	if strings.Contains(r.ServerPool.Name, "node") {
		found := false
		logger.Debug("Tag query: [%s] %s", "Name", fmt.Sprintf("%s.master", immutable.Name))
		logger.Debug("Tag query: [%s] %s", "KubernetesCluster", immutable.Name)
		for i := 0; i < MasterIPAttempts; i++ {
			logger.Debug("Attempting to lookup master IP for node registration..")
			input := &ec2.DescribeInstancesInput{
				Filters: []*ec2.Filter{
					{
						Name:   S("tag:Name"),
						Values: []*string{S(fmt.Sprintf("%s.master", immutable.Name))},
					},
					{
						Name:   S("tag:KubernetesCluster"),
						Values: []*string{S(immutable.Name)},
					},
				},
			}
			output, err := Sdk.Ec2.DescribeInstances(input)
			if err != nil {
				return nil, nil, err
			}
			lr := len(output.Reservations)
			if lr == 0 {
				logger.Debug("Found %d Reservations, hanging ", lr)
				time.Sleep(time.Duration(MasterIPSleepSecondsPerAttempt) * time.Second)
				continue
			}
			for _, reservation := range output.Reservations {
				for _, instance := range reservation.Instances {
					if instance.PublicIpAddress != nil {
						privip = *instance.PrivateIpAddress
						pubip = *instance.PublicIpAddress
						immutable.Values.ItemMap["INJECTEDMASTER"] = fmt.Sprintf("%s:%s", privip, immutable.KubernetesAPI.Port)
						immutable.KubernetesAPI.Endpoint = pubip
						logger.Info("Found public IP for master: [%s]", pubip)
						found = true
					}
				}
			}
			if found == true {
				break
			}
			time.Sleep(time.Duration(MasterIPSleepSecondsPerAttempt) * time.Second)
		}
		if !found {
			return nil, nil, fmt.Errorf("Unable to find Master IP")
		}
	}

	immutable.Values.ItemMap["INJECTEDPORT"] = immutable.KubernetesAPI.Port

	newResource := &Lc{}
	userData, err := script.BuildBootstrapScript(r.ServerPool.BootstrapScripts, immutable)
	if err != nil {
		return nil, nil, err
	}

	b64data := base64.StdEncoding.EncodeToString(userData)
	lcInput := &autoscaling.CreateLaunchConfigurationInput{
		AssociatePublicIpAddress: B(true),
		LaunchConfigurationName:  &expected.(*Lc).Name,
		ImageId:                  &expected.(*Lc).Image,
		InstanceType:             &expected.(*Lc).InstanceType,
		KeyName:                  &immutable.SSH.Identifier,
		SecurityGroups:           sgs,
		UserData:                 &b64data,
	}

	spotPrice, err := strconv.ParseFloat(*&expected.(*Lc).SpotPrice, 64)
	if *&expected.(*Lc).InstanceType != cluster.ServerPoolTypeMaster && err == nil && spotPrice > 0 {
		lcInput.SpotPrice = &expected.(*Lc).SpotPrice
	}
	_, err = Sdk.ASG.CreateLaunchConfiguration(lcInput)
	if err != nil {
		if awserr, ok := err.(awserr.Error); ok {
			switch awserr.Code() {
			case autoscaling.ErrCodeAlreadyExistsFault:
				logger.Debug(autoscaling.ErrCodeAlreadyExistsFault, awserr.Error())
			case autoscaling.ErrCodeLimitExceededFault:
				logger.Debug(autoscaling.ErrCodeLimitExceededFault, awserr.Error())
			case autoscaling.ErrCodeResourceContentionFault:
				logger.Debug(autoscaling.ErrCodeResourceContentionFault, awserr.Error())
			default:
				logger.Debug(awserr.Error())
			}
		} else {
			logger.Debug(err.Error())
		}
		return nil, nil, err
	}
	logger.Info("Created Launch Configuration [%s]", r.Name)
	newResource.Image = expected.(*Lc).Image
	newResource.InstanceType = expected.(*Lc).InstanceType
	newResource.Name = expected.(*Lc).Name
	newResource.Identifier = expected.(*Lc).Name
	newResource.BootstrapScripts = r.ServerPool.BootstrapScripts

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Delete(actual cloud.Resource, immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Delete")
	deleteResource := actual.(*Lc)
	if deleteResource.Name == "" {
		return nil, nil, fmt.Errorf("Unable to delete Launch Configuration resource without Name [%s]", deleteResource.Name)
	}
	input := &autoscaling.DeleteLaunchConfigurationInput{
		LaunchConfigurationName: &actual.(*Lc).Name,
	}
	_, err := Sdk.ASG.DeleteLaunchConfiguration(input)
	if err != nil {
		return nil, nil, err
	}
	logger.Info("Deleted Launch Configuration [%s]", actual.(*Lc).Name)

	// Kubernetes API
	// Todo (@kris-nova) this obviously isn't immutable
	immutable.KubernetesAPI.Endpoint = ""

	newResource := &Lc{}
	newResource.Name = actual.(*Lc).Name
	newResource.Tags = actual.(*Lc).Tags
	newResource.Image = actual.(*Lc).Image
	newResource.InstanceType = actual.(*Lc).InstanceType
	newResource.BootstrapScripts = actual.(*Lc).BootstrapScripts

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) immutableRender(newResource cloud.Resource, inaccurateCluster *cluster.Cluster) *cluster.Cluster {
	logger.Debug("lc.Render")
	newCluster := defaults.NewClusterDefaults(inaccurateCluster)
	serverPool := &cluster.ServerPool{}
	serverPool.Image = newResource.(*Lc).Image
	serverPool.Size = newResource.(*Lc).InstanceType
	serverPool.BootstrapScripts = newResource.(*Lc).BootstrapScripts
	found := false
	for i := 0; i < len(newCluster.ServerPools); i++ {
		if newCluster.ServerPools[i].Name == newResource.(*Lc).Name {
			newCluster.ServerPools[i].Image = newResource.(*Lc).Image
			newCluster.ServerPools[i].Size = newResource.(*Lc).InstanceType
			newCluster.ServerPools[i].BootstrapScripts = newResource.(*Lc).BootstrapScripts
			found = true
		}
	}
	if !found {
		newCluster.ServerPools = append(newCluster.ServerPools, serverPool)
	}

	return newCluster
}
<commit_msg>aws: add spotprice check<commit_after>// Copyright © 2017 The Kubicorn 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 resources

import (
	"encoding/base64"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/kris-nova/kubicorn/apis/cluster"
	"github.com/kris-nova/kubicorn/cloud"
	"github.com/kris-nova/kubicorn/cutil/compare"
	"github.com/kris-nova/kubicorn/cutil/defaults"
	"github.com/kris-nova/kubicorn/cutil/logger"
	"github.com/kris-nova/kubicorn/cutil/script"
)

var _ cloud.Resource = &Lc{}

type Lc struct {
	Shared
	InstanceType     string
	Image            string
	SpotPrice        string
	ServerPool       *cluster.ServerPool
	BootstrapScripts []string
}

const (
	MasterIPAttempts               = 40
	MasterIPSleepSecondsPerAttempt = 3
)

func (r *Lc) Actual(immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Actual")
	newResource := &Lc{
		Shared: Shared{
			Name: r.Name,
			Tags: make(map[string]string),
		},
	}
	if r.ServerPool.Identifier != "" {
		lcInput := &autoscaling.DescribeLaunchConfigurationsInput{
			LaunchConfigurationNames: []*string{&r.ServerPool.Identifier},
		}
		lcOutput, err := Sdk.ASG.DescribeLaunchConfigurations(lcInput)
		if err != nil {
			return nil, nil, err
		}
		llc := len(lcOutput.LaunchConfigurations)
		if llc != 1 {
			return nil, nil, fmt.Errorf("Found [%d] Launch Configurations for ID [%s]", llc, r.ServerPool.Identifier)
		}
		lc := lcOutput.LaunchConfigurations[0]
		newResource.Image = *lc.ImageId
		if lc.SpotPrice != nil {
			newResource.SpotPrice = *lc.SpotPrice
		}
		newResource.Identifier = *lc.LaunchConfigurationName
		newResource.Tags = map[string]string{
			"Name":              r.Name,
			"KubernetesCluster": immutable.Name,
		}
	} else {
		newResource.Image = r.ServerPool.Image
		newResource.InstanceType = r.ServerPool.Size
		if r.ServerPool.Type == cluster.ServerPoolTypeNode && r.ServerPool.AwsConfiguration != nil {
			newResource.SpotPrice = r.ServerPool.AwsConfiguration.SpotPrice
		}
	}
	newResource.BootstrapScripts = r.ServerPool.BootstrapScripts

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Expected(immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("asg.Expected")
	newResource := &Lc{
		Shared: Shared{
			Tags: map[string]string{
				"Name":              r.Name,
				"KubernetesCluster": immutable.Name,
			},
			Identifier: r.ServerPool.Identifier,
			Name:       r.Name,
		},
		InstanceType:     r.ServerPool.Size,
		Image:            r.ServerPool.Image,
		BootstrapScripts: r.ServerPool.BootstrapScripts,
	}
	if r.ServerPool.Type == cluster.ServerPoolTypeNode {
		newResource.SpotPrice = r.ServerPool.AwsConfiguration.SpotPrice
	}
	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Apply(actual, expected cloud.Resource, immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Apply")
	applyResource := expected.(*Lc)
	isEqual, err := compare.IsEqual(actual.(*Lc), expected.(*Lc))
	if err != nil {
		return nil, nil, err
	}
	if isEqual {
		return immutable, applyResource, nil
	}
	var sgs []*string
	found := false
	for _, serverPool := range immutable.ServerPools {
		if serverPool.Name == expected.(*Lc).Name {
			for _, firewall := range serverPool.Firewalls {
				sgs = append(sgs, &firewall.Identifier)
			}
			found = true
		}
	}
	if !found {
		return nil, nil, fmt.Errorf("Unable to lookup serverpool for Launch Configuration %s", r.Name)
	}

	// --- Hack in here for master IP
	privip := ""
	pubip := ""
	if strings.Contains(r.ServerPool.Name, "node") {
		found := false
		logger.Debug("Tag query: [%s] %s", "Name", fmt.Sprintf("%s.master", immutable.Name))
		logger.Debug("Tag query: [%s] %s", "KubernetesCluster", immutable.Name)
		for i := 0; i < MasterIPAttempts; i++ {
			logger.Debug("Attempting to lookup master IP for node registration..")
			input := &ec2.DescribeInstancesInput{
				Filters: []*ec2.Filter{
					{
						Name:   S("tag:Name"),
						Values: []*string{S(fmt.Sprintf("%s.master", immutable.Name))},
					},
					{
						Name:   S("tag:KubernetesCluster"),
						Values: []*string{S(immutable.Name)},
					},
				},
			}
			output, err := Sdk.Ec2.DescribeInstances(input)
			if err != nil {
				return nil, nil, err
			}
			lr := len(output.Reservations)
			if lr == 0 {
				logger.Debug("Found %d Reservations, hanging ", lr)
				time.Sleep(time.Duration(MasterIPSleepSecondsPerAttempt) * time.Second)
				continue
			}
			for _, reservation := range output.Reservations {
				for _, instance := range reservation.Instances {
					if instance.PublicIpAddress != nil {
						privip = *instance.PrivateIpAddress
						pubip = *instance.PublicIpAddress
						immutable.Values.ItemMap["INJECTEDMASTER"] = fmt.Sprintf("%s:%s", privip, immutable.KubernetesAPI.Port)
						immutable.KubernetesAPI.Endpoint = pubip
						logger.Info("Found public IP for master: [%s]", pubip)
						found = true
					}
				}
			}
			if found == true {
				break
			}
			time.Sleep(time.Duration(MasterIPSleepSecondsPerAttempt) * time.Second)
		}
		if !found {
			return nil, nil, fmt.Errorf("Unable to find Master IP")
		}
	}

	immutable.Values.ItemMap["INJECTEDPORT"] = immutable.KubernetesAPI.Port

	newResource := &Lc{}
	userData, err := script.BuildBootstrapScript(r.ServerPool.BootstrapScripts, immutable)
	if err != nil {
		return nil, nil, err
	}

	b64data := base64.StdEncoding.EncodeToString(userData)
	lcInput := &autoscaling.CreateLaunchConfigurationInput{
		AssociatePublicIpAddress: B(true),
		LaunchConfigurationName:  &expected.(*Lc).Name,
		ImageId:                  &expected.(*Lc).Image,
		InstanceType:             &expected.(*Lc).InstanceType,
		KeyName:                  &immutable.SSH.Identifier,
		SecurityGroups:           sgs,
		UserData:                 &b64data,
	}

	spotPrice, err := strconv.ParseFloat(*&expected.(*Lc).SpotPrice, 64)
	if *&expected.(*Lc).InstanceType != cluster.ServerPoolTypeMaster && err == nil && spotPrice > 0 {
		lcInput.SpotPrice = &expected.(*Lc).SpotPrice
	}
	_, err = Sdk.ASG.CreateLaunchConfiguration(lcInput)
	if err != nil {
		if awserr, ok := err.(awserr.Error); ok {
			switch awserr.Code() {
			case autoscaling.ErrCodeAlreadyExistsFault:
				logger.Debug(autoscaling.ErrCodeAlreadyExistsFault, awserr.Error())
			case autoscaling.ErrCodeLimitExceededFault:
				logger.Debug(autoscaling.ErrCodeLimitExceededFault, awserr.Error())
			case autoscaling.ErrCodeResourceContentionFault:
				logger.Debug(autoscaling.ErrCodeResourceContentionFault, awserr.Error())
			default:
				logger.Debug(awserr.Error())
			}
		} else {
			logger.Debug(err.Error())
		}
		return nil, nil, err
	}
	logger.Info("Created Launch Configuration [%s]", r.Name)
	newResource.Image = expected.(*Lc).Image
	newResource.InstanceType = expected.(*Lc).InstanceType
	newResource.Name = expected.(*Lc).Name
	newResource.Identifier = expected.(*Lc).Name
	newResource.BootstrapScripts = r.ServerPool.BootstrapScripts

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) Delete(actual cloud.Resource, immutable *cluster.Cluster) (*cluster.Cluster, cloud.Resource, error) {
	logger.Debug("lc.Delete")
	deleteResource := actual.(*Lc)
	if deleteResource.Name == "" {
		return nil, nil, fmt.Errorf("Unable to delete Launch Configuration resource without Name [%s]", deleteResource.Name)
	}
	input := &autoscaling.DeleteLaunchConfigurationInput{
		LaunchConfigurationName: &actual.(*Lc).Name,
	}
	_, err := Sdk.ASG.DeleteLaunchConfiguration(input)
	if err != nil {
		return nil, nil, err
	}
	logger.Info("Deleted Launch Configuration [%s]", actual.(*Lc).Name)

	// Kubernetes API
	// Todo (@kris-nova) this obviously isn't immutable
	immutable.KubernetesAPI.Endpoint = ""

	newResource := &Lc{}
	newResource.Name = actual.(*Lc).Name
	newResource.Tags = actual.(*Lc).Tags
	newResource.Image = actual.(*Lc).Image
	newResource.InstanceType = actual.(*Lc).InstanceType
	newResource.BootstrapScripts = actual.(*Lc).BootstrapScripts

	newCluster := r.immutableRender(newResource, immutable)
	return newCluster, newResource, nil
}

func (r *Lc) immutableRender(newResource cloud.Resource, inaccurateCluster *cluster.Cluster) *cluster.Cluster {
	logger.Debug("lc.Render")
	newCluster := defaults.NewClusterDefaults(inaccurateCluster)
	serverPool := &cluster.ServerPool{}
	serverPool.Image = newResource.(*Lc).Image
	serverPool.Size = newResource.(*Lc).InstanceType
	serverPool.BootstrapScripts = newResource.(*Lc).BootstrapScripts
	found := false
	for i := 0; i < len(newCluster.ServerPools); i++ {
		if newCluster.ServerPools[i].Name == newResource.(*Lc).Name {
			newCluster.ServerPools[i].Image = newResource.(*Lc).Image
			newCluster.ServerPools[i].Size = newResource.(*Lc).InstanceType
			newCluster.ServerPools[i].BootstrapScripts = newResource.(*Lc).BootstrapScripts
			found = true
		}
	}
	if !found {
		newCluster.ServerPools = append(newCluster.ServerPools, serverPool)
	}

	return newCluster
}
<|endoftext|>


<commit_before>package swarm

import (
	"fmt"

	"github.com/libp2p/go-libp2p-core/network"

	ma "github.com/multiformats/go-multiaddr"
)

// Listen sets up listeners for all of the given addresses.
// It returns as long as we successfully listen on at least *one* address.
func (s *Swarm) Listen(addrs ...ma.Multiaddr) error {
	errs := make([]error, len(addrs))
	var succeeded int
	for i, a := range addrs {
		if err := s.AddListenAddr(a); err != nil {
			errs[i] = err
		} else {
			succeeded++
		}
	}

	for i, e := range errs {
		if e != nil {
			log.Warningf("listen on %s failed: %s", addrs[i], errs[i])
		}
	}

	if succeeded == 0 && len(addrs) > 0 {
		return fmt.Errorf("failed to listen on any addresses: %s", errs)
	}

	return nil
}

// AddListenAddr tells the swarm to listen on a single address. Unlike Listen,
// this method does not attempt to filter out bad addresses.
func (s *Swarm) AddListenAddr(a ma.Multiaddr) error {
	tpt := s.TransportForListening(a)
	if tpt == nil {
		return ErrNoTransport
	}

	list, err := tpt.Listen(a)
	if err != nil {
		return err
	}

	s.listeners.Lock()
	if s.listeners.m == nil {
		s.listeners.Unlock()
		list.Close()
		return ErrSwarmClosed
	}
	s.refs.Add(1)
	s.listeners.m[list] = struct{}{}
	s.listeners.Unlock()

	maddr := list.Multiaddr()

	// signal to our notifiees on successful conn.
	s.notifyAll(func(n network.Notifiee) {
		n.Listen(s, maddr)
	})

	go func() {
		defer func() {
			list.Close()
			s.listeners.Lock()
			delete(s.listeners.m, list)
			s.listeners.Unlock()
			s.refs.Done()
		}()
		for {
			c, err := list.Accept()
			if err != nil {
				if s.ctx.Err() == nil {
					log.Errorf("swarm listener accept error: %s", err)
				}
				return
			}
			log.Debugf("swarm listener accepted connection: %s", c)
			s.refs.Add(1)
			go func() {
				defer s.refs.Done()
				_, err := s.addConn(c, network.DirInbound)
				if err != nil {
					// Probably just means that the swarm has been closed.
					log.Warningf("add conn failed: ", err)
					return
				}
			}()
		}
	}()
	return nil
}
<commit_msg>invalidate cache for InterfaceListenAddresses on Listen<commit_after>package swarm

import (
	"fmt"
	"time"

	"github.com/libp2p/go-libp2p-core/network"

	ma "github.com/multiformats/go-multiaddr"
)

// Listen sets up listeners for all of the given addresses.
// It returns as long as we successfully listen on at least *one* address.
func (s *Swarm) Listen(addrs ...ma.Multiaddr) error {
	errs := make([]error, len(addrs))
	var succeeded int
	for i, a := range addrs {
		if err := s.AddListenAddr(a); err != nil {
			errs[i] = err
		} else {
			succeeded++
		}
	}

	for i, e := range errs {
		if e != nil {
			log.Warningf("listen on %s failed: %s", addrs[i], errs[i])
		}
	}

	if succeeded == 0 && len(addrs) > 0 {
		return fmt.Errorf("failed to listen on any addresses: %s", errs)
	}

	return nil
}

// AddListenAddr tells the swarm to listen on a single address. Unlike Listen,
// this method does not attempt to filter out bad addresses.
func (s *Swarm) AddListenAddr(a ma.Multiaddr) error {
	tpt := s.TransportForListening(a)
	if tpt == nil {
		return ErrNoTransport
	}

	list, err := tpt.Listen(a)
	if err != nil {
		return err
	}

	s.listeners.Lock()
	if s.listeners.m == nil {
		s.listeners.Unlock()
		list.Close()
		return ErrSwarmClosed
	}
	s.refs.Add(1)
	s.listeners.m[list] = struct{}{}
	s.listeners.cacheEOL = time.Time{}
	s.listeners.Unlock()

	maddr := list.Multiaddr()

	// signal to our notifiees on successful conn.
	s.notifyAll(func(n network.Notifiee) {
		n.Listen(s, maddr)
	})

	go func() {
		defer func() {
			list.Close()
			s.listeners.Lock()
			delete(s.listeners.m, list)
			s.listeners.cacheEOL = time.Time{}
			s.listeners.Unlock()
			s.refs.Done()
		}()
		for {
			c, err := list.Accept()
			if err != nil {
				if s.ctx.Err() == nil {
					log.Errorf("swarm listener accept error: %s", err)
				}
				return
			}
			log.Debugf("swarm listener accepted connection: %s", c)
			s.refs.Add(1)
			go func() {
				defer s.refs.Done()
				_, err := s.addConn(c, network.DirInbound)
				if err != nil {
					// Probably just means that the swarm has been closed.
					log.Warningf("add conn failed: ", err)
					return
				}
			}()
		}
	}()
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"launchpad.net/loggo"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/juju"
	_ "launchpad.net/juju-core/provider/all"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/utils"
)

var logger = loggo.GetLogger("juju.plugins.updatebootstrap")

const updateBootstrapDoc = `
Patches all machines after state server has been restored from backup, to
update state server address to new location.
`

type updateBootstrapCommand struct {
	cmd.EnvCommandBase
}

func (c *updateBootstrapCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "juju-update-bootstrap",
		Purpose: "update all machines after recovering state server",
		Doc:     updateBootstrapDoc,
	}
}

func (c *updateBootstrapCommand) Run(ctx *cmd.Context) error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()
	stateAddr, err := GetStateAddress(conn.Environ)
	if err != nil {
		return err
	}
	fmt.Printf("using state address %v\n", stateAddr)
	return updateAllMachines(conn, stateAddr)
}

// GetStateAddress returns the address of one state server
func GetStateAddress(environ environs.Environ) (string, error) {
	// XXX: Can easily look up state server address using api instead
	stateInfo, _, err := environ.StateInfo()
	if err != nil {
		return "", err
	}
	return strings.Split(stateInfo.Addrs[0], ":")[0], nil
}

var agentAddressTemplate = `
set -ex
cd /var/lib/juju/agents
for agent in *
do
	initctl stop jujud-$agent
	sed -i.old -r "/^(stateaddresses|apiaddresses):/{
		n
		s/- .*(:[0-9]+)/- $ADDR\1/
	}" $agent/agent.conf
	if [[ $agent = unit-* ]]
	then
		sed -i -r 's/change-version: [0-9]+$/change-version: 0/' $agent/state/relations/*/*
	fi
	initctl start jujud-$agent
done
sed -i -r 's/^(:syslogtag, startswith, "juju-" @)(.*)(:[0-9]+)$/\1'$ADDR'\3/' /etc/rsyslog.d/*-juju*.conf
`

// renderScriptArg generates an ssh script argument to update state addresses
func renderScriptArg(stateAddr string) string {
	script := strings.Replace(agentAddressTemplate, "$ADDR", stateAddr, -1)
	return "sudo bash -c " + utils.ShQuote(script)
}

// runMachineUpdate connects via ssh to the machine and runs the update script
func runMachineUpdate(m *state.Machine, sshArg string) error {
	logger.Infof("updating machine: %v\n", m)
	addr := instance.SelectPublicAddress(m.Addresses())
	if addr == "" {
		return fmt.Errorf("no appropriate public address found")
	}
	args := []string{
		"-l", "ubuntu",
		"-T",
		"-o", "StrictHostKeyChecking no",
		"-o", "PasswordAuthentication no",
		addr,
		sshArg,
	}
	c := exec.Command("ssh", args...)
	if data, err := c.CombinedOutput(); err != nil {
		return fmt.Errorf("ssh command failed: %v (%q)", err, data)
	}
	return nil
}

// updateAllMachines finds all machines resets the stored state address
func updateAllMachines(conn *juju.Conn, stateAddr string) error {
	machines, err := conn.State.AllMachines()
	if err != nil {
		return err
	}
	done := make(chan error)
	for _, machine := range machines {
		machine := machine
		go func() {
			err := runMachineUpdate(machine, renderScriptArg(stateAddr))
			if err != nil {
				logger.Errorf("failed to update machine %s: %v", machine, err)
			} else {
				logger.Infof("updated machine %s", machine)
			}
			done <- err
		}()
	}
	err = nil
	for _ = range machines {
		if updateErr := <-done; updateErr != nil && err == nil {
			err = fmt.Errorf("machine update failed")
		}
	}
	return err
}

func Main(args []string) {
	if err := juju.InitJujuHome(); err != nil {
		fmt.Fprintf(os.Stderr, "error: %s\n", err)
		os.Exit(2)
	}
	command := updateBootstrapCommand{}
	os.Exit(cmd.Main(&command, cmd.DefaultContext(), args[1:]))
}

func main() {
	Main(os.Args)
}
<commit_msg>Explictly avoid updating state servers, and tweak set flags on script<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"launchpad.net/loggo"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/juju"
	_ "launchpad.net/juju-core/provider/all"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/utils"
)

var logger = loggo.GetLogger("juju.plugins.updatebootstrap")

const updateBootstrapDoc = `
Patches all machines after state server has been restored from backup, to
update state server address to new location.
`

type updateBootstrapCommand struct {
	cmd.EnvCommandBase
}

func (c *updateBootstrapCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "juju-update-bootstrap",
		Purpose: "update all machines after recovering state server",
		Doc:     updateBootstrapDoc,
	}
}

func (c *updateBootstrapCommand) Run(ctx *cmd.Context) error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()
	stateAddr, err := GetStateAddress(conn.Environ)
	if err != nil {
		return err
	}
	fmt.Printf("using state address %v\n", stateAddr)
	return updateAllMachines(conn, stateAddr)
}

// GetStateAddress returns the address of one state server
func GetStateAddress(environ environs.Environ) (string, error) {
	// XXX: Can easily look up state server address using api instead
	stateInfo, _, err := environ.StateInfo()
	if err != nil {
		return "", err
	}
	return strings.Split(stateInfo.Addrs[0], ":")[0], nil
}

var agentAddressTemplate = `
set -exu
cd /var/lib/juju/agents
for agent in *
do
	initctl stop jujud-$agent
	sed -i.old -r "/^(stateaddresses|apiaddresses):/{
		n
		s/- .*(:[0-9]+)/- $ADDR\1/
	}" $agent/agent.conf
	if [[ $agent = unit-* ]]
	then
		sed -i -r 's/change-version: [0-9]+$/change-version: 0/' $agent/state/relations/*/*
	fi
	initctl start jujud-$agent
done
sed -i -r 's/^(:syslogtag, startswith, "juju-" @)(.*)(:[0-9]+)$/\1'$ADDR'\3/' /etc/rsyslog.d/*-juju*.conf
`

// renderScriptArg generates an ssh script argument to update state addresses
func renderScriptArg(stateAddr string) string {
	script := strings.Replace(agentAddressTemplate, "$ADDR", stateAddr, -1)
	return "sudo bash -c " + utils.ShQuote(script)
}

// runMachineUpdate connects via ssh to the machine and runs the update script
func runMachineUpdate(m *state.Machine, sshArg string) error {
	logger.Infof("updating machine: %v\n", m)
	addr := instance.SelectPublicAddress(m.Addresses())
	if addr == "" {
		return fmt.Errorf("no appropriate public address found")
	}
	args := []string{
		"-l", "ubuntu",
		"-T",
		"-o", "StrictHostKeyChecking no",
		"-o", "PasswordAuthentication no",
		addr,
		sshArg,
	}
	c := exec.Command("ssh", args...)
	if data, err := c.CombinedOutput(); err != nil {
		return fmt.Errorf("ssh command failed: %v (%q)", err, data)
	}
	return nil
}

// updateAllMachines finds all machines resets the stored state address
func updateAllMachines(conn *juju.Conn, stateAddr string) error {
	machines, err := conn.State.AllMachines()
	if err != nil {
		return err
	}
	pendingMachineCount := 0
	done := make(chan error)
	for _, machine := range machines {
		// A newly resumed state server requires no updating, and more
		// than one state server is not yet support by this plugin.
		if machine.IsManager() {
			continue
		}
		pendingMachineCount += 1
		machine := machine
		go func() {
			err := runMachineUpdate(machine, renderScriptArg(stateAddr))
			if err != nil {
				logger.Errorf("failed to update machine %s: %v", machine, err)
			} else {
				logger.Infof("updated machine %s", machine)
			}
			done <- err
		}()
	}
	err = nil
	for ; pendingMachineCount > 0; pendingMachineCount-- {
		if updateErr := <-done; updateErr != nil && err == nil {
			err = fmt.Errorf("machine update failed")
		}
	}
	return err
}

func Main(args []string) {
	if err := juju.InitJujuHome(); err != nil {
		fmt.Fprintf(os.Stderr, "error: %s\n", err)
		os.Exit(2)
	}
	command := updateBootstrapCommand{}
	os.Exit(cmd.Main(&command, cmd.DefaultContext(), args[1:]))
}

func main() {
	Main(os.Args)
}
<|endoftext|>


<commit_before>package linda

import (
	"github.com/amlun/linda/linda/broker"
	_ "github.com/amlun/linda/linda/broker/redis"
	"github.com/amlun/linda/linda/core"
	"github.com/amlun/linda/linda/saver"
	_ "github.com/amlun/linda/linda/saver/cassandra"
)

type dispatcher struct {
	brokerURL string
	broker    broker.Broker
	saverURL  string
	saver     saver.Saver
}

func (d *dispatcher) Init() error {
	b, err := broker.NewBroker(d.brokerURL)
	if err != nil {
		return err
	}
	d.broker = b
	s, err := saver.NewSaver(d.saverURL)
	if err != nil {
		return err
	}
	d.saver = s
	return nil
}

func (d *dispatcher) Close() {
	d.broker.Close()
	d.saver.Close()
}

// push a [period] task to saver
func (d *dispatcher) PushTask(task core.Task) error {
	err := d.saver.PublishTask(&task)
	if err != nil {
		Logger.WithField("action", "PushTask").Errorf("push task error: [%s]", err)
		return err
	}
	Logger.WithField("action", "PushTask").WithField("task", task).Info("ok")
	return nil
}

// push a job to broker and saver
func (d *dispatcher) PushJob(job core.Job) error {
	err := d.broker.PushJob(&job)
	if err != nil {
		Logger.WithField("action", "PushJob").Errorf("push job to broker error: [%s]", err)
		return err
	}
	err = d.saver.PublishJob(&job)
	if err != nil {
		Logger.WithField("action", "PushJob").Errorf("push job to saver error: [%s]", err)
		return err
	}
	Logger.WithField("action", "PushJob").WithField("job", job).Info("ok")
	return nil
}

// get a job and delete it from the queue
func (d *dispatcher) GetJob(queue string) core.Job {
	var job core.Job
	err := d.broker.GetJob(queue, &job)
	if err != nil {
		Logger.WithField("action", "GetJob").Errorf("get job error: [%s]", err)
	} else {
		Logger.WithField("action", "GetJob").WithField("job", job).Info("ok")
	}
	return job

}
<commit_msg>first save the job then enquene the job<commit_after>package linda

import (
	"github.com/amlun/linda/linda/broker"
	_ "github.com/amlun/linda/linda/broker/redis"
	"github.com/amlun/linda/linda/core"
	"github.com/amlun/linda/linda/saver"
	_ "github.com/amlun/linda/linda/saver/cassandra"
)

type dispatcher struct {
	brokerURL string
	broker    broker.Broker
	saverURL  string
	saver     saver.Saver
}

func (d *dispatcher) Init() error {
	b, err := broker.NewBroker(d.brokerURL)
	if err != nil {
		return err
	}
	d.broker = b
	s, err := saver.NewSaver(d.saverURL)
	if err != nil {
		return err
	}
	d.saver = s
	return nil
}

func (d *dispatcher) Close() {
	d.broker.Close()
	d.saver.Close()
}

// push a [period] task to saver
func (d *dispatcher) PushTask(task core.Task) error {
	err := d.saver.PublishTask(&task)
	if err != nil {
		Logger.WithField("action", "PushTask").Errorf("push task error: [%s]", err)
		return err
	}
	Logger.WithField("action", "PushTask").WithField("task", task).Info("ok")
	return nil
}

// push a job to broker and saver
func (d *dispatcher) PushJob(job core.Job) error {
	err := d.saver.PublishJob(&job)
	if err != nil {
		Logger.WithField("action", "PushJob").Errorf("push job to saver error: [%s]", err)
		return err
	}
	err = d.broker.PushJob(&job)
	if err != nil {
		Logger.WithField("action", "PushJob").Errorf("push job to broker error: [%s]", err)
		return err
	}
	Logger.WithField("action", "PushJob").WithField("job", job).Info("ok")
	return nil
}

// get a job and delete it from the queue
func (d *dispatcher) GetJob(queue string) core.Job {
	var job core.Job
	err := d.broker.GetJob(queue, &job)
	if err != nil {
		Logger.WithField("action", "GetJob").Errorf("get job error: [%s]", err)
	} else {
		Logger.WithField("action", "GetJob").WithField("job", job).Info("ok")
	}
	return job

}
<|endoftext|>


<commit_before>/**
 * Just a place holder for now
 * User: danielharper
 * Date: 28/02/2013
 * Time: 21:26
 */
package apu

import (
	"components"
	"log"
	"types"
)

type APU struct{}

func NewAPU() *APU {
	var a *APU = new(APU)
	a.Reset()
	return a
}

func (apu *APU) Name() string {
	return "APU"
}

func (apu *APU) Read(Address types.Word) byte {
	return 0x00
}

func (apu *APU) Write(Address types.Word, Value byte) {

}

func (apu *APU) LinkIRQHandler(m components.IRQHandler) {

}

func (apu *APU) Reset() {
	log.Println(apu.Name()+": Resetting", apu.Name())
}
<commit_msg>Added temporary memory store for APU to fix a problem with Zelda<commit_after>/**
 * Just a place holder for now
 * User: danielharper
 * Date: 28/02/2013
 * Time: 21:26
 */
package apu

import (
	"components"
	"log"
	"types"
)

type APU struct {
	mem [0x41]byte
}

func NewAPU() *APU {
	var a *APU = new(APU)
	a.Reset()
	return a
}

func (apu *APU) Name() string {
	return "APU"
}

func (apu *APU) Read(addr types.Word) byte {
	if addr == 0xFF26 {
		return 0x00
	}
	return apu.mem[addr-0xFF00]
}

func (apu *APU) Write(addr types.Word, value byte) {
	apu.mem[addr-0xFF00] = value
}

func (apu *APU) LinkIRQHandler(m components.IRQHandler) {

}

func (apu *APU) Reset() {
	log.Println(apu.Name()+": Resetting", apu.Name())
}
<|endoftext|>


<commit_before>package plugins

import (
	"fmt"
	"regexp"
	"strings"

	"time"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/helpers"
	"github.com/Seklfreak/Robyul2/metrics"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/bwmarrin/discordgo"
	"github.com/getsentry/raven-go"
	"github.com/globalsign/mgo/bson"
	"github.com/sirupsen/logrus"
	"github.com/vmihailenco/msgpack"
)

type Gallery struct{}

func (g *Gallery) Commands() []string {
	return []string{
		"gallery",
	}
}

const (
	galleryUrlRegexText = `(<?https?:\/\/[^\s]+>?)`
)

var (
	galleryUrlRegex *regexp.Regexp
	galleries       []models.GalleryEntry
)

func (g *Gallery) Init(session *discordgo.Session) {
	galleryUrlRegex = regexp.MustCompile(galleryUrlRegexText)
	var err error
	galleries, err = g.GetGalleries()
	helpers.Relax(err)
}

func (g *Gallery) Uninit(session *discordgo.Session) {

}

func (g *Gallery) Action(command string, content string, msg *discordgo.Message, session *discordgo.Session) {
	if !helpers.ModuleIsAllowed(msg.ChannelID, msg.ID, msg.Author.ID, helpers.ModulePermGallery) {
		return
	}

	args := strings.Fields(content)
	if len(args) >= 1 {
		switch args[0] {
		case "add": // [p]gallery add <source channel> <target channel>
			helpers.RequireMod(msg, func() {
				if len(args) < 3 {
					helpers.SendMessage(msg.ChannelID, helpers.GetText("bot.arguments.too-few"))
					return
				}

				session.ChannelTyping(msg.ChannelID)

				channel, err := helpers.GetChannel(msg.ChannelID)
				helpers.Relax(err)
				guild, err := helpers.GetGuild(channel.GuildID)
				helpers.Relax(err)
				sourceChannel, err := helpers.GetChannelFromMention(msg, args[1])
				if err != nil || sourceChannel.ID == "" || sourceChannel.GuildID != channel.GuildID {
					helpers.SendMessage(msg.ChannelID, helpers.GetText("bot.arguments.invalid"))
					return
				}
				targetChannel, err := helpers.GetChannelFromMention(msg, args[2])
				if err != nil || targetChannel.ID == "" || targetChannel.GuildID != channel.GuildID {
					helpers.SendMessage(msg.ChannelID, helpers.GetText("bot.arguments.invalid"))
					return
				}

				targetChannelPermission, err := session.State.UserChannelPermissions(session.State.User.ID, targetChannel.ID)
				helpers.Relax(err)
				if targetChannelPermission&discordgo.PermissionAdministrator != discordgo.PermissionAdministrator &&
					targetChannelPermission&discordgo.PermissionManageWebhooks != discordgo.PermissionManageWebhooks {
					helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.mirror.add-channel-error-permissions"))
					return
				}

				newID, err := helpers.MDbInsert(models.GalleryTable, models.GalleryEntry{
					SourceChannelID: sourceChannel.ID,
					TargetChannelID: targetChannel.ID,
					GuildID:         channel.GuildID,
					AddedByUserID:   msg.Author.ID,
				})
				helpers.Relax(err)

				_, err = helpers.EventlogLog(time.Now(), channel.GuildID, helpers.MdbIdToHuman(newID),
					models.EventlogTargetTypeRobyulGallery, msg.Author.ID,
					models.EventlogTypeRobyulGalleryAdd, "",
					nil,
					[]models.ElasticEventlogOption{
						{
							Key:   "gallery_sourcechannelid",
							Value: sourceChannel.ID,
						},
						{
							Key:   "gallery_targetchannelid",
							Value: targetChannel.ID,
						},
					}, false)
				helpers.RelaxLog(err)

				cache.GetLogger().WithField("module", "galleries").Info(fmt.Sprintf("Added Gallery on Server %s (%s) posting from #%s (%s) to #%s (%s)",
					guild.Name, guild.ID, sourceChannel.Name, sourceChannel.ID, targetChannel.Name, targetChannel.ID))
				_, err = helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.gallery.add-success"))
				helpers.RelaxMessage(err, msg.ChannelID, msg.ID)

				galleries, err = g.GetGalleries()
				helpers.RelaxLog(err)
				return
			})
		case "list": // [p]gallery list
			session.ChannelTyping(msg.ChannelID)
			channel, err := helpers.GetChannel(msg.ChannelID)
			helpers.Relax(err)
			var entryBucket []models.GalleryEntry
			err = helpers.MDbIter(helpers.MdbCollection(models.GalleryTable).Find(bson.M{"guildid": channel.GuildID})).All(&entryBucket)
			helpers.Relax(err)

			if entryBucket == nil || len(entryBucket) <= 0 {
				helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.gallery.list-empty"))
				return
			}

			resultMessage := ":frame_photo: Galleries on this server:\n"
			for _, entry := range entryBucket {
				resultMessage += fmt.Sprintf("`%s`: posting from <#%s> to <#%s>\n",
					helpers.MdbIdToHuman(entry.ID), entry.SourceChannelID, entry.TargetChannelID)
			}
			resultMessage += fmt.Sprintf("Found **%d** Galleries in total.", len(entryBucket))

			_, err = helpers.SendMessage(msg.ChannelID, resultMessage)
			helpers.RelaxMessage(err, msg.ChannelID, msg.ID)
			return
		case "delete", "del", "remove": // [p]gallery delete <gallery id>
			helpers.RequireAdmin(msg, func() {
				session.ChannelTyping(msg.ChannelID)
				if len(args) < 2 {
					_, err := helpers.SendMessage(msg.ChannelID, helpers.GetText("bot.arguments.too-few"))
					helpers.Relax(err)
					return
				}

				channel, err := helpers.GetChannel(msg.ChannelID)
				helpers.Relax(err)

				var entryBucket models.GalleryEntry
				err = helpers.MdbOne(
					helpers.MdbCollection(models.GalleryTable).Find(bson.M{"guildid": channel.GuildID, "_id": helpers.HumanToMdbId(args[1])}),
					&entryBucket,
				)
				if helpers.IsMdbNotFound(err) {
					helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.gallery.delete-not-found"))
					return
				}
				helpers.Relax(err)

				err = helpers.MDbDelete(models.GalleryTable, entryBucket.ID)
				helpers.Relax(err)

				_, err = helpers.EventlogLog(time.Now(), entryBucket.GuildID, helpers.MdbIdToHuman(entryBucket.ID),
					models.EventlogTargetTypeRobyulGallery, msg.Author.ID,
					models.EventlogTypeRobyulGalleryRemove, "",
					nil,
					[]models.ElasticEventlogOption{
						{
							Key:   "gallery_sourcechannelid",
							Value: entryBucket.SourceChannelID,
						},
						{
							Key:   "gallery_targetchannelid",
							Value: entryBucket.TargetChannelID,
						},
					}, false)
				helpers.RelaxLog(err)

				cache.GetLogger().WithField("module", "galleries").Info(fmt.Sprintf("Deleted Gallery on Server #%s posting from #%s to #%s",
					channel.GuildID, entryBucket.SourceChannelID, entryBucket.TargetChannelID))
				_, err = helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.gallery.delete-success"))
				helpers.Relax(err)

				galleries, err = g.GetGalleries()
				helpers.RelaxLog(err)
				return
			})
		case "refresh": // [p]gallery refresh
			helpers.RequireBotAdmin(msg, func() {
				session.ChannelTyping(msg.ChannelID)
				var err error
				galleries, err = g.GetGalleries()
				helpers.RelaxLog(err)
				_, err = helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.gallery.refreshed-config"))
				helpers.Relax(err)
			})
		}
	}
}

func (g *Gallery) OnMessage(content string, msg *discordgo.Message, session *discordgo.Session) {
	go func() {
		defer helpers.Recover()
	TryNextGallery:
		for _, gallery := range galleries {
			if gallery.SourceChannelID == msg.ChannelID {
				// ignore bot messages
				if msg.Author.Bot == true {
					continue TryNextGallery
				}
				sourceChannel, err := helpers.GetChannel(msg.ChannelID)
				helpers.Relax(err)
				// ignore commands
				prefix := helpers.GetPrefixForServer(sourceChannel.GuildID)
				if prefix != "" {
					if strings.HasPrefix(content, prefix) {
						return
					}
				}
				var linksToRepost []string
				// get mirror attachements
				if len(msg.Attachments) > 0 {
					for _, attachement := range msg.Attachments {
						linksToRepost = append(linksToRepost, attachement.URL)
					}
				}
				// get mirror links
				if strings.Contains(msg.Content, "http") {
					linksFound := galleryUrlRegex.FindAllString(msg.Content, -1)
					if len(linksFound) > 0 {
						for _, linkFound := range linksFound {
							if strings.HasPrefix(linkFound, "<") == false && strings.HasSuffix(linkFound, ">") == false {
								linksToRepost = append(linksToRepost, linkFound)
							}
						}
					}
				}
				// get webhook
				webhooks, err := helpers.GetWebhooks(gallery.GuildID, gallery.TargetChannelID, 1)
				helpers.Relax(err)
				// post mirror links
				if len(linksToRepost) > 0 {
					for _, linkToRepost := range linksToRepost {
						result, err := helpers.WebhookExecuteWithResult(
							webhooks[0].ID,
							webhooks[0].Token,
							&discordgo.WebhookParams{
								Content:   fmt.Sprintf("posted %s", linkToRepost),
								Username:  msg.Author.Username,
								AvatarURL: helpers.GetAvatarUrl(msg.Author),
							},
						)
						if err != nil {
							if errD, ok := err.(*discordgo.RESTError); ok {
								if errD.Message.Code == 10015 {
									cache.GetLogger().WithField("module", "gallery").Warnf("Webhook for gallery #%s not found", gallery.ID)
									continue
								}
							}
							helpers.RelaxLog(err)
							continue
						}
						err = g.rememberPostedMessage(msg, result)
						if err != nil {
							raven.CaptureError(fmt.Errorf("%#v", err), map[string]string{})
						}
						metrics.GalleryPostsSent.Add(1)
					}
				}
			}
		}
	}()
}

type Gallery_PostedMessage struct {
	ChannelID string
	MessageID string
}

func (g *Gallery) rememberPostedMessage(sourceMessage *discordgo.Message, mirroredMessage *discordgo.Message) error {
	redis := cache.GetRedisClient()
	key := fmt.Sprintf("robyul2-discord:gallery:postedmessage:%s", sourceMessage.ID)

	item := new(Gallery_PostedMessage)
	item.ChannelID = mirroredMessage.ChannelID
	item.MessageID = mirroredMessage.ID

	itemBytes, err := msgpack.Marshal(&item)
	if err != nil {
		return err
	}

	_, err = redis.LPush(key, itemBytes).Result()
	if err != nil {
		return err
	}

	_, err = redis.Expire(key, time.Hour*1).Result()
	return err
}

func (g *Gallery) getRememberedMessages(sourceMessage *discordgo.Message) ([]Gallery_PostedMessage, error) {
	redis := cache.GetRedisClient()
	key := fmt.Sprintf("robyul2-discord:gallery:postedmessage:%s", sourceMessage.ID)

	length, err := redis.LLen(key).Result()
	if err != nil {
		return []Gallery_PostedMessage{}, err
	}

	if length <= 0 {
		return []Gallery_PostedMessage{}, err
	}

	result, err := redis.LRange(key, 0, length-1).Result()
	if err != nil {
		return []Gallery_PostedMessage{}, err
	}

	rememberedMessages := make([]Gallery_PostedMessage, 0)
	for _, messageData := range result {
		var message Gallery_PostedMessage
		err = msgpack.Unmarshal([]byte(messageData), &message)
		if err != nil {
			continue
		}
		rememberedMessages = append(rememberedMessages, message)
	}

	return rememberedMessages, nil
}

func (g *Gallery) OnGuildMemberAdd(member *discordgo.Member, session *discordgo.Session) {
}

func (g *Gallery) OnGuildMemberRemove(member *discordgo.Member, session *discordgo.Session) {
}

func (g *Gallery) GetGalleries() (galleries []models.GalleryEntry, err error) {
	err = helpers.MDbIterWithoutLogging(helpers.MdbCollection(models.GalleryTable).Find(nil)).All(&galleries)
	return
}

func (g *Gallery) OnReactionAdd(reaction *discordgo.MessageReactionAdd, session *discordgo.Session) {

}
func (g *Gallery) OnReactionRemove(reaction *discordgo.MessageReactionRemove, session *discordgo.Session) {

}
func (g *Gallery) OnGuildBanAdd(user *discordgo.GuildBanAdd, session *discordgo.Session) {

}
func (g *Gallery) OnGuildBanRemove(user *discordgo.GuildBanRemove, session *discordgo.Session) {

}
func (g *Gallery) OnMessageDelete(msg *discordgo.MessageDelete, session *discordgo.Session) {
	go func() {
		defer helpers.Recover()
		var err error
		var rememberedMessages []Gallery_PostedMessage

		for _, gallery := range galleries {
			if gallery.SourceChannelID == msg.ChannelID {
				rememberedMessages, err = g.getRememberedMessages(msg.Message)
				helpers.Relax(err)

				for _, messageData := range rememberedMessages {
					err = session.ChannelMessageDelete(messageData.ChannelID, messageData.MessageID)
					if err != nil {
						msgAuthorID := "N/A"
						if msg.Author != nil {
							msgAuthorID = msg.Author.ID
						}

						cache.GetLogger().WithFields(logrus.Fields{
							"module":            "gallery",
							"sourceChannelID":   msg.ChannelID,
							"sourceMessageID":   msg.ID,
							"sourceAuthorID":    msgAuthorID,
							"mirroredChannelID": messageData.ChannelID,
							"mirroredMessageID": messageData.MessageID,
						}).Warn(
							"Deleting mirrored message failed:", err.Error(),
						)
					}
				}
			}
		}
	}()
}
<commit_msg>[gallery] displays source channel in message 🏷<commit_after>package plugins

import (
	"fmt"
	"regexp"
	"strings"

	"time"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/helpers"
	"github.com/Seklfreak/Robyul2/metrics"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/bwmarrin/discordgo"
	"github.com/getsentry/raven-go"
	"github.com/globalsign/mgo/bson"
	"github.com/sirupsen/logrus"
	"github.com/vmihailenco/msgpack"
)

type Gallery struct{}

func (g *Gallery) Commands() []string {
	return []string{
		"gallery",
	}
}

const (
	galleryUrlRegexText = `(<?https?:\/\/[^\s]+>?)`
)

var (
	galleryUrlRegex *regexp.Regexp
	galleries       []models.GalleryEntry
)

func (g *Gallery) Init(session *discordgo.Session) {
	galleryUrlRegex = regexp.MustCompile(galleryUrlRegexText)
	var err error
	galleries, err = g.GetGalleries()
	helpers.Relax(err)
}

func (g *Gallery) Uninit(session *discordgo.Session) {

}

func (g *Gallery) Action(command string, content string, msg *discordgo.Message, session *discordgo.Session) {
	if !helpers.ModuleIsAllowed(msg.ChannelID, msg.ID, msg.Author.ID, helpers.ModulePermGallery) {
		return
	}

	args := strings.Fields(content)
	if len(args) >= 1 {
		switch args[0] {
		case "add": // [p]gallery add <source channel> <target channel>
			helpers.RequireMod(msg, func() {
				if len(args) < 3 {
					helpers.SendMessage(msg.ChannelID, helpers.GetText("bot.arguments.too-few"))
					return
				}

				session.ChannelTyping(msg.ChannelID)

				channel, err := helpers.GetChannel(msg.ChannelID)
				helpers.Relax(err)
				guild, err := helpers.GetGuild(channel.GuildID)
				helpers.Relax(err)
				sourceChannel, err := helpers.GetChannelFromMention(msg, args[1])
				if err != nil || sourceChannel.ID == "" || sourceChannel.GuildID != channel.GuildID {
					helpers.SendMessage(msg.ChannelID, helpers.GetText("bot.arguments.invalid"))
					return
				}
				targetChannel, err := helpers.GetChannelFromMention(msg, args[2])
				if err != nil || targetChannel.ID == "" || targetChannel.GuildID != channel.GuildID {
					helpers.SendMessage(msg.ChannelID, helpers.GetText("bot.arguments.invalid"))
					return
				}

				targetChannelPermission, err := session.State.UserChannelPermissions(session.State.User.ID, targetChannel.ID)
				helpers.Relax(err)
				if targetChannelPermission&discordgo.PermissionAdministrator != discordgo.PermissionAdministrator &&
					targetChannelPermission&discordgo.PermissionManageWebhooks != discordgo.PermissionManageWebhooks {
					helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.mirror.add-channel-error-permissions"))
					return
				}

				newID, err := helpers.MDbInsert(models.GalleryTable, models.GalleryEntry{
					SourceChannelID: sourceChannel.ID,
					TargetChannelID: targetChannel.ID,
					GuildID:         channel.GuildID,
					AddedByUserID:   msg.Author.ID,
				})
				helpers.Relax(err)

				_, err = helpers.EventlogLog(time.Now(), channel.GuildID, helpers.MdbIdToHuman(newID),
					models.EventlogTargetTypeRobyulGallery, msg.Author.ID,
					models.EventlogTypeRobyulGalleryAdd, "",
					nil,
					[]models.ElasticEventlogOption{
						{
							Key:   "gallery_sourcechannelid",
							Value: sourceChannel.ID,
						},
						{
							Key:   "gallery_targetchannelid",
							Value: targetChannel.ID,
						},
					}, false)
				helpers.RelaxLog(err)

				cache.GetLogger().WithField("module", "galleries").Info(fmt.Sprintf("Added Gallery on Server %s (%s) posting from #%s (%s) to #%s (%s)",
					guild.Name, guild.ID, sourceChannel.Name, sourceChannel.ID, targetChannel.Name, targetChannel.ID))
				_, err = helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.gallery.add-success"))
				helpers.RelaxMessage(err, msg.ChannelID, msg.ID)

				galleries, err = g.GetGalleries()
				helpers.RelaxLog(err)
				return
			})
		case "list": // [p]gallery list
			session.ChannelTyping(msg.ChannelID)
			channel, err := helpers.GetChannel(msg.ChannelID)
			helpers.Relax(err)
			var entryBucket []models.GalleryEntry
			err = helpers.MDbIter(helpers.MdbCollection(models.GalleryTable).Find(bson.M{"guildid": channel.GuildID})).All(&entryBucket)
			helpers.Relax(err)

			if entryBucket == nil || len(entryBucket) <= 0 {
				helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.gallery.list-empty"))
				return
			}

			resultMessage := ":frame_photo: Galleries on this server:\n"
			for _, entry := range entryBucket {
				resultMessage += fmt.Sprintf("`%s`: posting from <#%s> to <#%s>\n",
					helpers.MdbIdToHuman(entry.ID), entry.SourceChannelID, entry.TargetChannelID)
			}
			resultMessage += fmt.Sprintf("Found **%d** Galleries in total.", len(entryBucket))

			_, err = helpers.SendMessage(msg.ChannelID, resultMessage)
			helpers.RelaxMessage(err, msg.ChannelID, msg.ID)
			return
		case "delete", "del", "remove": // [p]gallery delete <gallery id>
			helpers.RequireAdmin(msg, func() {
				session.ChannelTyping(msg.ChannelID)
				if len(args) < 2 {
					_, err := helpers.SendMessage(msg.ChannelID, helpers.GetText("bot.arguments.too-few"))
					helpers.Relax(err)
					return
				}

				channel, err := helpers.GetChannel(msg.ChannelID)
				helpers.Relax(err)

				var entryBucket models.GalleryEntry
				err = helpers.MdbOne(
					helpers.MdbCollection(models.GalleryTable).Find(bson.M{"guildid": channel.GuildID, "_id": helpers.HumanToMdbId(args[1])}),
					&entryBucket,
				)
				if helpers.IsMdbNotFound(err) {
					helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.gallery.delete-not-found"))
					return
				}
				helpers.Relax(err)

				err = helpers.MDbDelete(models.GalleryTable, entryBucket.ID)
				helpers.Relax(err)

				_, err = helpers.EventlogLog(time.Now(), entryBucket.GuildID, helpers.MdbIdToHuman(entryBucket.ID),
					models.EventlogTargetTypeRobyulGallery, msg.Author.ID,
					models.EventlogTypeRobyulGalleryRemove, "",
					nil,
					[]models.ElasticEventlogOption{
						{
							Key:   "gallery_sourcechannelid",
							Value: entryBucket.SourceChannelID,
						},
						{
							Key:   "gallery_targetchannelid",
							Value: entryBucket.TargetChannelID,
						},
					}, false)
				helpers.RelaxLog(err)

				cache.GetLogger().WithField("module", "galleries").Info(fmt.Sprintf("Deleted Gallery on Server #%s posting from #%s to #%s",
					channel.GuildID, entryBucket.SourceChannelID, entryBucket.TargetChannelID))
				_, err = helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.gallery.delete-success"))
				helpers.Relax(err)

				galleries, err = g.GetGalleries()
				helpers.RelaxLog(err)
				return
			})
		case "refresh": // [p]gallery refresh
			helpers.RequireBotAdmin(msg, func() {
				session.ChannelTyping(msg.ChannelID)
				var err error
				galleries, err = g.GetGalleries()
				helpers.RelaxLog(err)
				_, err = helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.gallery.refreshed-config"))
				helpers.Relax(err)
			})
		}
	}
}

func (g *Gallery) OnMessage(content string, msg *discordgo.Message, session *discordgo.Session) {
	go func() {
		defer helpers.Recover()
	TryNextGallery:
		for _, gallery := range galleries {
			if gallery.SourceChannelID == msg.ChannelID {
				// ignore bot messages
				if msg.Author.Bot == true {
					continue TryNextGallery
				}
				sourceChannel, err := helpers.GetChannel(msg.ChannelID)
				helpers.Relax(err)
				// ignore commands
				prefix := helpers.GetPrefixForServer(sourceChannel.GuildID)
				if prefix != "" {
					if strings.HasPrefix(content, prefix) {
						return
					}
				}
				var linksToRepost []string
				// get mirror attachements
				if len(msg.Attachments) > 0 {
					for _, attachement := range msg.Attachments {
						linksToRepost = append(linksToRepost, attachement.URL)
					}
				}
				// get mirror links
				if strings.Contains(msg.Content, "http") {
					linksFound := galleryUrlRegex.FindAllString(msg.Content, -1)
					if len(linksFound) > 0 {
						for _, linkFound := range linksFound {
							if strings.HasPrefix(linkFound, "<") == false && strings.HasSuffix(linkFound, ">") == false {
								linksToRepost = append(linksToRepost, linkFound)
							}
						}
					}
				}
				// get webhook
				webhooks, err := helpers.GetWebhooks(gallery.GuildID, gallery.TargetChannelID, 1)
				helpers.Relax(err)
				// post mirror links
				if len(linksToRepost) > 0 {
					for _, linkToRepost := range linksToRepost {
						result, err := helpers.WebhookExecuteWithResult(
							webhooks[0].ID,
							webhooks[0].Token,
							&discordgo.WebhookParams{
								Content:   fmt.Sprintf("posted %s in <#%s>", linkToRepost, gallery.SourceChannelID),
								Username:  msg.Author.Username,
								AvatarURL: helpers.GetAvatarUrl(msg.Author),
							},
						)
						if err != nil {
							if errD, ok := err.(*discordgo.RESTError); ok {
								if errD.Message.Code == 10015 {
									cache.GetLogger().WithField("module", "gallery").Warnf("Webhook for gallery #%s not found", gallery.ID)
									continue
								}
							}
							helpers.RelaxLog(err)
							continue
						}
						err = g.rememberPostedMessage(msg, result)
						if err != nil {
							raven.CaptureError(fmt.Errorf("%#v", err), map[string]string{})
						}
						metrics.GalleryPostsSent.Add(1)
					}
				}
			}
		}
	}()
}

type Gallery_PostedMessage struct {
	ChannelID string
	MessageID string
}

func (g *Gallery) rememberPostedMessage(sourceMessage *discordgo.Message, mirroredMessage *discordgo.Message) error {
	redis := cache.GetRedisClient()
	key := fmt.Sprintf("robyul2-discord:gallery:postedmessage:%s", sourceMessage.ID)

	item := new(Gallery_PostedMessage)
	item.ChannelID = mirroredMessage.ChannelID
	item.MessageID = mirroredMessage.ID

	itemBytes, err := msgpack.Marshal(&item)
	if err != nil {
		return err
	}

	_, err = redis.LPush(key, itemBytes).Result()
	if err != nil {
		return err
	}

	_, err = redis.Expire(key, time.Hour*1).Result()
	return err
}

func (g *Gallery) getRememberedMessages(sourceMessage *discordgo.Message) ([]Gallery_PostedMessage, error) {
	redis := cache.GetRedisClient()
	key := fmt.Sprintf("robyul2-discord:gallery:postedmessage:%s", sourceMessage.ID)

	length, err := redis.LLen(key).Result()
	if err != nil {
		return []Gallery_PostedMessage{}, err
	}

	if length <= 0 {
		return []Gallery_PostedMessage{}, err
	}

	result, err := redis.LRange(key, 0, length-1).Result()
	if err != nil {
		return []Gallery_PostedMessage{}, err
	}

	rememberedMessages := make([]Gallery_PostedMessage, 0)
	for _, messageData := range result {
		var message Gallery_PostedMessage
		err = msgpack.Unmarshal([]byte(messageData), &message)
		if err != nil {
			continue
		}
		rememberedMessages = append(rememberedMessages, message)
	}

	return rememberedMessages, nil
}

func (g *Gallery) OnGuildMemberAdd(member *discordgo.Member, session *discordgo.Session) {
}

func (g *Gallery) OnGuildMemberRemove(member *discordgo.Member, session *discordgo.Session) {
}

func (g *Gallery) GetGalleries() (galleries []models.GalleryEntry, err error) {
	err = helpers.MDbIterWithoutLogging(helpers.MdbCollection(models.GalleryTable).Find(nil)).All(&galleries)
	return
}

func (g *Gallery) OnReactionAdd(reaction *discordgo.MessageReactionAdd, session *discordgo.Session) {

}
func (g *Gallery) OnReactionRemove(reaction *discordgo.MessageReactionRemove, session *discordgo.Session) {

}
func (g *Gallery) OnGuildBanAdd(user *discordgo.GuildBanAdd, session *discordgo.Session) {

}
func (g *Gallery) OnGuildBanRemove(user *discordgo.GuildBanRemove, session *discordgo.Session) {

}
func (g *Gallery) OnMessageDelete(msg *discordgo.MessageDelete, session *discordgo.Session) {
	go func() {
		defer helpers.Recover()
		var err error
		var rememberedMessages []Gallery_PostedMessage

		for _, gallery := range galleries {
			if gallery.SourceChannelID == msg.ChannelID {
				rememberedMessages, err = g.getRememberedMessages(msg.Message)
				helpers.Relax(err)

				for _, messageData := range rememberedMessages {
					err = session.ChannelMessageDelete(messageData.ChannelID, messageData.MessageID)
					if err != nil {
						msgAuthorID := "N/A"
						if msg.Author != nil {
							msgAuthorID = msg.Author.ID
						}

						cache.GetLogger().WithFields(logrus.Fields{
							"module":            "gallery",
							"sourceChannelID":   msg.ChannelID,
							"sourceMessageID":   msg.ID,
							"sourceAuthorID":    msgAuthorID,
							"mirroredChannelID": messageData.ChannelID,
							"mirroredMessageID": messageData.MessageID,
						}).Warn(
							"Deleting mirrored message failed:", err.Error(),
						)
					}
				}
			}
		}
	}()
}
<|endoftext|>