<commit_before>package lclip

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"sort"
	"testing"

	"github.com/mitchellh/go-homedir"
)

var tempPath string

func TestMain(m *testing.M) {
	f, err := ioutil.TempFile(os.TempDir(), "test")
	if err != nil {
		fmt.Fprintln(os.Stderr, "lclip_test:", err)
		return
	}
	f.Close()
	tempPath = f.Name()

	e := m.Run()
	defer os.Exit(e)

	os.Remove(tempPath)
}

func TestDefaultPath(t *testing.T) {
	h, err := homedir.Dir()
	if err != nil {
		t.Fatal(err)
	}
	expect := filepath.Join(h, ".lclip.json")
	actual, err := DefaultPath()
	if err != nil {
		t.Errorf("DefaultPath returns %q; want nil", err)
	}
	if actual != expect {
		t.Errorf("DefaultPath = %q; want %q",
			actual, expect)
	}
}

type GetTextTest struct {
	Src string
	Dst string
}

var indexTestsGetText = []GetTextTest{
	{Src: "foo", Dst: "bar"},
	{Src: "hoge", Dst: "piyo"},
}

func TestGetText(t *testing.T) {
	w := bytes.NewBuffer(make([]byte, 0))
	m := make(map[string]string)
	for _, test := range indexTestsGetText {
		m[test.Src] = test.Dst
	}
	if err := json.NewEncoder(w).Encode(m); err != nil {
		t.Fatal(err)
	}
	if err := ioutil.WriteFile(tempPath, w.Bytes(), 0644); err != nil {
		t.Fatal(err)
	}

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Errorf("NewClipboard returns %q; want nil", err)
	}
	defer c.Close()
	for _, test := range indexTestsGetText {
		expect := test.Dst
		actual := c.Get(test.Src)
		if actual != expect {
			t.Errorf("Get(%q) = %q; want %q",
				test.Src, actual, expect)
		}
	}
}

type SetTextTest struct {
	Label string
	Data  string
}

var indexTestsSetText = []SetTextTest{
	{Label: "a", Data: "aaa"},
	{Label: "abc", Data: "def"},
	{Label: "", Data: ""},
}

func TestSetText(t *testing.T) {
	if err := ioutil.WriteFile(tempPath, []byte(`{}`), 0644); err != nil {
		t.Fatal(err)
	}

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Errorf("NewClipboard returns %q; want nil", err)
	}
	defer c.Close()
	for _, test := range indexTestsSetText {
		c.Set(test.Label, test.Data)
		expect := test.Data
		actual := c.Get(test.Label)
		if actual != expect {
			t.Errorf("after Set(%q, %q), Get(%q) = %q; want %q",
				test.Label, test.Data,
				test.Label, actual, expect)
		}
	}
}

var indexTestsLabels = [][]string{
	{"foo", "bar", "baz"},
	{"hoge", "piyo", "fuga"},
}

func TestListLabels(t *testing.T) {
	empty := []byte(`{}`)
	for _, labels := range indexTestsLabels {
		if err := ioutil.WriteFile(tempPath, empty, 0644); err != nil {
			t.Fatal(err)
		}

		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Errorf("NewClipboard returns %q; want nil", err)
		}
		for _, label := range labels {
			c.Set(label, "")
		}

		expect := append(make([]string, 0, len(labels)), labels...)
		actual := c.Labels()
		sort.Strings(expect)
		sort.Strings(actual)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("got %q; want %q", actual, expect)
		}
		if err := c.Close(); err != nil {
			t.Error("Close returns %q; want nil", err)
		}
	}
}

func TestSaveText(t *testing.T) {
	if err := ioutil.WriteFile(tempPath, []byte(`{}`), 0644); err != nil {
		t.Fatal(err)
	}

	k, v := "key", "value"
	{
		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Errorf("NewClipboard returns %q; want nil", err)
		}
		c.Set(k, v)
		if err := c.Close(); err != nil {
			t.Error("Close returns %q; want nil", err)
		}
	}
	{
		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Errorf("NewClipboard returns %q; want nil", err)
		}
		defer c.Close()
		expect := v
		actual := c.Get(k)
		if actual != expect {
			t.Errorf("Get(%q) = %q; want %q",
				k, actual, expect)
		}
	}
}
<commit_msg>Add test for create storage file if not exists<commit_after>package lclip

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"sort"
	"testing"

	"github.com/mitchellh/go-homedir"
)

var tempPath string

func TestMain(m *testing.M) {
	f, err := ioutil.TempFile(os.TempDir(), "test")
	if err != nil {
		fmt.Fprintln(os.Stderr, "lclip_test:", err)
		return
	}
	f.Close()
	tempPath = f.Name()

	e := m.Run()
	defer os.Exit(e)

	os.Remove(tempPath)
}

func TestDefaultPath(t *testing.T) {
	h, err := homedir.Dir()
	if err != nil {
		t.Fatal(err)
	}
	expect := filepath.Join(h, ".lclip.json")
	actual, err := DefaultPath()
	if err != nil {
		t.Errorf("DefaultPath returns %q; want nil", err)
	}
	if actual != expect {
		t.Errorf("DefaultPath = %q; want %q",
			actual, expect)
	}
}

func TestCreateStorageFileIfNotExists(t *testing.T) {
	f, err := ioutil.TempFile(os.TempDir(), "test")
	if err != nil {
		t.Fatal(err)
	}
	if err = f.Close(); err != nil {
		t.Fatal(err)
	}

	path := f.Name()
	if err = os.Remove(path); err != nil {
		t.Fatal(err)
	}
	defer os.Remove(path)

	c, err := NewClipboard(path)
	if err != nil {
		t.Errorf("NewClipboard returns %q; want nil", err)
	}
	defer c.Close()

	expect := []byte("{}\n")
	actual, err := ioutil.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(actual, expect) {
		t.Errorf("got %q; want %q",
			actual, expect)
	}
}

type GetTextTest struct {
	Src string
	Dst string
}

var indexTestsGetText = []GetTextTest{
	{Src: "foo", Dst: "bar"},
	{Src: "hoge", Dst: "piyo"},
}

func TestGetText(t *testing.T) {
	w := bytes.NewBuffer(make([]byte, 0))
	m := make(map[string]string)
	for _, test := range indexTestsGetText {
		m[test.Src] = test.Dst
	}
	if err := json.NewEncoder(w).Encode(m); err != nil {
		t.Fatal(err)
	}
	if err := ioutil.WriteFile(tempPath, w.Bytes(), 0644); err != nil {
		t.Fatal(err)
	}

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Errorf("NewClipboard returns %q; want nil", err)
	}
	defer c.Close()
	for _, test := range indexTestsGetText {
		expect := test.Dst
		actual := c.Get(test.Src)
		if actual != expect {
			t.Errorf("Get(%q) = %q; want %q",
				test.Src, actual, expect)
		}
	}
}

type SetTextTest struct {
	Label string
	Data  string
}

var indexTestsSetText = []SetTextTest{
	{Label: "a", Data: "aaa"},
	{Label: "abc", Data: "def"},
	{Label: "", Data: ""},
}

func TestSetText(t *testing.T) {
	if err := ioutil.WriteFile(tempPath, []byte(`{}`), 0644); err != nil {
		t.Fatal(err)
	}

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Errorf("NewClipboard returns %q; want nil", err)
	}
	defer c.Close()
	for _, test := range indexTestsSetText {
		c.Set(test.Label, test.Data)
		expect := test.Data
		actual := c.Get(test.Label)
		if actual != expect {
			t.Errorf("after Set(%q, %q), Get(%q) = %q; want %q",
				test.Label, test.Data,
				test.Label, actual, expect)
		}
	}
}

var indexTestsLabels = [][]string{
	{"foo", "bar", "baz"},
	{"hoge", "piyo", "fuga"},
}

func TestListLabels(t *testing.T) {
	empty := []byte(`{}`)
	for _, labels := range indexTestsLabels {
		if err := ioutil.WriteFile(tempPath, empty, 0644); err != nil {
			t.Fatal(err)
		}

		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Errorf("NewClipboard returns %q; want nil", err)
		}
		for _, label := range labels {
			c.Set(label, "")
		}

		expect := append(make([]string, 0, len(labels)), labels...)
		actual := c.Labels()
		sort.Strings(expect)
		sort.Strings(actual)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("got %q; want %q", actual, expect)
		}
		if err := c.Close(); err != nil {
			t.Error("Close returns %q; want nil", err)
		}
	}
}

func TestSaveText(t *testing.T) {
	if err := ioutil.WriteFile(tempPath, []byte(`{}`), 0644); err != nil {
		t.Fatal(err)
	}

	k, v := "key", "value"
	{
		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Errorf("NewClipboard returns %q; want nil", err)
		}
		c.Set(k, v)
		if err := c.Close(); err != nil {
			t.Error("Close returns %q; want nil", err)
		}
	}
	{
		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Errorf("NewClipboard returns %q; want nil", err)
		}
		defer c.Close()
		expect := v
		actual := c.Get(k)
		if actual != expect {
			t.Errorf("Get(%q) = %q; want %q",
				k, actual, expect)
		}
	}
}
<|endoftext|>


<commit_before>package integration_test

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/cloudfoundry/libbuildpack/cutlass"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("CF Python Buildpack", func() {
	var app *cutlass.App

	AfterEach(func() {
		if app != nil {
			app.Destroy()
		}
		app = nil
	})

	CleanAnsi := func(s string) string {
		r := strings.NewReplacer("\033[31;1m", "", "\033[33;1m", "", "\033[34;1m", "", "\033[0m", "")
		return r.Replace(s)
	}

	Context("an app that uses miniconda and python 2", func() {
		BeforeEach(func() {
			app = cutlass.New(filepath.Join(bpDir, "fixtures", "miniconda_python_2"))
			app.Disk = "2G"
			app.Memory = "1G"
		})

		It("deploys", func() {
			PushAppAndConfirm(app)

			body, err := app.GetBody("/")
			Expect(err).ToNot(HaveOccurred())
			Expect(body).To(ContainSubstring("numpy: 1.10.4"))
			Expect(body).To(ContainSubstring("scipy: 0.17.0"))
			Expect(body).To(ContainSubstring("sklearn: 0.17.1"))
			Expect(body).To(ContainSubstring("pandas: 0.18.0"))
			Expect(body).To(ContainSubstring("python-version2"))
		})

		AssertUsesProxyDuringStagingIfPresent("miniconda_python_2")
	})

	Context("an app that uses miniconda and python 3", func() {
		var fixtureDir string
		BeforeEach(func() {
			var err error
			fixtureDir, err = cutlass.CopyFixture(filepath.Join(bpDir, "fixtures", "miniconda_python_3"))
			Expect(err).ToNot(HaveOccurred())
			app = cutlass.New(fixtureDir)
			app.Disk = "2G"
			app.Memory = "1G"
		})
		AfterEach(func() { _ = os.RemoveAll(fixtureDir) })

		It("keeps track of environment.yml", func() {
			PushAppAndConfirm(app)

			Expect(app.Stdout.String()).To(ContainSubstring("scipy"))

			body, err := app.GetBody("/")
			Expect(err).ToNot(HaveOccurred())
			Expect(body).To(ContainSubstring("numpy: 1.10.4"))
			Expect(body).To(ContainSubstring("scipy: 0.17.0"))
			Expect(body).To(ContainSubstring("sklearn: 0.17.1"))
			Expect(body).To(ContainSubstring("pandas: 0.18.0"))
			Expect(body).To(ContainSubstring("python-version3"))
		})

		It("doesn't re-download unchanged dependencies", func() {
			PushAppAndConfirm(app)
			PushAppAndConfirm(app)

			// Check that scipy was not re-installed in the logs
			Expect(app.Stdout.String()).ToNot(ContainSubstring("scipy"))
		})

		It("it updates dependencies if environment.yml changes", func() {
			PushAppAndConfirm(app)
			Expect(app.GetBody("/")).To(ContainSubstring("numpy: 1.10.4"))
			Expect(app.GetBody("/")).ToNot(ContainSubstring("numpy: 1.11.0"))

			input, err := ioutil.ReadFile(filepath.Join(fixtureDir, "environment.yml"))
			Expect(err).ToNot(HaveOccurred())
			output := strings.Replace(string(input), "numpy=1.10.4", "numpy=1.11.0", 1)
			Expect(ioutil.WriteFile(filepath.Join(fixtureDir, "environment.yml"), []byte(output), 0644)).To(Succeed())

			PushAppAndConfirm(app)
			Expect(app.GetBody("/")).To(ContainSubstring("numpy: 1.11.0"))
		})

		AssertUsesProxyDuringStagingIfPresent("miniconda_python_3")
	})

	Context("an app that uses miniconda and specifies python 2 in runtime.txt but python3 in the environment.yml", func() {
		BeforeEach(func() {
			app = cutlass.New(filepath.Join(bpDir, "fixtures", "miniconda_python_2_3"))
			app.Disk = "2G"
			app.Memory = "1G"
		})

		It("deploys", func() {
			PushAppAndConfirm(app)
			Expect(app.GetBody("/")).To(ContainSubstring("python-version3"))
			Expect(CleanAnsi(app.Stdout.String())).To(ContainSubstring("**WARNING** you have specified the version of Python runtime both in 'runtime.txt' and 'environment.yml'. You should remove one of the two versions"))
		})

		AssertUsesProxyDuringStagingIfPresent("miniconda_python_2_3")
	})
})
<commit_msg>Reset app log between pushes in integration test<commit_after>package integration_test

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/cloudfoundry/libbuildpack/cutlass"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("CF Python Buildpack", func() {
	var app *cutlass.App

	AfterEach(func() {
		if app != nil {
			app.Destroy()
		}
		app = nil
	})

	CleanAnsi := func(s string) string {
		r := strings.NewReplacer("\033[31;1m", "", "\033[33;1m", "", "\033[34;1m", "", "\033[0m", "")
		return r.Replace(s)
	}

	Context("an app that uses miniconda and python 2", func() {
		BeforeEach(func() {
			app = cutlass.New(filepath.Join(bpDir, "fixtures", "miniconda_python_2"))
			app.Disk = "2G"
			app.Memory = "1G"
		})

		It("deploys", func() {
			PushAppAndConfirm(app)

			body, err := app.GetBody("/")
			Expect(err).ToNot(HaveOccurred())
			Expect(body).To(ContainSubstring("numpy: 1.10.4"))
			Expect(body).To(ContainSubstring("scipy: 0.17.0"))
			Expect(body).To(ContainSubstring("sklearn: 0.17.1"))
			Expect(body).To(ContainSubstring("pandas: 0.18.0"))
			Expect(body).To(ContainSubstring("python-version2"))
		})

		AssertUsesProxyDuringStagingIfPresent("miniconda_python_2")
	})

	Context("an app that uses miniconda and python 3", func() {
		var fixtureDir string
		BeforeEach(func() {
			var err error
			fixtureDir, err = cutlass.CopyFixture(filepath.Join(bpDir, "fixtures", "miniconda_python_3"))
			Expect(err).ToNot(HaveOccurred())
			app = cutlass.New(fixtureDir)
			app.Disk = "2G"
			app.Memory = "1G"
		})
		AfterEach(func() { _ = os.RemoveAll(fixtureDir) })

		It("keeps track of environment.yml", func() {
			PushAppAndConfirm(app)

			Expect(app.Stdout.String()).To(ContainSubstring("scipy"))

			body, err := app.GetBody("/")
			Expect(err).ToNot(HaveOccurred())
			Expect(body).To(ContainSubstring("numpy: 1.10.4"))
			Expect(body).To(ContainSubstring("scipy: 0.17.0"))
			Expect(body).To(ContainSubstring("sklearn: 0.17.1"))
			Expect(body).To(ContainSubstring("pandas: 0.18.0"))
			Expect(body).To(ContainSubstring("python-version3"))
		})

		It("doesn't re-download unchanged dependencies", func() {
			PushAppAndConfirm(app)
			Expect(app.Stdout.String()).To(ContainSubstring("scipy"))

			app.Stdout.Reset()

			PushAppAndConfirm(app)
			// Check that scipy was not re-installed in the logs
			Expect(app.Stdout.String()).ToNot(ContainSubstring("scipy"))
		})

		It("it updates dependencies if environment.yml changes", func() {
			PushAppAndConfirm(app)
			Expect(app.GetBody("/")).To(ContainSubstring("numpy: 1.10.4"))
			Expect(app.GetBody("/")).ToNot(ContainSubstring("numpy: 1.11.0"))

			input, err := ioutil.ReadFile(filepath.Join(fixtureDir, "environment.yml"))
			Expect(err).ToNot(HaveOccurred())
			output := strings.Replace(string(input), "numpy=1.10.4", "numpy=1.11.0", 1)
			Expect(ioutil.WriteFile(filepath.Join(fixtureDir, "environment.yml"), []byte(output), 0644)).To(Succeed())

			PushAppAndConfirm(app)
			Expect(app.GetBody("/")).To(ContainSubstring("numpy: 1.11.0"))
		})

		AssertUsesProxyDuringStagingIfPresent("miniconda_python_3")
	})

	Context("an app that uses miniconda and specifies python 2 in runtime.txt but python3 in the environment.yml", func() {
		BeforeEach(func() {
			app = cutlass.New(filepath.Join(bpDir, "fixtures", "miniconda_python_2_3"))
			app.Disk = "2G"
			app.Memory = "1G"
		})

		It("deploys", func() {
			PushAppAndConfirm(app)
			Expect(app.GetBody("/")).To(ContainSubstring("python-version3"))
			Expect(CleanAnsi(app.Stdout.String())).To(ContainSubstring("**WARNING** you have specified the version of Python runtime both in 'runtime.txt' and 'environment.yml'. You should remove one of the two versions"))
		})

		AssertUsesProxyDuringStagingIfPresent("miniconda_python_2_3")
	})
})
<|endoftext|>


<commit_before>package web

import (
	"bytes"
	"compress/gzip"
	"compress/zlib"
	"encoding/json"
	"encoding/xml"
    //"encoding/base64"
	"fmt"
	"io"
	"os"
	"reflect"
	"strings"
)

/**
Generic {json,xml} encoding interface
*/
type Encoder interface {
	Encode(interface{}) error
}

/**
This is really only useful for specific requests and should probably
be thought out a little more. If a client accepts * / *, how do we
determine what the best Content-Type is for the particular data?
Right now we would just send the original content
*/
func MarshalResponse(ctx *Context, content interface{}) (interface{}, error) {
	var encoded bytes.Buffer
	var enc Encoder

	ctx.SetHeader("Content-Type", "text/plain", true)
	if len(ctx.Request.Header["Accept"]) > 0 {
		for _, accepts := range ctx.Request.Header["Accept"] {

			// Look for a JSON request
			if strings.Index(accepts, "application/json") >= 0 {
				enc = json.NewEncoder(&encoded)
				ctx.SetHeader("Content-Type", "application/json", true)
			}

			// Look for an XML request
			if strings.Index(accepts, "text/xml") >= 0 ||
				strings.Index(accepts, "application/xml") >= 0 {
				if reflect.TypeOf(content).Kind() == reflect.Map {
					ctx.NotAcceptable("Can not encode datatype")
					err := WebError{"Can not encode datatype"}
					return content, err
				}
				enc = xml.NewEncoder(&encoded)
				ctx.SetHeader("Content-Type", "text/xml", true)
			} else if strings.Index(accepts, "image/jpeg") >= 0 {
                ctx.SetHeader("Content-Type", "image/jpeg", true)
                ctx.SetHeader("Content-Length", fmt.Sprintf("%d", len(content.([]byte))), true)
                return content, nil
            } else if strings.Index(accepts, "video/") >= 0 {
				// Setup the connection to receive a vdeo stream
                fmt.Println("Starting stream...")
                if strings.Index(accepts, "video/mp4") >= 0 {
    				ctx.SetHeader("Content-Type", "video/mp4", true)
				    ctx.SetHeader("Content-Disposition", "inline; filename=\"motion.mp4\"", true)
                } else {
    				ctx.SetHeader("Content-Type", "video/mp4", true)
				    ctx.SetHeader("Content-Disposition", "inline; filename=\"motion.mp4\"", true)
                }
				ctx.SetHeader("Connection", "close", true)
				ctx.SetHeader("Accept-Ranges", "bytes", true)

				// We should be passed a File pointer to stream from
				var c *os.File
				t := reflect.TypeOf(content)
				if t.Kind() != reflect.Ptr || t.Elem().String() != "os.File" {
                    ctx.Abort(500, "Can not stream data source")
					return nil, WebError{"Can not stream data source"}
				}

                // Hand out the correct size
				c = content.(*os.File)
				stat, _ := c.Stat()
				ctx.SetHeader("Content-Length", fmt.Sprintf("%d", stat.Size()), true)
                contentrange := fmt.Sprintf("bytes 0-%d", stat.Size())
				ctx.SetHeader("Content-Range", contentrange, true)

                // Do a buffer copy, which will stream the data to the client
                fmt.Println("Starting copy")
//                enc := base64.NewEncoder(base64.StdEncoding, ctx.ResponseWriter)
//				n, err := io.Copy(enc, c)
				n, err := io.Copy(ctx.ResponseWriter, c)
                if err != nil {
                    fmt.Println(err)
                    return nil, WebError{err.Error()}
                }
                fmt.Println("Done copying ", n)
				return nil, nil
			}
		}
	}

	if enc != nil {
		err := enc.Encode(content)
		if err != nil {
			return content, err
		}
		return encoded.Bytes(), nil
	}

	// If we don't have a MIME type handler, just return the
	// original content
	return content, nil
}

/**
Attempts to encode the response according to the client's Accept-Encoding
header. If there is an error, or if the encoding requests aren't supported
then the original content is returned.

Encoding type:
 * deflate (zlib stream)
 * gzip

This should be the last module loaded
*/
func EncodeResponse(ctx *Context, content interface{}) (interface{}, error) {
	var compressed bytes.Buffer
	var output io.WriteCloser

	if len(ctx.Request.Header["Accept-Encoding"]) > 0 {
		for _, opt := range ctx.Request.Header["Accept-Encoding"] {
			if strings.Index(opt, "gzip") >= 0 {
				output = gzip.NewWriter(&compressed)
				ctx.SetHeader("Content-Encoding", "gzip", true)
			}
			if strings.Index(opt, "deflate") >= 0 {
				output = zlib.NewWriter(&compressed)
				ctx.SetHeader("Content-Encoding", "deflate", true)
			}
		}
	}

	if output != nil {
		_, err := output.Write(content.([]byte))
		if err != nil {
			return content, err
		}
		err = output.Close()
		return compressed.Bytes(), nil
	}

	return content, nil
}
<commit_msg>Horrible merge for last minute changes<commit_after>package web

import (
	"bytes"
	"compress/gzip"
	"compress/zlib"
	"encoding/json"
	"encoding/xml"
    //"encoding/base64"
	"fmt"
	"io"
	"os"
	"reflect"
	"strings"
)

/**
Generic {json,xml} encoding interface
*/
type Encoder interface {
	Encode(interface{}) error
}

/**
This is really only useful for specific requests and should probably
be thought out a little more. If a client accepts * / *, how do we
determine what the best Content-Type is for the particular data?
Right now we would just send the original content
*/
func MarshalResponse(ctx *Context, content interface{}) (interface{}, error) {
	var encoded bytes.Buffer
	var enc Encoder

	ctx.SetHeader("Content-Type", "text/plain", true)
	if len(ctx.Request.Header["Accept"]) > 0 {
		for _, accepts := range ctx.Request.Header["Accept"] {

			// Look for a JSON request
			if strings.Index(accepts, "application/json") >= 0 {
				enc = json.NewEncoder(&encoded)
				ctx.SetHeader("Content-Type", "application/json", true)
			}

			// Look for an XML request
			if strings.Index(accepts, "text/xml") >= 0 ||
				strings.Index(accepts, "application/xml") >= 0 {
				if reflect.TypeOf(content).Kind() == reflect.Map {
					ctx.NotAcceptable("Can not encode datatype")
					err := WebError{"Can not encode datatype"}
					return content, err
				}
				enc = xml.NewEncoder(&encoded)
				ctx.SetHeader("Content-Type", "text/xml", true)
			} else if strings.Index(accepts, "image/jpeg") >= 0 {
                ctx.SetHeader("Content-Type", "image/jpeg", true)
                ctx.SetHeader("Content-Length", fmt.Sprintf("%d", len(content.([]byte))), true)
                return content, nil
            } else if strings.Index(accepts, "video/") >= 0 {
				// Setup the connection to receive a vdeo stream
                fmt.Println("Starting stream...")
                if strings.Index(accepts, "video/mp4") >= 0 {
    				ctx.SetHeader("Content-Type", "video/mp4", true)
				    ctx.SetHeader("Content-Disposition", "inline; filename=\"motion.mp4\"", true)
                } else {
    				ctx.SetHeader("Content-Type", "video/ogg", true)
				    ctx.SetHeader("Content-Disposition", "inline; filename=\"motion.ogg\"", true)
                }
				ctx.SetHeader("Connection", "close", true)
				ctx.SetHeader("Accept-Ranges", "bytes", true)

				// We should be passed a File pointer to stream from
				var c *os.File
				t := reflect.TypeOf(content)
				if t.Kind() != reflect.Ptr || t.Elem().String() != "os.File" {
                    ctx.Abort(500, "Can not stream data source")
					return nil, WebError{"Can not stream data source"}
				}

                // Hand out the correct size
				c = content.(*os.File)
				stat, _ := c.Stat()
				ctx.SetHeader("Content-Length", fmt.Sprintf("%d", stat.Size()), true)
                contentrange := fmt.Sprintf("bytes 0-%d", stat.Size())
				ctx.SetHeader("Content-Range", contentrange, true)

                // Do a buffer copy, which will stream the data to the client
                fmt.Println("Starting copy")
//                enc := base64.NewEncoder(base64.StdEncoding, ctx.ResponseWriter)
//				n, err := io.Copy(enc, c)
				n, err := io.Copy(ctx.ResponseWriter, c)
                if err != nil {
                    fmt.Println(err)
                    return nil, WebError{err.Error()}
                }
                fmt.Println("Done copying ", n)
				return nil, nil
			}
		}
	}

	if enc != nil {
		err := enc.Encode(content)
		if err != nil {
			return content, err
		}
		return encoded.Bytes(), nil
	}

	// If we don't have a MIME type handler, just return the
	// original content
	return content, nil
}

/**
Attempts to encode the response according to the client's Accept-Encoding
header. If there is an error, or if the encoding requests aren't supported
then the original content is returned.

Encoding type:
 * deflate (zlib stream)
 * gzip

This should be the last module loaded
*/
func EncodeResponse(ctx *Context, content interface{}) (interface{}, error) {
	var compressed bytes.Buffer
	var output io.WriteCloser

	if len(ctx.Request.Header["Accept-Encoding"]) > 0 {
		for _, opt := range ctx.Request.Header["Accept-Encoding"] {
			if strings.Index(opt, "gzip") >= 0 {
				output = gzip.NewWriter(&compressed)
				ctx.SetHeader("Content-Encoding", "gzip", true)
			}
			if strings.Index(opt, "deflate") >= 0 {
				output = zlib.NewWriter(&compressed)
				ctx.SetHeader("Content-Encoding", "deflate", true)
			}
		}
	}

	if output != nil {
		_, err := output.Write(content.([]byte))
		if err != nil {
			return content, err
		}
		err = output.Close()
		return compressed.Bytes(), nil
	}

	return content, nil
}
<|endoftext|>


<commit_before>// package records contains functions to generate resource records from
// mesos master states to serve through a dns server
package records

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"

	"github.com/mesosphere/mesos-dns/logging"
)

// rrs is a type of question names to resource records answers
type rrs map[string][]string

type slave struct {
	Id       string `json:"id"`
	Hostname string `json:"hostname"`
}

// Slaves is a mapping of id to hostname read in from state.json
type Slaves []slave

// Resources holds our SRV ports
type Resources struct {
	Ports string `json:"ports"`
}

// Tasks holds mesos task information read in from state.json
type Tasks []struct {
	FrameworkId string `json:"framework_id"`
	Id          string `json:"id"`
	Name        string `json:"name"`
	SlaveId     string `json:"slave_id"`
	State       string `json:"state"`
	Resources   `json:"resources"`
}

// Frameworks holds mesos frameworks information read in from state.json
type Frameworks []struct {
	Tasks `json:"tasks"`
	Name  string `json:"name"`
}

// StateJSON is a representation of mesos master state.json
type StateJSON struct {
	Frameworks `json:"frameworks"`
	Slaves     `json:"slaves"`
	Leader     string `json:"leader"`
}

// RecordGenerator is a tmp mapping of resource records and slaves
// maybe de-dupe
// prob. want to break apart
// refactor me - prob. not needed
type RecordGenerator struct {
	As   rrs
	SRVs rrs
	Slaves
}

// hostBySlaveId looks up a hostname by slave_id
func (rg *RecordGenerator) hostBySlaveId(slaveId string) (string, error) {
	for i := 0; i < len(rg.Slaves); i++ {
		if rg.Slaves[i].Id == slaveId {
			return rg.Slaves[i].Hostname, nil
		}
	}

	return "", errors.New("not found")
}

// loadFromMaster loads state.json from mesos master
func (rg *RecordGenerator) loadFromMaster(ip string, port string) (sj StateJSON) {
	// tls ?
	url := "http://" + ip + ":" + port + "/master/state.json"

	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		logging.Error.Println(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		logging.Error.Println(err)
	}

	err = json.Unmarshal(body, &sj)
	if err != nil {
		logging.Error.Println(err)
	}

	return sj
}

// leaderIP returns the ip for the mesos master
func leaderIP(leader string) string {
	pair := strings.Split(leader, "@")[1]
	return strings.Split(pair, ":")[0]
}

// loadWrap catches an attempt to load state.json from a mesos master
// attempts can fail from down server or mesos master secondary
// it also reloads from a different master if the master it attempted to
// load from was not the leader
func (rg *RecordGenerator) loadWrap(ip string, port string) (StateJSON, error) {
	var err error
	var sj StateJSON

	defer func() {
		if rec := recover(); rec != nil {
			err = errors.New("can't connect to mesos")
		}

	}()

	logging.VeryVerbose.Println("reloading from master " + ip)
	sj = rg.loadFromMaster(ip, port)

	if rip := leaderIP(sj.Leader); rip != ip {
		logging.VeryVerbose.Println("master changed to " + ip)
		sj = rg.loadFromMaster(rip, port)
	}

	return sj, err
}

// yankPorts takes an array of port ranges
func yankPorts(ports string) []string {
	rhs := strings.Split(ports, "[")[1]
	lhs := strings.Split(rhs, "]")[0]

	yports := []string{}

	mports := strings.Split(lhs, ",")
	for i := 0; i < len(mports); i++ {
		tmp := strings.TrimSpace(mports[i])
		pz := strings.Split(tmp, "-")
		lo, _ := strconv.Atoi(pz[0])
		hi, _ := strconv.Atoi(pz[1])

		for t := lo; t <= hi; t++ {
			yports = append(yports, strconv.Itoa(t))
		}
	}

	return yports
}

// findMaster tries each master and looks for the leader
// if no leader responds it errors
func (rg *RecordGenerator) findMaster(c Config) (StateJSON, error) {
	var sj StateJSON

	if len(c.leader) != 0 {
		logging.VeryVerbose.Println("Zookeeper says the leader is: ", c.leader)
		ip, port, err := getProto(c.leader)
		if err != nil {
			logging.Error.Println(err)
		}

		sj, _ = rg.loadWrap(ip, port)
		if sj.Leader == "" {
			logging.VeryVerbose.Println("Zookeeper is wrong about leader")
			if len(c.Masters) == 0 {
				return sj, errors.New("no master")
			}
		} else {
			return sj, nil
		}
	}

	// try each listed mesos master before dying
	logging.VeryVerbose.Println("Falling back to Masters config field: ", c.Masters)
	for i := 0; i < len(c.Masters); i++ {
		ip, port, err := getProto(c.Masters[i])
		if err != nil {
			logging.Error.Println(err)
		}

		sj, _ = rg.loadWrap(ip, port)

		if sj.Leader == "" {
			logging.VeryVerbose.Println("not a leader - trying next one")

			if len(c.Masters)-1 == i {
				return sj, errors.New("no master")
			}

		} else {
			return sj, nil
		}

	}

	return sj, nil
}

// should be able to accept
// ip:port
// zk://host1:port1,host2:port2,.../path
// zk://username:password@host1:port1,host2:port2,.../path
// file:///path/to/file (where file contains one of the above)
func getProto(pair string) (string, string, error) {
	h := strings.Split(pair, ":")
	return h[0], h[1], nil
}

// ParseState parses a state.json from a mesos master
// it sets the resource records map for the resolver
// with the following format
//
//  _<tag>.<service>.<framework>._<protocol>..mesos
// it also tries different mesos masters if one is not up
// this will shudown if it can't connect to a mesos master
func (rg *RecordGenerator) ParseState(config Config) {

	// try each listed mesos master before dying
	sj, err := rg.findMaster(config)
	if err != nil {
		logging.Error.Println("no master")
		return
	}

	rg.InsertState(sj, config.Domain, config.Mname, config.Listener, config.Masters)
}

// cleanName sanitizes invalid characters
func cleanName(tname string) string {
	return stripInvalid(tname)
}

// stripInvalid remove any non-valid hostname characters
func stripInvalid(tname string) string {
	reg, err := regexp.Compile("[^\\w-.\\.]")
	if err != nil {
		logging.Error.Println(err)
	}

	s := reg.ReplaceAllString(tname, "")

	return strings.ToLower(strings.Replace(s, "_", "", -1))
}

// InsertState transforms a StateJSON into RecordGenerator RRs
func (rg *RecordGenerator) InsertState(sj StateJSON, domain string, mname string,
	listener string, masters []string) error {
	rg.Slaves = sj.Slaves

	rg.SRVs = make(rrs)
	rg.As = make(rrs)

	f := sj.Frameworks

	// complete crap - refactor me
	for i := 0; i < len(f); i++ {
		fname := f[i].Name
		fname = cleanName(fname)

		for x := 0; x < len(f[i].Tasks); x++ {
			task := f[i].Tasks[x]

			host, err := rg.hostBySlaveId(task.SlaveId)
			if err == nil && (task.State == "TASK_RUNNING") {

				tname := cleanName(task.Name)
				tail := fname + "." + domain + "."

				// hack - what to do?
				if task.Resources.Ports != "" {
					sports := yankPorts(task.Resources.Ports)

					// FIXME - 3 nested loops
					for s := 0; s < len(sports); s++ {
						var srvhost string = tname + "." + fname + "." + domain + ":" + sports[s]

						tcp := "_" + tname + "._tcp." + tail
						udp := "_" + tname + "._udp." + tail

						rg.insertRR(tcp, srvhost, "SRV")
						rg.insertRR(udp, srvhost, "SRV")
					}

				}

				arec := tname + "." + tail
				rg.insertRR(arec, host, "A")

			}
		}
	}

	rg.listenerRecord(listener, mname)
	rg.masterRecord(domain, masters, sj.Leader)
	return nil
}

// listenerRecord sets the A record for the mesos-dns server in case
// there is a request for it's hostname (eg: from SOA mname)
func (rg *RecordGenerator) listenerRecord(listener string, mname string) {
	if listener == "0.0.0.0" {
		rg.setFromLocal(listener, mname)
	} else if listener == "127.0.0.1" {
		rg.insertRR(mname, "127.0.0.1", "A")
	} else {
		rg.insertRR(mname, listener, "A")
	}
}

// masterRecord sets A records for the mesos masters and an A record
// for the leading master
func (rg *RecordGenerator) masterRecord(domain string, masters []string, leader string) {
	logging.VeryVerbose.Println("MasterRecord called" + leader)

	// create records for leader
	// A records
	h := strings.Split(leader, "@")
	ip, port, err := getProto(h[0])
	if err != nil {
		logging.Error.Println(err)
	}
	arec := "leader." + domain + "."
	rg.insertRR(arec, ip, "A")
	arec = "master." + domain + "."
	rg.insertRR(arec, ip, "A")
	// SRV records
	tcp := "_leader._tcp." + domain + "."
	udp := "_leader._udp." + domain + "."
	host := "leader." + domain + ":" + port
	rg.insertRR(tcp, host, "SRV")
	rg.insertRR(udp, host, "SRV")
	tcp = "_master._tcp." + domain + "."
	udp = "_master._udp." + domain + "."
	host = "master." + domain + ":" + port
	rg.insertRR(tcp, host, "SRV")
	rg.insertRR(udp, host, "SRV")

	// if there is a list of masters, insert that as well
	for i := 0; i < len(masters); i++ {

		// skip leader
		if leader == masters[i] {
			continue
		}

		ip, port, err := getProto(masters[i])
		if err != nil {
			logging.Error.Println(err)
		}

		// A records (master and masterN)
		arec := "master." + domain + "."
		rg.insertRR(arec, ip, "A")
		arec = "master" + strconv.Itoa(i) + "." + domain + "."
		rg.insertRR(arec, ip, "A")

		// SRV records
		tcp := "_master._tcp." + domain + "."
		udp := "_master._udp." + domain + "."
		host := "master." + domain + ":" + port
		rg.insertRR(tcp, host, "SRV")
		rg.insertRR(udp, host, "SRV")
	}
}

// setFromLocal generates A records for each local interface we are
// listening on - if this causes problems you should explicitly set the
// listener address in config.json
func (rg *RecordGenerator) setFromLocal(host string, mname string) {

	ifaces, err := net.Interfaces()
	if err != nil {
		logging.Error.Println(err)
	}

	// handle err
	for _, i := range ifaces {

		addrs, err := i.Addrs()
		if err != nil {
			logging.Error.Println(err)
		}

		for _, addr := range addrs {
			var ip net.IP
			switch v := addr.(type) {
			case *net.IPNet:
				ip = v.IP
			case *net.IPAddr:
				ip = v.IP
			}

			if ip == nil || ip.IsLoopback() {
				continue
			}

			ip = ip.To4()
			if ip == nil {
				continue
			}

			rg.insertRR(mname, ip.String(), "A")
		}

	}
}

func stripHost(hostip string) string {
	return strings.Split(hostip, ":")[0]
}

// insertRR inserts host to name's map
// refactor me
func (rg *RecordGenerator) insertRR(name string, host string, rtype string) {
	logging.VeryVerbose.Println("[" + rtype + "]\t" + name + ": " + host)

	if rtype == "A" {
		if val, ok := rg.As[name]; ok {

			h := stripHost(host)
			for _, b := range val {
				if stripHost(b) == h {
					return
				}
			}

			rg.As[name] = append(val, host)
		} else {
			rg.As[name] = []string{host}
		}
	} else {
		if val, ok := rg.SRVs[name]; ok {
			rg.SRVs[name] = append(val, host)
		} else {
			rg.SRVs[name] = []string{host}
		}
	}
}
<commit_msg>Fixed string.<commit_after>// package records contains functions to generate resource records from
// mesos master states to serve through a dns server
package records

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"

	"github.com/mesosphere/mesos-dns/logging"
)

// rrs is a type of question names to resource records answers
type rrs map[string][]string

type slave struct {
	Id       string `json:"id"`
	Hostname string `json:"hostname"`
}

// Slaves is a mapping of id to hostname read in from state.json
type Slaves []slave

// Resources holds our SRV ports
type Resources struct {
	Ports string `json:"ports"`
}

// Tasks holds mesos task information read in from state.json
type Tasks []struct {
	FrameworkId string `json:"framework_id"`
	Id          string `json:"id"`
	Name        string `json:"name"`
	SlaveId     string `json:"slave_id"`
	State       string `json:"state"`
	Resources   `json:"resources"`
}

// Frameworks holds mesos frameworks information read in from state.json
type Frameworks []struct {
	Tasks `json:"tasks"`
	Name  string `json:"name"`
}

// StateJSON is a representation of mesos master state.json
type StateJSON struct {
	Frameworks `json:"frameworks"`
	Slaves     `json:"slaves"`
	Leader     string `json:"leader"`
}

// RecordGenerator is a tmp mapping of resource records and slaves
// maybe de-dupe
// prob. want to break apart
// refactor me - prob. not needed
type RecordGenerator struct {
	As   rrs
	SRVs rrs
	Slaves
}

// hostBySlaveId looks up a hostname by slave_id
func (rg *RecordGenerator) hostBySlaveId(slaveId string) (string, error) {
	for i := 0; i < len(rg.Slaves); i++ {
		if rg.Slaves[i].Id == slaveId {
			return rg.Slaves[i].Hostname, nil
		}
	}

	return "", errors.New("not found")
}

// loadFromMaster loads state.json from mesos master
func (rg *RecordGenerator) loadFromMaster(ip string, port string) (sj StateJSON) {
	// tls ?
	url := "http://" + ip + ":" + port + "/master/state.json"

	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		logging.Error.Println(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		logging.Error.Println(err)
	}

	err = json.Unmarshal(body, &sj)
	if err != nil {
		logging.Error.Println(err)
	}

	return sj
}

// leaderIP returns the ip for the mesos master
func leaderIP(leader string) string {
	pair := strings.Split(leader, "@")[1]
	return strings.Split(pair, ":")[0]
}

// loadWrap catches an attempt to load state.json from a mesos master
// attempts can fail from down server or mesos master secondary
// it also reloads from a different master if the master it attempted to
// load from was not the leader
func (rg *RecordGenerator) loadWrap(ip string, port string) (StateJSON, error) {
	var err error
	var sj StateJSON

	defer func() {
		if rec := recover(); rec != nil {
			err = errors.New("can't connect to mesos")
		}

	}()

	logging.VeryVerbose.Println("reloading from master " + ip)
	sj = rg.loadFromMaster(ip, port)

	if rip := leaderIP(sj.Leader); rip != ip {
		logging.VeryVerbose.Println("master changed to " + ip)
		sj = rg.loadFromMaster(rip, port)
	}

	return sj, err
}

// yankPorts takes an array of port ranges
func yankPorts(ports string) []string {
	rhs := strings.Split(ports, "[")[1]
	lhs := strings.Split(rhs, "]")[0]

	yports := []string{}

	mports := strings.Split(lhs, ",")
	for i := 0; i < len(mports); i++ {
		tmp := strings.TrimSpace(mports[i])
		pz := strings.Split(tmp, "-")
		lo, _ := strconv.Atoi(pz[0])
		hi, _ := strconv.Atoi(pz[1])

		for t := lo; t <= hi; t++ {
			yports = append(yports, strconv.Itoa(t))
		}
	}

	return yports
}

// findMaster tries each master and looks for the leader
// if no leader responds it errors
func (rg *RecordGenerator) findMaster(c Config) (StateJSON, error) {
	var sj StateJSON

	if len(c.leader) != 0 {
		logging.VeryVerbose.Println("Zookeeper says the leader is: ", c.leader)
		ip, port, err := getProto(c.leader)
		if err != nil {
			logging.Error.Println(err)
		}

		sj, _ = rg.loadWrap(ip, port)
		if sj.Leader == "" {
			logging.VeryVerbose.Println("Zookeeper is wrong about leader")
			if len(c.Masters) == 0 {
				return sj, errors.New("no master")
			}
		} else {
			return sj, nil
		}
	}

	// try each listed mesos master before dying
	logging.VeryVerbose.Println("Falling back to Masters config field: ", c.Masters)
	for i := 0; i < len(c.Masters); i++ {
		ip, port, err := getProto(c.Masters[i])
		if err != nil {
			logging.Error.Println(err)
		}

		sj, _ = rg.loadWrap(ip, port)

		if sj.Leader == "" {
			logging.VeryVerbose.Println("not a leader - trying next one")

			if len(c.Masters)-1 == i {
				return sj, errors.New("no master")
			}

		} else {
			return sj, nil
		}

	}

	return sj, nil
}

// should be able to accept
// ip:port
// zk://host1:port1,host2:port2,.../path
// zk://username:password@host1:port1,host2:port2,.../path
// file:///path/to/file (where file contains one of the above)
func getProto(pair string) (string, string, error) {
	h := strings.Split(pair, ":")
	return h[0], h[1], nil
}

// ParseState parses a state.json from a mesos master
// it sets the resource records map for the resolver
// with the following format
//
//  _<tag>.<service>.<framework>._<protocol>..mesos
// it also tries different mesos masters if one is not up
// this will shudown if it can't connect to a mesos master
func (rg *RecordGenerator) ParseState(config Config) {

	// try each listed mesos master before dying
	sj, err := rg.findMaster(config)
	if err != nil {
		logging.Error.Println("no master")
		return
	}

	rg.InsertState(sj, config.Domain, config.Mname, config.Listener, config.Masters)
}

// cleanName sanitizes invalid characters
func cleanName(tname string) string {
	return stripInvalid(tname)
}

// stripInvalid remove any non-valid hostname characters
func stripInvalid(tname string) string {
	reg, err := regexp.Compile("[^\\w-.\\.]")
	if err != nil {
		logging.Error.Println(err)
	}

	s := reg.ReplaceAllString(tname, "")

	return strings.ToLower(strings.Replace(s, "_", "", -1))
}

// InsertState transforms a StateJSON into RecordGenerator RRs
func (rg *RecordGenerator) InsertState(sj StateJSON, domain string, mname string,
	listener string, masters []string) error {
	rg.Slaves = sj.Slaves

	rg.SRVs = make(rrs)
	rg.As = make(rrs)

	f := sj.Frameworks

	// complete crap - refactor me
	for i := 0; i < len(f); i++ {
		fname := f[i].Name
		fname = cleanName(fname)

		for x := 0; x < len(f[i].Tasks); x++ {
			task := f[i].Tasks[x]

			host, err := rg.hostBySlaveId(task.SlaveId)
			if err == nil && (task.State == "TASK_RUNNING") {

				tname := cleanName(task.Name)
				tail := fname + "." + domain + "."

				// hack - what to do?
				if task.Resources.Ports != "" {
					sports := yankPorts(task.Resources.Ports)

					// FIXME - 3 nested loops
					for s := 0; s < len(sports); s++ {
						var srvhost string = tname + "." + fname + "." + domain + ":" + sports[s]

						tcp := "_" + tname + "._tcp." + tail
						udp := "_" + tname + "._udp." + tail

						rg.insertRR(tcp, srvhost, "SRV")
						rg.insertRR(udp, srvhost, "SRV")
					}

				}

				arec := tname + "." + tail
				rg.insertRR(arec, host, "A")

			}
		}
	}

	rg.listenerRecord(listener, mname)
	rg.masterRecord(domain, masters, sj.Leader)
	return nil
}

// listenerRecord sets the A record for the mesos-dns server in case
// there is a request for it's hostname (eg: from SOA mname)
func (rg *RecordGenerator) listenerRecord(listener string, mname string) {
	if listener == "0.0.0.0" {
		rg.setFromLocal(listener, mname)
	} else if listener == "127.0.0.1" {
		rg.insertRR(mname, "127.0.0.1", "A")
	} else {
		rg.insertRR(mname, listener, "A")
	}
}

// masterRecord sets A records for the mesos masters and an A record
// for the leading master
func (rg *RecordGenerator) masterRecord(domain string, masters []string, leader string) {
	logging.VeryVerbose.Println("MasterRecord called " + leader)

	// create records for leader
	// A records
	h := strings.Split(leader, "@")
	if len(h) < 2 {
		logging.Error.Println(leader)
	}
	ip, port, err := getProto(h[1])
	if err != nil {
		logging.Error.Println(err)
	}
	arec := "leader." + domain + "."
	rg.insertRR(arec, ip, "A")
	arec = "master." + domain + "."
	rg.insertRR(arec, ip, "A")
	// SRV records
	tcp := "_leader._tcp." + domain + "."
	udp := "_leader._udp." + domain + "."
	host := "leader." + domain + ":" + port
	rg.insertRR(tcp, host, "SRV")
	rg.insertRR(udp, host, "SRV")
	tcp = "_master._tcp." + domain + "."
	udp = "_master._udp." + domain + "."
	host = "master." + domain + ":" + port
	rg.insertRR(tcp, host, "SRV")
	rg.insertRR(udp, host, "SRV")

	// if there is a list of masters, insert that as well
	for i := 0; i < len(masters); i++ {

		// skip leader
		if leader == masters[i] {
			continue
		}

		ip, port, err := getProto(masters[i])
		if err != nil {
			logging.Error.Println(err)
		}

		// A records (master and masterN)
		arec := "master." + domain + "."
		rg.insertRR(arec, ip, "A")
		arec = "master" + strconv.Itoa(i) + "." + domain + "."
		rg.insertRR(arec, ip, "A")

		// SRV records
		tcp := "_master._tcp." + domain + "."
		udp := "_master._udp." + domain + "."
		host := "master." + domain + ":" + port
		rg.insertRR(tcp, host, "SRV")
		rg.insertRR(udp, host, "SRV")
	}
}

// setFromLocal generates A records for each local interface we are
// listening on - if this causes problems you should explicitly set the
// listener address in config.json
func (rg *RecordGenerator) setFromLocal(host string, mname string) {

	ifaces, err := net.Interfaces()
	if err != nil {
		logging.Error.Println(err)
	}

	// handle err
	for _, i := range ifaces {

		addrs, err := i.Addrs()
		if err != nil {
			logging.Error.Println(err)
		}

		for _, addr := range addrs {
			var ip net.IP
			switch v := addr.(type) {
			case *net.IPNet:
				ip = v.IP
			case *net.IPAddr:
				ip = v.IP
			}

			if ip == nil || ip.IsLoopback() {
				continue
			}

			ip = ip.To4()
			if ip == nil {
				continue
			}

			rg.insertRR(mname, ip.String(), "A")
		}

	}
}

func stripHost(hostip string) string {
	return strings.Split(hostip, ":")[0]
}

// insertRR inserts host to name's map
// refactor me
func (rg *RecordGenerator) insertRR(name string, host string, rtype string) {
	logging.VeryVerbose.Println("[" + rtype + "]\t" + name + ": " + host)

	if rtype == "A" {
		if val, ok := rg.As[name]; ok {

			h := stripHost(host)
			for _, b := range val {
				if stripHost(b) == h {
					return
				}
			}

			rg.As[name] = append(val, host)
		} else {
			rg.As[name] = []string{host}
		}
	} else {
		if val, ok := rg.SRVs[name]; ok {
			rg.SRVs[name] = append(val, host)
		} else {
			rg.SRVs[name] = []string{host}
		}
	}
}
<|endoftext|>


<commit_before>package redirect

import (
	"fmt"
	"net/http"
)

// Redirect - stops following redirects after specified number of times
// compatible with http.Client
type Redirect struct {
	StopAfter int
}

// DefaultCheckRedirect - default check redirect function from http.Client
func (r *Redirect) DefaultCheckRedirect(req *http.Request, via []*http.Request) error {
	if len(via) >= r.StopAfter {
		return fmt.Errorf("Stopped after %d redirects")
	}
	return nil
}
<commit_msg>Fix redirect.go missing Errorf param<commit_after>package redirect

import (
	"fmt"
	"net/http"
)

// Redirect - stops following redirects after specified number of times
// compatible with http.Client
type Redirect struct {
	StopAfter int
}

// DefaultCheckRedirect - default check redirect function from http.Client
func (r *Redirect) DefaultCheckRedirect(req *http.Request, via []*http.Request) error {
	if len(via) >= r.StopAfter {
		return fmt.Errorf("Stopped after %d redirects", len(via))
	}
	return nil
}
<|endoftext|>


<commit_before>package mysql

import (
	"crypto/rand"
	"crypto/sha1"
	"encoding/binary"
	"fmt"
	"io"
	"runtime"
	"strings"

	"github.com/juju/errors"
	"github.com/siddontang/go/hack"
)

func Pstack() string {
	buf := make([]byte, 1024)
	n := runtime.Stack(buf, false)
	return string(buf[0:n])
}

func CalcPassword(scramble, password []byte) []byte {
	if len(password) == 0 {
		return nil
	}

	// stage1Hash = SHA1(password)
	crypt := sha1.New()
	crypt.Write(password)
	stage1 := crypt.Sum(nil)

	// scrambleHash = SHA1(scramble + SHA1(stage1Hash))
	// inner Hash
	crypt.Reset()
	crypt.Write(stage1)
	hash := crypt.Sum(nil)

	// outer Hash
	crypt.Reset()
	crypt.Write(scramble)
	crypt.Write(hash)
	scramble = crypt.Sum(nil)

	// token = scrambleHash XOR stage1Hash
	for i := range scramble {
		scramble[i] ^= stage1[i]
	}
	return scramble
}

func RandomBuf(size int) ([]byte, error) {
	buf := make([]byte, size)

	if _, err := io.ReadFull(rand.Reader, buf); err != nil {
		return nil, errors.Trace(err)
	}

	// avoid to generate '\0'
	for i, b := range buf {
		if uint8(b) == 0 {
			buf[i] = '0'
		}
	}

	return buf, nil
}

// little endian
func FixedLengthInt(buf []byte) uint64 {
	var num uint64 = 0
	for i, b := range buf {
		num |= uint64(b) << (uint(i) * 8)
	}
	return num
}

// big endian
func BFixedLengthInt(buf []byte) uint64 {
	var num uint64 = 0
	for i, b := range buf {
		num |= uint64(b) << (uint(len(buf)-i-1) * 8)
	}
	return num
}

func LengthEncodedInt(b []byte) (num uint64, isNull bool, n int) {
	switch b[0] {

	// 251: NULL
	case 0xfb:
		n = 1
		isNull = true
		return

	// 252: value of following 2
	case 0xfc:
		num = uint64(b[1]) | uint64(b[2])<<8
		n = 3
		return

	// 253: value of following 3
	case 0xfd:
		num = uint64(b[1]) | uint64(b[2])<<8 | uint64(b[3])<<16
		n = 4
		return

	// 254: value of following 8
	case 0xfe:
		num = uint64(b[1]) | uint64(b[2])<<8 | uint64(b[3])<<16 |
			uint64(b[4])<<24 | uint64(b[5])<<32 | uint64(b[6])<<40 |
			uint64(b[7])<<48 | uint64(b[8])<<56
		n = 9
		return
	}

	// 0-250: value of first byte
	num = uint64(b[0])
	n = 1
	return
}

func PutLengthEncodedInt(n uint64) []byte {
	switch {
	case n <= 250:
		return []byte{byte(n)}

	case n <= 0xffff:
		return []byte{0xfc, byte(n), byte(n >> 8)}

	case n <= 0xffffff:
		return []byte{0xfd, byte(n), byte(n >> 8), byte(n >> 16)}

	case n <= 0xffffffffffffffff:
		return []byte{0xfe, byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24),
			byte(n >> 32), byte(n >> 40), byte(n >> 48), byte(n >> 56)}
	}
	return nil
}

func LengthEnodedString(b []byte) ([]byte, bool, int, error) {
	// Get length
	num, isNull, n := LengthEncodedInt(b)
	if num < 1 {
		return nil, isNull, n, nil
	}

	n += int(num)

	// Check data length
	if len(b) >= n {
		return b[n-int(num) : n], false, n, nil
	}
	return nil, false, n, io.EOF
}

func SkipLengthEnodedString(b []byte) (int, error) {
	// Get length
	num, _, n := LengthEncodedInt(b)
	if num < 1 {
		return n, nil
	}

	n += int(num)

	// Check data length
	if len(b) >= n {
		return n, nil
	}
	return n, io.EOF
}

func PutLengthEncodedString(b []byte) []byte {
	data := make([]byte, 0, len(b)+9)
	data = append(data, PutLengthEncodedInt(uint64(len(b)))...)
	data = append(data, b...)
	return data
}

func Uint16ToBytes(n uint16) []byte {
	return []byte{
		byte(n),
		byte(n >> 8),
	}
}

func Uint32ToBytes(n uint32) []byte {
	return []byte{
		byte(n),
		byte(n >> 8),
		byte(n >> 16),
		byte(n >> 24),
	}
}

func Uint64ToBytes(n uint64) []byte {
	return []byte{
		byte(n),
		byte(n >> 8),
		byte(n >> 16),
		byte(n >> 24),
		byte(n >> 32),
		byte(n >> 40),
		byte(n >> 48),
		byte(n >> 56),
	}
}

func FormatBinaryDate(n int, data []byte) ([]byte, error) {
	switch n {
	case 0:
		return []byte("0000-00-00"), nil
	case 4:
		return []byte(fmt.Sprintf("%04d-%02d-%02d",
			binary.LittleEndian.Uint16(data[:2]),
			data[2],
			data[3])), nil
	default:
		return nil, errors.Errorf("invalid date packet length %d", n)
	}
}

func FormatBinaryDateTime(n int, data []byte) ([]byte, error) {
	switch n {
	case 0:
		return []byte("0000-00-00 00:00:00"), nil
	case 4:
		return []byte(fmt.Sprintf("%04d-%02d-%02d 00:00:00",
			binary.LittleEndian.Uint16(data[:2]),
			data[2],
			data[3])), nil
	case 7:
		return []byte(fmt.Sprintf(
			"%04d-%02d-%02d %02d:%02d:%02d",
			binary.LittleEndian.Uint16(data[:2]),
			data[2],
			data[3],
			data[4],
			data[5],
			data[6])), nil
	case 11:
		return []byte(fmt.Sprintf(
			"%04d-%02d-%02d %02d:%02d:%02d.%06d",
			binary.LittleEndian.Uint16(data[:2]),
			data[2],
			data[3],
			data[4],
			data[5],
			data[6],
			binary.LittleEndian.Uint32(data[7:11]))), nil
	default:
		return nil, errors.Errorf("invalid datetime packet length %d", n)
	}
}

func FormatBinaryTime(n int, data []byte) ([]byte, error) {
	if n == 0 {
		return []byte("0000-00-00"), nil
	}

	var sign byte
	if data[0] == 1 {
		sign = byte('-')
	}

	switch n {
	case 8:
		return []byte(fmt.Sprintf(
			"%c%02d:%02d:%02d",
			sign,
			uint16(data[1])*24+uint16(data[5]),
			data[6],
			data[7],
		)), nil
	case 12:
		return []byte(fmt.Sprintf(
			"%c%02d:%02d:%02d.%06d",
			sign,
			uint16(data[1])*24+uint16(data[5]),
			data[6],
			data[7],
			binary.LittleEndian.Uint32(data[8:12]),
		)), nil
	default:
		return nil, errors.Errorf("invalid time packet length %d", n)
	}
}

var (
	DONTESCAPE = byte(255)

	EncodeMap [256]byte
)

// only support utf-8
func Escape(sql string) string {
	dest := make([]byte, 0, 2*len(sql))

	for _, w := range hack.Slice(sql) {
		if c := EncodeMap[w]; c == DONTESCAPE {
			dest = append(dest, w)
		} else {
			dest = append(dest, '\\', c)
		}
	}

	return string(dest)
}

func GetNetProto(addr string) string {
	if strings.Contains(addr, "/") {
		return "unix"
	} else {
		return "tcp"
	}
}

// ErrorEqual returns a boolean indicating whether err1 is equal to err2.
func ErrorEqual(err1, err2 error) bool {
	e1 := errors.Cause(err1)
	e2 := errors.Cause(err2)

	if e1 == e2 {
		return true
	}

	if e1 == nil || e2 == nil {
		return e1 == e2
	}

	return e1.Error() == e2.Error()
}

var encodeRef = map[byte]byte{
	'\x00': '0',
	'\'':   '\'',
	'"':    '"',
	'\b':   'b',
	'\n':   'n',
	'\r':   'r',
	'\t':   't',
	26:     'Z', // ctl-Z
	'\\':   '\\',
}

func init() {
	for i := range EncodeMap {
		EncodeMap[i] = DONTESCAPE
	}
	for i := range EncodeMap {
		if to, ok := encodeRef[byte(i)]; ok {
			EncodeMap[byte(i)] = to
		}
	}
}
<commit_msg>disable log highlight<commit_after>package mysql

import (
	"crypto/rand"
	"crypto/sha1"
	"encoding/binary"
	"fmt"
	"io"
	"runtime"
	"strings"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/siddontang/go/hack"
)

func Pstack() string {
	buf := make([]byte, 1024)
	n := runtime.Stack(buf, false)
	return string(buf[0:n])
}

func CalcPassword(scramble, password []byte) []byte {
	if len(password) == 0 {
		return nil
	}

	// stage1Hash = SHA1(password)
	crypt := sha1.New()
	crypt.Write(password)
	stage1 := crypt.Sum(nil)

	// scrambleHash = SHA1(scramble + SHA1(stage1Hash))
	// inner Hash
	crypt.Reset()
	crypt.Write(stage1)
	hash := crypt.Sum(nil)

	// outer Hash
	crypt.Reset()
	crypt.Write(scramble)
	crypt.Write(hash)
	scramble = crypt.Sum(nil)

	// token = scrambleHash XOR stage1Hash
	for i := range scramble {
		scramble[i] ^= stage1[i]
	}
	return scramble
}

func RandomBuf(size int) ([]byte, error) {
	buf := make([]byte, size)

	if _, err := io.ReadFull(rand.Reader, buf); err != nil {
		return nil, errors.Trace(err)
	}

	// avoid to generate '\0'
	for i, b := range buf {
		if uint8(b) == 0 {
			buf[i] = '0'
		}
	}

	return buf, nil
}

// little endian
func FixedLengthInt(buf []byte) uint64 {
	var num uint64 = 0
	for i, b := range buf {
		num |= uint64(b) << (uint(i) * 8)
	}
	return num
}

// big endian
func BFixedLengthInt(buf []byte) uint64 {
	var num uint64 = 0
	for i, b := range buf {
		num |= uint64(b) << (uint(len(buf)-i-1) * 8)
	}
	return num
}

func LengthEncodedInt(b []byte) (num uint64, isNull bool, n int) {
	switch b[0] {

	// 251: NULL
	case 0xfb:
		n = 1
		isNull = true
		return

	// 252: value of following 2
	case 0xfc:
		num = uint64(b[1]) | uint64(b[2])<<8
		n = 3
		return

	// 253: value of following 3
	case 0xfd:
		num = uint64(b[1]) | uint64(b[2])<<8 | uint64(b[3])<<16
		n = 4
		return

	// 254: value of following 8
	case 0xfe:
		num = uint64(b[1]) | uint64(b[2])<<8 | uint64(b[3])<<16 |
			uint64(b[4])<<24 | uint64(b[5])<<32 | uint64(b[6])<<40 |
			uint64(b[7])<<48 | uint64(b[8])<<56
		n = 9
		return
	}

	// 0-250: value of first byte
	num = uint64(b[0])
	n = 1
	return
}

func PutLengthEncodedInt(n uint64) []byte {
	switch {
	case n <= 250:
		return []byte{byte(n)}

	case n <= 0xffff:
		return []byte{0xfc, byte(n), byte(n >> 8)}

	case n <= 0xffffff:
		return []byte{0xfd, byte(n), byte(n >> 8), byte(n >> 16)}

	case n <= 0xffffffffffffffff:
		return []byte{0xfe, byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24),
			byte(n >> 32), byte(n >> 40), byte(n >> 48), byte(n >> 56)}
	}
	return nil
}

func LengthEnodedString(b []byte) ([]byte, bool, int, error) {
	// Get length
	num, isNull, n := LengthEncodedInt(b)
	if num < 1 {
		return nil, isNull, n, nil
	}

	n += int(num)

	// Check data length
	if len(b) >= n {
		return b[n-int(num) : n], false, n, nil
	}
	return nil, false, n, io.EOF
}

func SkipLengthEnodedString(b []byte) (int, error) {
	// Get length
	num, _, n := LengthEncodedInt(b)
	if num < 1 {
		return n, nil
	}

	n += int(num)

	// Check data length
	if len(b) >= n {
		return n, nil
	}
	return n, io.EOF
}

func PutLengthEncodedString(b []byte) []byte {
	data := make([]byte, 0, len(b)+9)
	data = append(data, PutLengthEncodedInt(uint64(len(b)))...)
	data = append(data, b...)
	return data
}

func Uint16ToBytes(n uint16) []byte {
	return []byte{
		byte(n),
		byte(n >> 8),
	}
}

func Uint32ToBytes(n uint32) []byte {
	return []byte{
		byte(n),
		byte(n >> 8),
		byte(n >> 16),
		byte(n >> 24),
	}
}

func Uint64ToBytes(n uint64) []byte {
	return []byte{
		byte(n),
		byte(n >> 8),
		byte(n >> 16),
		byte(n >> 24),
		byte(n >> 32),
		byte(n >> 40),
		byte(n >> 48),
		byte(n >> 56),
	}
}

func FormatBinaryDate(n int, data []byte) ([]byte, error) {
	switch n {
	case 0:
		return []byte("0000-00-00"), nil
	case 4:
		return []byte(fmt.Sprintf("%04d-%02d-%02d",
			binary.LittleEndian.Uint16(data[:2]),
			data[2],
			data[3])), nil
	default:
		return nil, errors.Errorf("invalid date packet length %d", n)
	}
}

func FormatBinaryDateTime(n int, data []byte) ([]byte, error) {
	switch n {
	case 0:
		return []byte("0000-00-00 00:00:00"), nil
	case 4:
		return []byte(fmt.Sprintf("%04d-%02d-%02d 00:00:00",
			binary.LittleEndian.Uint16(data[:2]),
			data[2],
			data[3])), nil
	case 7:
		return []byte(fmt.Sprintf(
			"%04d-%02d-%02d %02d:%02d:%02d",
			binary.LittleEndian.Uint16(data[:2]),
			data[2],
			data[3],
			data[4],
			data[5],
			data[6])), nil
	case 11:
		return []byte(fmt.Sprintf(
			"%04d-%02d-%02d %02d:%02d:%02d.%06d",
			binary.LittleEndian.Uint16(data[:2]),
			data[2],
			data[3],
			data[4],
			data[5],
			data[6],
			binary.LittleEndian.Uint32(data[7:11]))), nil
	default:
		return nil, errors.Errorf("invalid datetime packet length %d", n)
	}
}

func FormatBinaryTime(n int, data []byte) ([]byte, error) {
	if n == 0 {
		return []byte("0000-00-00"), nil
	}

	var sign byte
	if data[0] == 1 {
		sign = byte('-')
	}

	switch n {
	case 8:
		return []byte(fmt.Sprintf(
			"%c%02d:%02d:%02d",
			sign,
			uint16(data[1])*24+uint16(data[5]),
			data[6],
			data[7],
		)), nil
	case 12:
		return []byte(fmt.Sprintf(
			"%c%02d:%02d:%02d.%06d",
			sign,
			uint16(data[1])*24+uint16(data[5]),
			data[6],
			data[7],
			binary.LittleEndian.Uint32(data[8:12]),
		)), nil
	default:
		return nil, errors.Errorf("invalid time packet length %d", n)
	}
}

var (
	DONTESCAPE = byte(255)

	EncodeMap [256]byte
)

// only support utf-8
func Escape(sql string) string {
	dest := make([]byte, 0, 2*len(sql))

	for _, w := range hack.Slice(sql) {
		if c := EncodeMap[w]; c == DONTESCAPE {
			dest = append(dest, w)
		} else {
			dest = append(dest, '\\', c)
		}
	}

	return string(dest)
}

func GetNetProto(addr string) string {
	if strings.Contains(addr, "/") {
		return "unix"
	} else {
		return "tcp"
	}
}

// ErrorEqual returns a boolean indicating whether err1 is equal to err2.
func ErrorEqual(err1, err2 error) bool {
	e1 := errors.Cause(err1)
	e2 := errors.Cause(err2)

	if e1 == e2 {
		return true
	}

	if e1 == nil || e2 == nil {
		return e1 == e2
	}

	return e1.Error() == e2.Error()
}

var encodeRef = map[byte]byte{
	'\x00': '0',
	'\'':   '\'',
	'"':    '"',
	'\b':   'b',
	'\n':   'n',
	'\r':   'r',
	'\t':   't',
	26:     'Z', // ctl-Z
	'\\':   '\\',
}

func init() {
	for i := range EncodeMap {
		EncodeMap[i] = DONTESCAPE
	}
	for i := range EncodeMap {
		if to, ok := encodeRef[byte(i)]; ok {
			EncodeMap[byte(i)] = to
		}
	}

	// Disable highlight by default
	log.SetHighlighting(false)
}
<|endoftext|>


<commit_before>package config

import (
	"io/ioutil"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"gopkg.in/yaml.v2"
)

// Define Application Settings Structure
type CassabonConfig struct {
	Logging struct {
		Logdir   string // Log Directory
		Loglevel string // Level to log at.
	}
	Cassandra struct {
		Hosts []string // List of hostnames or IP addresses of Cassandra ring
		Port  string   // Cassandra port
	}
	API struct {
		Address         string // HTTP API listens on this address
		Port            string // HTTP API listens on this port
		HealthCheckFile string // Location of healthcheck file.
	}
	Redis struct {
		Index RedisSettings // Settings for Redis Index
		Queue RedisSettings // Settings for Redis Queue
	}
	Carbon struct {
		Address  string   // Address for Carbon Receiver to listen on
		Port     string   // Port for Carbon Receiver to listen on
		Protocol string   // "tcp", "udp" or "both" are acceptable
		Peers    []string // All servers in the Cassabon array, as "ip:port"
	}
	Statsd   StatsdSettings
	Rollups  map[string]RollupSettings // Map of regex and rollups
	Channels struct {
		DataStoreChanLen  int // Length of the DataStore channel
		IndexStoreChanLen int // Length of the Indexer channel
		GopherChanLen     int // Length of the StatGopher channel
	}
	Parameters struct {
		Listener struct {
			TCPTimeout int
			UDPTimeout int
		}
	}
}

// Definition of each rollup
type RollupSettings struct {
	Retention   []string
	Aggregation string
}

// Redis struct for redis connection information
type RedisSettings struct {
	Sentinel bool     // True if sentinel, false if standalone.
	Addr     []string // List of addresses in host:port format
	DB       int64    // Redis DB number for the index.
	Pwd      string   // Password for Redis.
	Master   string   // Master config name for sentinel settings.
}

type StatsdSettings struct {
	Host   string // Host or IP address of statsd server
	Port   string // Port that statsd server listens on
	Events struct {
		ReceiveOK struct {
			Key        string
			SampleRate float32
		}
		ReceiveFail struct {
			Key        string
			SampleRate float32
		}
	}
}

// rawCassabonConfig is the decoded YAML from the configuration file.
var rawCassabonConfig *CassabonConfig

func ReadConfigurationFile(configFile string, haveLogger bool) error {

	// Load config file.
	// This happens before the logger is initialized because we also read in
	// logger configuration, causing a cycle. So, we panic on errors.
	yamlConfig, err := ioutil.ReadFile(configFile)
	if err != nil {
		if haveLogger {
			G.Log.System.LogError("Configuration not available: %v", err)
			return err
		} else {
			panic(err)
		}
	}

	// Unmarshal config file into config struct.
	err = yaml.Unmarshal(yamlConfig, &rawCassabonConfig)
	if err != nil {
		if haveLogger {
			G.Log.System.LogError("Configuration parse error: %v", err)
			return err
		} else {
			panic(err)
		}
	}
	return nil
}

func ParseStartupValues() {

	// Fill in arguments not provided on the command line.
	if G.Log.Logdir == "" {
		G.Log.Logdir = rawCassabonConfig.Logging.Logdir
	}

	// We need the log level at startup to initialize the loggers.
	G.Log.Loglevel = rawCassabonConfig.Logging.Loglevel

	// Copy in the statsd configuration.
	G.Statsd = rawCassabonConfig.Statsd

	// Copy in and sanitize the channel lengths.
	G.Channels.DataStoreChanLen = rawCassabonConfig.Channels.DataStoreChanLen
	if G.Channels.DataStoreChanLen < 10 {
		G.Channels.DataStoreChanLen = 10
	}
	if G.Channels.DataStoreChanLen > 1000 {
		G.Channels.DataStoreChanLen = 1000
	}
	G.Channels.IndexStoreChanLen = rawCassabonConfig.Channels.IndexStoreChanLen
	if G.Channels.IndexStoreChanLen < 10 {
		G.Channels.IndexStoreChanLen = 10
	}
	if G.Channels.IndexStoreChanLen > 1000 {
		G.Channels.IndexStoreChanLen = 1000
	}
	G.Channels.GopherChanLen = rawCassabonConfig.Channels.GopherChanLen
	if G.Channels.GopherChanLen < 10 {
		G.Channels.GopherChanLen = 10
	}
	if G.Channels.GopherChanLen > 1000 {
		G.Channels.GopherChanLen = 1000
	}
}

func ParseRefreshableValues() {

	// We can dynamically change the logging level.
	G.Log.Loglevel = rawCassabonConfig.Logging.Loglevel

	// Copy in the addresses of the services we offer.
	G.API.Address = rawCassabonConfig.API.Address
	G.API.Port = rawCassabonConfig.API.Port
	G.API.HealthCheckFile = rawCassabonConfig.API.HealthCheckFile

	G.Carbon.Address = rawCassabonConfig.Carbon.Address
	G.Carbon.Port = rawCassabonConfig.Carbon.Port
	G.Carbon.Protocol = rawCassabonConfig.Carbon.Protocol
	G.Carbon.Peers = rawCassabonConfig.Carbon.Peers
	// Ensure a canonical order for peers, as we will allocate metrics paths by index.
	sort.Strings(G.Carbon.Peers)

	// Copy in the addresses of the services we use.
	G.Cassandra = rawCassabonConfig.Cassandra
	G.Redis.Index = rawCassabonConfig.Redis.Index
	G.Redis.Queue = rawCassabonConfig.Redis.Queue

	// Copy in and sanitize the Listener internal parameters.
	G.Parameters.Listener.TCPTimeout = rawCassabonConfig.Parameters.Listener.TCPTimeout
	if G.Parameters.Listener.TCPTimeout < 1 {
		G.Parameters.Listener.TCPTimeout = 1
	}
	if G.Parameters.Listener.TCPTimeout > 30 {
		G.Parameters.Listener.TCPTimeout = 30
	}
	G.Parameters.Listener.UDPTimeout = rawCassabonConfig.Parameters.Listener.UDPTimeout
	if G.Parameters.Listener.UDPTimeout < 1 {
		G.Parameters.Listener.UDPTimeout = 1
	}
	if G.Parameters.Listener.UDPTimeout > 30 {
		G.Parameters.Listener.UDPTimeout = 30
	}

	// Validate, copy in and normalize the rollup definitions.
	G.RollupPriority = make([]string, 0, len(rawCassabonConfig.Rollups))
	G.Rollup = make(map[string]RollupDef)

	var method RollupMethod
	var window, retention time.Duration
	var err error
	var intRetention int64
	var rd *RollupDef
	re := regexp.MustCompile("([0-9]+)([a-z])")

	// Inspect each rollup found in the configuration.
	// Note: YAML decode has already folded duplicate path expressions.
	for expression, v := range rawCassabonConfig.Rollups {

		// Validate and decode the aggregation method.
		switch strings.ToLower(v.Aggregation) {
		case "average":
			method = AVERAGE
		case "max":
			method = MAX
		case "min":
			method = MIN
		case "sum":
			method = SUM
		default:
			G.Log.System.LogWarn("Invalid aggregation method for \"%s\": %s", expression, v.Aggregation)
			continue
		}

		// Build up the rollup definitions for this regular expression.
		rd = new(RollupDef)
		rd.Method = method
		rd.Windows = make([]RollupWindow, 0)
		if expression != CATCHALL_EXPRESSION {
			if re, err := regexp.Compile(expression); err == nil {
				rd.Expression = re
			} else {
				G.Log.System.LogWarn("Malformed regular expression for \"%s\": %v", expression, err)
				continue
			}
		}

		// Parse and validate each window:retention pair.
		for _, s := range v.Retention {

			// Split the value on the colon between the parts.
			couplet := strings.Split(s, ":")
			if len(couplet) != 2 {
				G.Log.System.LogWarn("Malformed definition for \"%s\": %s", expression, s)
				continue
			}

			// Convert the window to a time.Duration.
			if window, err = time.ParseDuration(couplet[0]); err != nil {
				G.Log.System.LogWarn("Malformed window for \"%s\": %s %s", expression, s, couplet[0])
				continue
			}

			// Don't permit windows shorter than 1 second.
			if window < time.Second {
				G.Log.System.LogWarn("Duration less than minimum 1 second for \"%s\": %v", expression, window)
				continue
			}

			// Convert the retention to a time.Duration (max: 720 years).
			// ParseDuration doesn't handle anything longer than hours, so do it manually.
			matches := re.FindStringSubmatch(couplet[1]) // "1d" -> [ 1d 1 d ]
			if len(matches) != 3 {
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}
			if intRetention, err = strconv.ParseInt(matches[1], 10, 64); err != nil {
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}
			switch matches[2] {
			case "m":
				retention = time.Minute * time.Duration(intRetention)
			case "h":
				retention = time.Hour * time.Duration(intRetention)
			case "d":
				retention = time.Hour * 24 * time.Duration(intRetention)
			case "w":
				retention = time.Hour * 24 * 7 * time.Duration(intRetention)
			case "y":
				retention = time.Hour * 24 * 365 * time.Duration(intRetention)
			default:
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}

			// Append to the rollups for this expression.
			rd.Windows = append(rd.Windows, RollupWindow{window, retention})
		}

		// If any of the rollup window definitions were valid, add expression to the list.
		if len(rd.Windows) > 0 {

			// Sort windows into ascending duration order, and validate.
			sort.Sort(ByWindow(rd.Windows))
			shortestDuration := rd.Windows[0].Window
			allExactMultiples := true
			for i, v := range rd.Windows {
				if i > 0 {
					remainder := v.Window % shortestDuration
					if remainder != 0 {
						G.Log.System.LogWarn(
							"Next duration is not a multiple for \"%s\": %v %% %v remainder is %v",
							expression, v.Window, shortestDuration, remainder)
						allExactMultiples = false
					}
				}
			}

			// If all durations are exact multiples of the shortest duration, save this expression.
			if allExactMultiples {
				G.Rollup[expression] = *rd
				G.RollupPriority = append(G.RollupPriority, expression)
			} else {
				G.Log.System.LogWarn("Rollup expression rejected due to previous errors: \"%s\"", expression)
			}
		}
	}

	// Sort the path expressions into priority order.
	sort.Sort(ByPriority(G.RollupPriority))
}

// ByWindow is used to sort retention definitions by window duration.
type ByWindow []RollupWindow

// Implementation of sort.Interface.
func (w ByWindow) Len() int {
	return len(w)
}
func (w ByWindow) Swap(i, j int) {
	w[i], w[j] = w[j], w[i]
}
func (w ByWindow) Less(i, j int) bool {
	return w[i].Window < w[j].Window
}

// ByPriority is used to provide a consistent order for processing path expressions.
type ByPriority []string

// Implementation of sort.Interface.
func (p ByPriority) Len() int {
	return len(p)
}
func (p ByPriority) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}
func (p ByPriority) Less(i, j int) bool {

	// "default" is always last in priority.
	if p[i] == CATCHALL_EXPRESSION {
		return false
	}
	if p[j] == CATCHALL_EXPRESSION {
		return true
	}

	// Longer strings sort earlier.
	if len(p[i]) > len(p[j]) {
		return true
	} else if len(p[i]) < len(p[j]) {
		return false
	}

	// Same-length strings are ordered lexically.
	return p[i] < p[j]
}
<commit_msg>Ensure that local address is in peer list<commit_after>package config

import (
	"io/ioutil"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"gopkg.in/yaml.v2"
)

// Define Application Settings Structure
type CassabonConfig struct {
	Logging struct {
		Logdir   string // Log Directory
		Loglevel string // Level to log at.
	}
	Cassandra struct {
		Hosts []string // List of hostnames or IP addresses of Cassandra ring
		Port  string   // Cassandra port
	}
	API struct {
		Address         string // HTTP API listens on this address
		Port            string // HTTP API listens on this port
		HealthCheckFile string // Location of healthcheck file.
	}
	Redis struct {
		Index RedisSettings // Settings for Redis Index
		Queue RedisSettings // Settings for Redis Queue
	}
	Carbon struct {
		Address  string   // Address for Carbon Receiver to listen on
		Port     string   // Port for Carbon Receiver to listen on
		Protocol string   // "tcp", "udp" or "both" are acceptable
		Peers    []string // All servers in the Cassabon array, as "ip:port"
	}
	Statsd   StatsdSettings
	Rollups  map[string]RollupSettings // Map of regex and rollups
	Channels struct {
		DataStoreChanLen  int // Length of the DataStore channel
		IndexStoreChanLen int // Length of the Indexer channel
		GopherChanLen     int // Length of the StatGopher channel
	}
	Parameters struct {
		Listener struct {
			TCPTimeout int
			UDPTimeout int
		}
	}
}

// Definition of each rollup
type RollupSettings struct {
	Retention   []string
	Aggregation string
}

// Redis struct for redis connection information
type RedisSettings struct {
	Sentinel bool     // True if sentinel, false if standalone.
	Addr     []string // List of addresses in host:port format
	DB       int64    // Redis DB number for the index.
	Pwd      string   // Password for Redis.
	Master   string   // Master config name for sentinel settings.
}

type StatsdSettings struct {
	Host   string // Host or IP address of statsd server
	Port   string // Port that statsd server listens on
	Events struct {
		ReceiveOK struct {
			Key        string
			SampleRate float32
		}
		ReceiveFail struct {
			Key        string
			SampleRate float32
		}
	}
}

// rawCassabonConfig is the decoded YAML from the configuration file.
var rawCassabonConfig *CassabonConfig

func ReadConfigurationFile(configFile string, haveLogger bool) error {

	// Load config file.
	// This happens before the logger is initialized because we also read in
	// logger configuration, causing a cycle. So, we panic on errors.
	yamlConfig, err := ioutil.ReadFile(configFile)
	if err != nil {
		if haveLogger {
			G.Log.System.LogError("Configuration not available: %v", err)
			return err
		} else {
			panic(err)
		}
	}

	// Unmarshal config file into config struct.
	err = yaml.Unmarshal(yamlConfig, &rawCassabonConfig)
	if err != nil {
		if haveLogger {
			G.Log.System.LogError("Configuration parse error: %v", err)
			return err
		} else {
			panic(err)
		}
	}
	return nil
}

func ParseStartupValues() {

	// Fill in arguments not provided on the command line.
	if G.Log.Logdir == "" {
		G.Log.Logdir = rawCassabonConfig.Logging.Logdir
	}

	// We need the log level at startup to initialize the loggers.
	G.Log.Loglevel = rawCassabonConfig.Logging.Loglevel

	// Copy in the statsd configuration.
	G.Statsd = rawCassabonConfig.Statsd

	// Copy in and sanitize the channel lengths.
	G.Channels.DataStoreChanLen = rawCassabonConfig.Channels.DataStoreChanLen
	if G.Channels.DataStoreChanLen < 10 {
		G.Channels.DataStoreChanLen = 10
	}
	if G.Channels.DataStoreChanLen > 1000 {
		G.Channels.DataStoreChanLen = 1000
	}
	G.Channels.IndexStoreChanLen = rawCassabonConfig.Channels.IndexStoreChanLen
	if G.Channels.IndexStoreChanLen < 10 {
		G.Channels.IndexStoreChanLen = 10
	}
	if G.Channels.IndexStoreChanLen > 1000 {
		G.Channels.IndexStoreChanLen = 1000
	}
	G.Channels.GopherChanLen = rawCassabonConfig.Channels.GopherChanLen
	if G.Channels.GopherChanLen < 10 {
		G.Channels.GopherChanLen = 10
	}
	if G.Channels.GopherChanLen > 1000 {
		G.Channels.GopherChanLen = 1000
	}
}

func ParseRefreshableValues() {

	// We can dynamically change the logging level.
	G.Log.Loglevel = rawCassabonConfig.Logging.Loglevel

	// Copy in the addresses of the services we offer.
	G.API.Address = rawCassabonConfig.API.Address
	G.API.Port = rawCassabonConfig.API.Port
	G.API.HealthCheckFile = rawCassabonConfig.API.HealthCheckFile

	G.Carbon.Address = rawCassabonConfig.Carbon.Address
	G.Carbon.Port = rawCassabonConfig.Carbon.Port
	G.Carbon.Protocol = rawCassabonConfig.Carbon.Protocol
	G.Carbon.Peers = rawCassabonConfig.Carbon.Peers
	// Ensure a canonical order for peers, as we will allocate metrics paths by index.
	sort.Strings(G.Carbon.Peers)
	// Ensure that the local address:port is in the peer list.
	localHostPort := G.Carbon.Address + ":" + G.Carbon.Port
	for _, v := range G.Carbon.Peers {
		if v == localHostPort {
			localHostPort = ""
			break
		}
	}
	if localHostPort != "" {
		G.Log.System.LogFatal("Local host:port %s is not in peer list: %v", localHostPort, G.Carbon.Peers)
		os.Exit(2)
	}

	// Copy in the addresses of the services we use.
	G.Cassandra = rawCassabonConfig.Cassandra
	G.Redis.Index = rawCassabonConfig.Redis.Index
	G.Redis.Queue = rawCassabonConfig.Redis.Queue

	// Copy in and sanitize the Listener internal parameters.
	G.Parameters.Listener.TCPTimeout = rawCassabonConfig.Parameters.Listener.TCPTimeout
	if G.Parameters.Listener.TCPTimeout < 1 {
		G.Parameters.Listener.TCPTimeout = 1
	}
	if G.Parameters.Listener.TCPTimeout > 30 {
		G.Parameters.Listener.TCPTimeout = 30
	}
	G.Parameters.Listener.UDPTimeout = rawCassabonConfig.Parameters.Listener.UDPTimeout
	if G.Parameters.Listener.UDPTimeout < 1 {
		G.Parameters.Listener.UDPTimeout = 1
	}
	if G.Parameters.Listener.UDPTimeout > 30 {
		G.Parameters.Listener.UDPTimeout = 30
	}

	// Validate, copy in and normalize the rollup definitions.
	G.RollupPriority = make([]string, 0, len(rawCassabonConfig.Rollups))
	G.Rollup = make(map[string]RollupDef)

	var method RollupMethod
	var window, retention time.Duration
	var err error
	var intRetention int64
	var rd *RollupDef
	re := regexp.MustCompile("([0-9]+)([a-z])")

	// Inspect each rollup found in the configuration.
	// Note: YAML decode has already folded duplicate path expressions.
	for expression, v := range rawCassabonConfig.Rollups {

		// Validate and decode the aggregation method.
		switch strings.ToLower(v.Aggregation) {
		case "average":
			method = AVERAGE
		case "max":
			method = MAX
		case "min":
			method = MIN
		case "sum":
			method = SUM
		default:
			G.Log.System.LogWarn("Invalid aggregation method for \"%s\": %s", expression, v.Aggregation)
			continue
		}

		// Build up the rollup definitions for this regular expression.
		rd = new(RollupDef)
		rd.Method = method
		rd.Windows = make([]RollupWindow, 0)
		if expression != CATCHALL_EXPRESSION {
			if re, err := regexp.Compile(expression); err == nil {
				rd.Expression = re
			} else {
				G.Log.System.LogWarn("Malformed regular expression for \"%s\": %v", expression, err)
				continue
			}
		}

		// Parse and validate each window:retention pair.
		for _, s := range v.Retention {

			// Split the value on the colon between the parts.
			couplet := strings.Split(s, ":")
			if len(couplet) != 2 {
				G.Log.System.LogWarn("Malformed definition for \"%s\": %s", expression, s)
				continue
			}

			// Convert the window to a time.Duration.
			if window, err = time.ParseDuration(couplet[0]); err != nil {
				G.Log.System.LogWarn("Malformed window for \"%s\": %s %s", expression, s, couplet[0])
				continue
			}

			// Don't permit windows shorter than 1 second.
			if window < time.Second {
				G.Log.System.LogWarn("Duration less than minimum 1 second for \"%s\": %v", expression, window)
				continue
			}

			// Convert the retention to a time.Duration (max: 720 years).
			// ParseDuration doesn't handle anything longer than hours, so do it manually.
			matches := re.FindStringSubmatch(couplet[1]) // "1d" -> [ 1d 1 d ]
			if len(matches) != 3 {
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}
			if intRetention, err = strconv.ParseInt(matches[1], 10, 64); err != nil {
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}
			switch matches[2] {
			case "m":
				retention = time.Minute * time.Duration(intRetention)
			case "h":
				retention = time.Hour * time.Duration(intRetention)
			case "d":
				retention = time.Hour * 24 * time.Duration(intRetention)
			case "w":
				retention = time.Hour * 24 * 7 * time.Duration(intRetention)
			case "y":
				retention = time.Hour * 24 * 365 * time.Duration(intRetention)
			default:
				G.Log.System.LogWarn("Malformed retention for \"%s\": %s %s", expression, s, couplet[1])
				continue
			}

			// Append to the rollups for this expression.
			rd.Windows = append(rd.Windows, RollupWindow{window, retention})
		}

		// If any of the rollup window definitions were valid, add expression to the list.
		if len(rd.Windows) > 0 {

			// Sort windows into ascending duration order, and validate.
			sort.Sort(ByWindow(rd.Windows))
			shortestDuration := rd.Windows[0].Window
			allExactMultiples := true
			for i, v := range rd.Windows {
				if i > 0 {
					remainder := v.Window % shortestDuration
					if remainder != 0 {
						G.Log.System.LogWarn(
							"Next duration is not a multiple for \"%s\": %v %% %v remainder is %v",
							expression, v.Window, shortestDuration, remainder)
						allExactMultiples = false
					}
				}
			}

			// If all durations are exact multiples of the shortest duration, save this expression.
			if allExactMultiples {
				G.Rollup[expression] = *rd
				G.RollupPriority = append(G.RollupPriority, expression)
			} else {
				G.Log.System.LogWarn("Rollup expression rejected due to previous errors: \"%s\"", expression)
			}
		}
	}

	// Sort the path expressions into priority order.
	sort.Sort(ByPriority(G.RollupPriority))
}

// ByWindow is used to sort retention definitions by window duration.
type ByWindow []RollupWindow

// Implementation of sort.Interface.
func (w ByWindow) Len() int {
	return len(w)
}
func (w ByWindow) Swap(i, j int) {
	w[i], w[j] = w[j], w[i]
}
func (w ByWindow) Less(i, j int) bool {
	return w[i].Window < w[j].Window
}

// ByPriority is used to provide a consistent order for processing path expressions.
type ByPriority []string

// Implementation of sort.Interface.
func (p ByPriority) Len() int {
	return len(p)
}
func (p ByPriority) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}
func (p ByPriority) Less(i, j int) bool {

	// "default" is always last in priority.
	if p[i] == CATCHALL_EXPRESSION {
		return false
	}
	if p[j] == CATCHALL_EXPRESSION {
		return true
	}

	// Longer strings sort earlier.
	if len(p[i]) > len(p[j]) {
		return true
	} else if len(p[i]) < len(p[j]) {
		return false
	}

	// Same-length strings are ordered lexically.
	return p[i] < p[j]
}
<|endoftext|>


<commit_before>// Command mealplanner is the main entry point of the application. It simply
// runs the *mux.Router provided by mphandlers.CreateMux() as an HTTP server.
package main

import (
	"flag"
	"fmt"
	"github.com/kierdavis/mealplanner/mpdb"
	"github.com/kierdavis/mealplanner/mphandlers"
	"net/http"
	"os"

	_ "github.com/go-sql-driver/mysql"
)

var (
	dbSource = flag.String("dbsource", "", "database source, in the form USER:PASS@unix(/PATH/TO/SOCKET)/DB or USER:PASS@tcp(HOST:PORT)/DB")
	host     = flag.String("host", "", "hostname to listen on")
	port     = flag.Int("port", 8080, "port to listen on")
	debug    = flag.Bool("debug", false, "debug mode")
	testdata = flag.Bool("testdata", false, "clear the database and insert test data")
)

func main() {
	flag.Parse()

	mpdb.Source = *dbSource
	if mpdb.Source == "" {
		fmt.Fprintf(os.Stderr, "Please specify a non-empty -dbsource flag.\n")
		os.Exit(1)
	}

	err := mpdb.InitDB(*testdata)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Database Error: %s\n", err)
		os.Exit(1)
	}

	listenAddr := fmt.Sprintf("%s:%d", *host, *port)
	m := mphandlers.CreateMux()

	app := http.Handler(m)
	if *debug {
		app = mphandlers.LoggingHandler{app}

		fmt.Printf("Listening on %s\n", listenAddr)
	}

	err = http.ListenAndServe(listenAddr, app)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Server Error: %s\n", err)
		os.Exit(1)
	}
}
<commit_msg>Allow use of env var for db source<commit_after>// Command mealplanner is the main entry point of the application. It simply
// runs the *mux.Router provided by mphandlers.CreateMux() as an HTTP server.
package main

import (
	"flag"
	"fmt"
	"github.com/kierdavis/mealplanner/mpdb"
	"github.com/kierdavis/mealplanner/mphandlers"
	"net/http"
	"os"

	_ "github.com/go-sql-driver/mysql"
)

var (
	dbSource = flag.String("dbsource", "", "database source, in the form USER:PASS@unix(/PATH/TO/SOCKET)/DB or USER:PASS@tcp(HOST:PORT)/DB")
	host     = flag.String("host", "", "hostname to listen on")
	port     = flag.Int("port", 8080, "port to listen on")
	debug    = flag.Bool("debug", false, "debug mode")
	testdata = flag.Bool("testdata", false, "clear the database and insert test data")
)

func main() {
	flag.Parse()

	source := *dbSource
	if source == "" {
		source = os.Getenv("MPDBSOURCE")
		if source == "" {
			fmt.Fprintf(os.Stderr, "Please specify a non-empty -dbsource flag or set the MPDBSOURCE environment variable.\n")
			os.Exit(1)
		}
	}

	mpdb.Source = source

	err := mpdb.InitDB(*testdata)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Database Error: %s\n", err)
		os.Exit(1)
	}

	listenAddr := fmt.Sprintf("%s:%d", *host, *port)
	m := mphandlers.CreateMux()

	app := http.Handler(m)
	if *debug {
		app = mphandlers.LoggingHandler{app}

		fmt.Printf("Listening on %s\n", listenAddr)
	}

	err = http.ListenAndServe(listenAddr, app)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Server Error: %s\n", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

   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 com

import (
	"errors"
	"log"
	"os"
	"path/filepath"

	"github.com/admpub/fsnotify"
)

var (
	DefaultMonitor      = NewMonitor()
	MonitorEventEmptyFn = func(string) {}
)

func NewMonitor() *MonitorEvent {
	return &MonitorEvent{
		Create:  MonitorEventEmptyFn,
		Delete:  MonitorEventEmptyFn,
		Modify:  MonitorEventEmptyFn,
		Chmod:   MonitorEventEmptyFn,
		Rename:  MonitorEventEmptyFn,
		filters: []func(string) bool{},
	}
}

//MonitorEvent 监控事件函数
type MonitorEvent struct {
	//文件事件
	Create func(string) //创建
	Delete func(string) //删除（包含文件夹和文件。因为已经删除，无法确定是文件夹还是文件）
	Modify func(string) //修改（包含修改权限。如果是文件夹，则内部的文件被更改也会触发此事件）
	Chmod  func(string) //修改权限（windows不支持）
	Rename func(string) //重命名

	//其它
	Channel chan bool //管道
	Debug   bool
	watcher *fsnotify.Watcher
	filters []func(string) bool
}

func (m *MonitorEvent) AddFilter(args ...func(string) bool) *MonitorEvent {
	if m.filters == nil {
		m.filters = []func(string) bool{}
	}
	m.filters = append(m.filters, args...)
	return m
}

func (m *MonitorEvent) SetFilters(args ...func(string) bool) *MonitorEvent {
	m.filters = args
	return m
}

func (m *MonitorEvent) Watch(args ...func(string) bool) *MonitorEvent {
	m.SetFilters(args...)
	go func() {
		m.backendListen()
		<-m.Channel
	}()
	return m
}

func (m *MonitorEvent) Close() error {
	if m.Channel != nil {
		close(m.Channel)
	}
	if m.watcher != nil {
		return m.watcher.Close()
	}
	return nil
}

func (m *MonitorEvent) Watcher() *fsnotify.Watcher {
	if m.watcher == nil {
		var err error
		m.watcher, err = fsnotify.NewWatcher()
		m.Channel = make(chan bool)
		m.filters = []func(string) bool{}
		if err != nil {
			log.Panic(err)
		}
	}
	return m.watcher
}

func (m *MonitorEvent) backendListen() *MonitorEvent {
	go m.listen()
	return m
}

func (m *MonitorEvent) AddDir(dir string) error {
	f, err := os.Stat(dir)
	if err != nil {
		return err
	}
	if !f.IsDir() {
		return errors.New(dir + ` is not dir.`)
	}
	err = filepath.Walk(dir, func(f string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			if m.Debug {
				log.Println(`[Monitor]`, `Add Watch:`, f)
			}
			return m.Watcher().Add(f)
		}
		return nil
	})
	return err
}

func (m *MonitorEvent) AddFile(file string) error {
	return m.Watcher().Add(file)
}

func (m *MonitorEvent) Remove(fileOrDir string) error {
	if m.watcher != nil {
		return m.watcher.Remove(fileOrDir)
	}
	return nil
}

func (m *MonitorEvent) listen() {
	for {
		watcher := m.Watcher()
		select {
		case ev := <-watcher.Events:
			if m.Debug {
				log.Println(`[Monitor]`, `Trigger Event:`, ev)
			}
			if m.filters != nil {
				var skip bool
				for _, filter := range m.filters {
					if !filter(ev.Name) {
						skip = true
						break
					}
				}
				if skip {
					break
				}
			}
			switch {
			case ev.Op&fsnotify.Create == fsnotify.Create:
				if m.IsDir(ev.Name) {
					watcher.Add(ev.Name)
				}
				if m.Create != nil {
					m.Create(ev.Name)
				}
			case ev.Op&fsnotify.Remove == fsnotify.Remove:
				if m.IsDir(ev.Name) {
					watcher.Remove(ev.Name)
				}
				if m.Delete != nil {
					m.Delete(ev.Name)
				}
			case ev.Op&fsnotify.Write == fsnotify.Write:
				if m.Modify != nil {
					m.Modify(ev.Name)
				}
			case ev.Op&fsnotify.Rename == fsnotify.Rename:
				watcher.Remove(ev.Name)
				if m.Rename != nil {
					m.Rename(ev.Name)
				}
			case ev.Op&fsnotify.Chmod == fsnotify.Chmod:
				if m.Chmod != nil {
					m.Chmod(ev.Name)
				}
			}
		case err := <-watcher.Errors:
			if err != nil {
				log.Println("Watcher error:", err)
			}
		}
	}
}

func (m *MonitorEvent) IsDir(path string) bool {
	d, e := os.Stat(path)
	if e != nil {
		return false
	}
	return d.IsDir()
}

//Monitor 文件监测
func Monitor(rootDir string, callback *MonitorEvent, args ...func(string) bool) error {
	watcher := callback.Watcher()
	defer watcher.Close()
	callback.Watch(args...)
	err := callback.AddDir(rootDir)
	if err != nil {
		callback.Close()
		return err
	}
	return nil
}
<commit_msg>update<commit_after>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

   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 com

import (
	"errors"
	"log"
	"os"
	"path/filepath"

	"github.com/admpub/fsnotify"
)

var (
	DefaultMonitor      = NewMonitor()
	MonitorEventEmptyFn = func(string) {}
)

func NewMonitor() *MonitorEvent {
	return &MonitorEvent{
		Create:  MonitorEventEmptyFn,
		Delete:  MonitorEventEmptyFn,
		Modify:  MonitorEventEmptyFn,
		Chmod:   MonitorEventEmptyFn,
		Rename:  MonitorEventEmptyFn,
		filters: []func(string) bool{},
	}
}

//MonitorEvent 监控事件函数
type MonitorEvent struct {
	//文件事件
	Create func(string) //创建
	Delete func(string) //删除（包含文件夹和文件。因为已经删除，无法确定是文件夹还是文件）
	Modify func(string) //修改（包含修改权限。如果是文件夹，则内部的文件被更改也会触发此事件）
	Chmod  func(string) //修改权限（windows不支持）
	Rename func(string) //重命名

	//其它
	Channel chan bool //管道
	Debug   bool
	watcher *fsnotify.Watcher
	filters []func(string) bool
}

func (m *MonitorEvent) AddFilter(args ...func(string) bool) *MonitorEvent {
	if m.filters == nil {
		m.filters = []func(string) bool{}
	}
	m.filters = append(m.filters, args...)
	return m
}

func (m *MonitorEvent) SetFilters(args ...func(string) bool) *MonitorEvent {
	m.filters = args
	return m
}

func (m *MonitorEvent) Watch(args ...func(string) bool) *MonitorEvent {
	m.SetFilters(args...)
	go func() {
		m.backendListen()
		<-m.Channel
	}()
	return m
}

func (m *MonitorEvent) Close() error {
	if m.Channel != nil {
		close(m.Channel)
	}
	if m.watcher != nil {
		return m.watcher.Close()
	}
	return nil
}

func (m *MonitorEvent) Watcher() *fsnotify.Watcher {
	if m.watcher == nil {
		var err error
		m.watcher, err = fsnotify.NewWatcher()
		m.Channel = make(chan bool)
		m.filters = []func(string) bool{}
		if err != nil {
			log.Panic(err)
		}
	}
	return m.watcher
}

func (m *MonitorEvent) backendListen() *MonitorEvent {
	go m.listen()
	return m
}

func (m *MonitorEvent) AddDir(dir string) error {
	f, err := os.Stat(dir)
	if err != nil {
		return err
	}
	if !f.IsDir() {
		return errors.New(dir + ` is not dir.`)
	}
	err = filepath.Walk(dir, func(f string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			if m.Debug {
				log.Println(`[Monitor]`, `Add Watch:`, f)
			}
			return m.Watcher().Add(f)
		}
		return nil
	})
	return err
}

func (m *MonitorEvent) AddFile(file string) error {
	return m.Watcher().Add(file)
}

func (m *MonitorEvent) Remove(fileOrDir string) error {
	if m.watcher != nil {
		return m.watcher.Remove(fileOrDir)
	}
	return nil
}

func (m *MonitorEvent) listen() {
	for {
		watcher := m.Watcher()
		select {
		case ev, ok := <-watcher.Events:
			if !ok {
				return
			}
			if m.Debug {
				log.Println(`[Monitor]`, `Trigger Event:`, ev)
			}
			if m.filters != nil {
				var skip bool
				for _, filter := range m.filters {
					if !filter(ev.Name) {
						skip = true
						break
					}
				}
				if skip {
					break
				}
			}
			switch {
			case ev.Op&fsnotify.Create == fsnotify.Create:
				if m.IsDir(ev.Name) {
					watcher.Add(ev.Name)
				}
				if m.Create != nil {
					m.Create(ev.Name)
				}
			case ev.Op&fsnotify.Remove == fsnotify.Remove:
				if m.IsDir(ev.Name) {
					watcher.Remove(ev.Name)
				}
				if m.Delete != nil {
					m.Delete(ev.Name)
				}
			case ev.Op&fsnotify.Write == fsnotify.Write:
				if m.Modify != nil {
					m.Modify(ev.Name)
				}
			case ev.Op&fsnotify.Rename == fsnotify.Rename:
				watcher.Remove(ev.Name)
				if m.Rename != nil {
					m.Rename(ev.Name)
				}
			case ev.Op&fsnotify.Chmod == fsnotify.Chmod:
				if m.Chmod != nil {
					m.Chmod(ev.Name)
				}
			}
		case err, ok := <-watcher.Errors:
			if !ok {
				return
			}
			if err != nil {
				log.Println("Watcher error:", err)
			}
		}
	}
}

func (m *MonitorEvent) IsDir(path string) bool {
	d, e := os.Stat(path)
	if e != nil {
		return false
	}
	return d.IsDir()
}

//Monitor 文件监测
func Monitor(rootDir string, callback *MonitorEvent, args ...func(string) bool) error {
	watcher := callback.Watcher()
	defer watcher.Close()
	callback.Watch(args...)
	err := callback.AddDir(rootDir)
	if err != nil {
		callback.Close()
		return err
	}
	return nil
}
<|endoftext|>