<commit_before>package lxdatabase_test

import (
	"github.com/layer-x/layerx-commons/lxdatabase"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"os"
	"os/exec"
	"runtime"
	"time"
)

var binaryUrl string
var fileName string
var extract *exec.Cmd
var run *exec.Cmd

//not working, must run etcd manually first
func startETCD() {
	//start a test etcd server (will not work if etcd already running on host
	if runtime.GOOS == "darwin" {
		binaryUrl = "https://github.com/coreos/etcd/releases/download/v2.2.2/etcd-v2.2.2-darwin-amd64.zip"
		fileName = "etcd-v2.2.2-darwin-amd64.zip"
		extract = exec.Command("unzip", fileName, "-d", "etcd")
		run = exec.Command("etcd/etcd-v2.2.2-darwin-amd64/etcd")
	}
	if runtime.GOOS == "linux" {
		binaryUrl = "https://github.com/coreos/etcd/releases/download/v2.2.2/etcd-v2.2.2-linux-amd64.tar.gz"
		fileName = "etcd-v2.2.2-linux-amd64.tar.gz"
		exec.Command("mkdir", "etcd").Run()
		extract = exec.Command("tar", "xzvf", fileName, "-C", "etcd")
		run = exec.Command("etcd/etcd-v2.2.2-linux-amd64/etcd")
	}
	exec.Command("curl", "-L", binaryUrl, "-o", fileName).Run()
	extract.Run()
	go func() {
		run.Run()
	}()
	//5 seconds to initialize etcd
	time.Sleep(5 * time.Second)
}

var _ = Describe("Lxdb", func() {
	startETCD()

	BeforeEach(func() {
		Describe("initializes without error", func() {
			err := lxdatabase.Init([]string{"http://127.0.0.1:2379"})
			Expect(err).To(BeNil())
		})
	})
	Describe("lxdatabase.Set(key, val)", func() {
		It("sets keys", func() {
			err := lxdatabase.Set("foo", "bar")
			Expect(err).To(BeNil())
		})
	})
	Describe("lxdatabase.Get(key)", func() {
		It("gets values", func() {
			val, err := lxdatabase.Get("foo")
			Expect(err).To(BeNil())
			Expect(val).To(Equal("bar"))
		})
	})
	Describe("lxdatabase.Rm(key)", func() {
		It("deletes keys", func() {
			err := lxdatabase.Rm("foo")
			Expect(err).To(BeNil())
			val, err := lxdatabase.Get("foo")
			Expect(val).To(Equal(""))
			Expect(err).To(Not(BeNil()))
			Expect(err.Error()).To(ContainSubstring("Key not found"))
		})
	})
	Describe("lxdatabase.Mkdir(key)", func() {
		It("makes directories", func() {
			err := lxdatabase.Mkdir("foo_dir")
			Expect(err).To(BeNil())
		})
	})
	Describe("lxdatabase.Ls(key)", func() {
		It("lists directories", func() {
			keys, err := lxdatabase.Ls("foo_dir")
			Expect(err).To(BeNil())
			Expect(keys).To(BeEmpty())
			err = lxdatabase.Set("foo_dir/foo", "bar")
			Expect(err).To(BeNil())
			keys, err = lxdatabase.Ls("foo_dir")
			Expect(err).To(BeNil())
			Expect(keys).To(ContainElement("bar"))
		})
	})
	Describe("cleanup", func() {
		It("cleans up etcd", func() {
			os.RemoveAll(fileName)
			os.RemoveAll("etcd")
			os.RemoveAll("default.etcd")
			exec.Command("pkill", "etcd")
		})
	})
})
<commit_msg>remove comment<commit_after>package lxdatabase_test

import (
	"github.com/layer-x/layerx-commons/lxdatabase"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"os"
	"os/exec"
	"runtime"
	"time"
)

var binaryUrl string
var fileName string
var extract *exec.Cmd
var run *exec.Cmd

func startETCD() {
	//start a test etcd server (will not work if etcd already running on host
	if runtime.GOOS == "darwin" {
		binaryUrl = "https://github.com/coreos/etcd/releases/download/v2.2.2/etcd-v2.2.2-darwin-amd64.zip"
		fileName = "etcd-v2.2.2-darwin-amd64.zip"
		extract = exec.Command("unzip", fileName, "-d", "etcd")
		run = exec.Command("etcd/etcd-v2.2.2-darwin-amd64/etcd")
	}
	if runtime.GOOS == "linux" {
		binaryUrl = "https://github.com/coreos/etcd/releases/download/v2.2.2/etcd-v2.2.2-linux-amd64.tar.gz"
		fileName = "etcd-v2.2.2-linux-amd64.tar.gz"
		exec.Command("mkdir", "etcd").Run()
		extract = exec.Command("tar", "xzvf", fileName, "-C", "etcd")
		run = exec.Command("etcd/etcd-v2.2.2-linux-amd64/etcd")
	}
	exec.Command("curl", "-L", binaryUrl, "-o", fileName).Run()
	extract.Run()
	go func() {
		run.Run()
	}()
	//5 seconds to initialize etcd
	time.Sleep(5 * time.Second)
}

var _ = Describe("Lxdb", func() {
	startETCD()

	BeforeEach(func() {
		Describe("initializes without error", func() {
			err := lxdatabase.Init([]string{"http://127.0.0.1:2379"})
			Expect(err).To(BeNil())
		})
	})
	Describe("lxdatabase.Set(key, val)", func() {
		It("sets keys", func() {
			err := lxdatabase.Set("foo", "bar")
			Expect(err).To(BeNil())
		})
	})
	Describe("lxdatabase.Get(key)", func() {
		It("gets values", func() {
			val, err := lxdatabase.Get("foo")
			Expect(err).To(BeNil())
			Expect(val).To(Equal("bar"))
		})
	})
	Describe("lxdatabase.Rm(key)", func() {
		It("deletes keys", func() {
			err := lxdatabase.Rm("foo")
			Expect(err).To(BeNil())
			val, err := lxdatabase.Get("foo")
			Expect(val).To(Equal(""))
			Expect(err).To(Not(BeNil()))
			Expect(err.Error()).To(ContainSubstring("Key not found"))
		})
	})
	Describe("lxdatabase.Mkdir(key)", func() {
		It("makes directories", func() {
			err := lxdatabase.Mkdir("foo_dir")
			Expect(err).To(BeNil())
		})
	})
	Describe("lxdatabase.Ls(key)", func() {
		It("lists directories", func() {
			keys, err := lxdatabase.Ls("foo_dir")
			Expect(err).To(BeNil())
			Expect(keys).To(BeEmpty())
			err = lxdatabase.Set("foo_dir/foo", "bar")
			Expect(err).To(BeNil())
			keys, err = lxdatabase.Ls("foo_dir")
			Expect(err).To(BeNil())
			Expect(keys).To(ContainElement("bar"))
		})
	})
	Describe("cleanup", func() {
		It("cleans up etcd", func() {
			os.RemoveAll(fileName)
			os.RemoveAll("etcd")
			os.RemoveAll("default.etcd")
			exec.Command("pkill", "etcd")
		})
	})
})
<|endoftext|>


<commit_before>package upload

import (
	"crypto/md5"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// Error Incomplete returned by uploader when loaded non-last chunk.
var Incomplete = errors.New("Incomplete")

// Structure describes the state of the original file.
type OriginalFile struct {
	BaseMime string
	Filepath string
	Filename string
	Size     int64
}

func (ofile *OriginalFile) Ext() string {
	return strings.ToLower(filepath.Ext(ofile.Filename))
}

// Downloading files from the received request.
// The root directory of storage, storage,  used to temporarily store chunks.
// Returns an array of the original files and error.
// If you load a portion of the file, chunk, it will be stored in err error Incomplete,
// and in an array of a single file. File size will fit the current size.
func Process(cookie_name string, req *http.Request, storage string) ([]*OriginalFile, error) {
	meta, err := ParseMeta(cookie_name, req)
	if err != nil {
		return nil, err
	}

	body, err := NewBody(req.Header.Get("X-File"), req.Body)
	if err != nil {
		return nil, err
	}
	up := &Uploader{Root: storage, Meta: meta, Body: body}
	files, err := up.SaveFiles()
	if err == Incomplete {
		return files, err
	}
	if err != nil {
		return nil, err
	}

	return files, nil
}

// Upload manager.
type Uploader struct {
	Root string
	Meta *Meta
	Body *Body
}

// Function SaveFiles sequentially loads the original files or chunk's.
func (up *Uploader) SaveFiles() ([]*OriginalFile, error) {
	files := make([]*OriginalFile, 0)
	for {
		ofile, err := up.SaveFile()
		if err == io.EOF {
			break
		}

		if err == Incomplete {
			files = append(files, ofile)
			return files, err
		}

		if err != nil {
			return nil, err
		}

		files = append(files, ofile)
	}

	return files, nil
}

// Function loads one or download the original file chunk.
// Asks for the starting position in the body of the request to read the next file.
// Asks for a temporary file.
// Writes data from the request body into a temporary file.
// Specifies the size of the resulting temporary file.
// If the query specified header Content-Range,
// and the size of the resulting file does not match, it returns an error Incomplete.
// Otherwise, defines the basic mime type, and returns the original file.
func (up *Uploader) SaveFile() (*OriginalFile, error) {
	body, filename, err := up.Reader()
	if err != nil {
		return nil, err
	}

	temp_file, err := up.TempFile()
	if err != nil {
		return nil, err
	}
	defer temp_file.Close()

	if err = up.Write(temp_file, body); err != nil {
		return nil, err
	}

	fi, err := temp_file.Stat()
	if err != nil {
		return nil, err
	}
	if fi.Size() == 0 {
		return nil, errors.New("Empty file")
	}
	ofile := &OriginalFile{Filename: filename, Filepath: temp_file.Name(), Size: fi.Size()}

	if up.Meta.Range != nil && ofile.Size != up.Meta.Range.Size {
		return ofile, Incomplete
	}

	ofile.BaseMime, err = IdentifyMime(ofile.Filepath)
	if err != nil {
		return nil, err
	}

	return ofile, nil
}

// Returns the reader to read the file or chunk of request body and the original file name.
// If the request header Conent-Type is multipart/form-data, returns the next copy part.
// If all of part read the case of binary loading read the request body, an error is returned io.EOF.
func (up *Uploader) Reader() (io.Reader, string, error) {
	if up.Meta.MediaType == "multipart/form-data" {
		if up.Body.MR == nil {
			up.Body.MR = multipart.NewReader(up.Body.Body, up.Meta.Boundary)
		}
		for {
			part, err := up.Body.MR.NextPart()
			if err != nil {
				return nil, "", err
			}
			if part.FormName() == "file" {
				return part, part.FileName(), nil
			}
			if part.FormName() == "files[]" {
				return part, part.FileName(), nil
			}
		}
	}

	if up.Body.Available == false {
		return nil, "", io.EOF
	}

	up.Body.Available = false

	return up.Body.Body, up.Meta.Filename, nil
}

// Returns a temporary file to download the file or resume chunk.
func (up *Uploader) TempFile() (*os.File, error) {
	if up.Meta.Range == nil {
		return TempFile()
	}
	return TempFileChunks(up.Meta.Range.Start, up.Root, up.Meta.UploadSid, up.Meta.Filename)
}

// Returns the newly created temporary file.
func TempFile() (*os.File, error) {
	return ioutil.TempFile(os.TempDir(), "pavo")
}

// Returns a temporary file to download chunk.
// To calculate a unique file name used cookie named pavo and the original file name.
// File located in the directory chunks storage root directory.
// Before returning the file pointer is shifted by the value of offset,
// in a situation where the pieces are loaded from the second to the last.
func TempFileChunks(offset int64, storage, upload_sid, user_filename string) (*os.File, error) {
	hasher := md5.New()
	hasher.Write([]byte(upload_sid + user_filename))
	filename := hex.EncodeToString(hasher.Sum(nil))

	path := filepath.Join(storage, "chunks")

	err := os.MkdirAll(path, 0755)
	if err != nil {
		return nil, err
	}

	file, err := os.OpenFile(filepath.Join(path, filename), os.O_CREATE|os.O_WRONLY, 0664)
	if err != nil {
		return nil, err
	}

	if _, err = file.Seek(offset, 0); err != nil {
		return nil, err
	}

	return file, nil
}

// The function writes a temporary file value from reader.
func (up *Uploader) Write(temp_file *os.File, body io.Reader) error {
	var err error
	if up.Meta.Range == nil {
		_, err = io.Copy(temp_file, body)
	} else {
		chunk_size := up.Meta.Range.End - up.Meta.Range.Start + 1
		_, err = io.CopyN(temp_file, body, chunk_size)
	}
	return err
}

// Get base mime type
//		$ file --mime-type pic.jpg
//		pic.jpg: image/jpeg
func IdentifyMime(file string) (string, error) {
	out, err := exec.Command("file", "--mime-type", file).CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("identify: err: %s; detail: %s", err, string(out))
	}

	mime := strings.Split(strings.Split(string(out), ": ")[1], "/")[0]

	return mime, nil
}
<commit_msg>debug symbol<commit_after>package upload

import (
	"crypto/md5"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// Error Incomplete returned by uploader when loaded non-last chunk.
var Incomplete = errors.New("Incomplete")

// Structure describes the state of the original file.
type OriginalFile struct {
	BaseMime string
	Filepath string
	Filename string
	Size     int64
}

func (ofile *OriginalFile) Ext() string {
	return strings.ToLower(filepath.Ext(ofile.Filename))
}

// Downloading files from the received request.
// The root directory of storage, storage,  used to temporarily store chunks.
// Returns an array of the original files and error.
// If you load a portion of the file, chunk, it will be stored in err error Incomplete,
// and in an array of a single file. File size will fit the current size.
func Process(cookie_name string, req *http.Request, storage string) ([]*OriginalFile, error) {
	meta, err := ParseMeta(cookie_name, req)
	if err != nil {
		return nil, err
	}

	body, err := NewBody(req.Header.Get("X-File"), req.Body)
	if err != nil {
		println(err.Error())
		return nil, err
	}
	up := &Uploader{Root: storage, Meta: meta, Body: body}
	files, err := up.SaveFiles()
	if err == Incomplete {
		return files, err
	}
	if err != nil {
		return nil, err
	}

	return files, nil
}

// Upload manager.
type Uploader struct {
	Root string
	Meta *Meta
	Body *Body
}

// Function SaveFiles sequentially loads the original files or chunk's.
func (up *Uploader) SaveFiles() ([]*OriginalFile, error) {
	files := make([]*OriginalFile, 0)
	for {
		ofile, err := up.SaveFile()
		if err == io.EOF {
			break
		}

		if err == Incomplete {
			files = append(files, ofile)
			return files, err
		}

		if err != nil {
			return nil, err
		}

		files = append(files, ofile)
	}

	return files, nil
}

// Function loads one or download the original file chunk.
// Asks for the starting position in the body of the request to read the next file.
// Asks for a temporary file.
// Writes data from the request body into a temporary file.
// Specifies the size of the resulting temporary file.
// If the query specified header Content-Range,
// and the size of the resulting file does not match, it returns an error Incomplete.
// Otherwise, defines the basic mime type, and returns the original file.
func (up *Uploader) SaveFile() (*OriginalFile, error) {
	body, filename, err := up.Reader()
	if err != nil {
		return nil, err
	}

	temp_file, err := up.TempFile()
	if err != nil {
		return nil, err
	}
	defer temp_file.Close()

	if err = up.Write(temp_file, body); err != nil {
		return nil, err
	}

	fi, err := temp_file.Stat()
	if err != nil {
		return nil, err
	}
	if fi.Size() == 0 {
		return nil, errors.New("Empty file")
	}
	ofile := &OriginalFile{Filename: filename, Filepath: temp_file.Name(), Size: fi.Size()}

	if up.Meta.Range != nil && ofile.Size != up.Meta.Range.Size {
		return ofile, Incomplete
	}

	ofile.BaseMime, err = IdentifyMime(ofile.Filepath)
	if err != nil {
		return nil, err
	}

	return ofile, nil
}

// Returns the reader to read the file or chunk of request body and the original file name.
// If the request header Conent-Type is multipart/form-data, returns the next copy part.
// If all of part read the case of binary loading read the request body, an error is returned io.EOF.
func (up *Uploader) Reader() (io.Reader, string, error) {
	if up.Meta.MediaType == "multipart/form-data" {
		if up.Body.MR == nil {
			up.Body.MR = multipart.NewReader(up.Body.Body, up.Meta.Boundary)
		}
		for {
			part, err := up.Body.MR.NextPart()
			if err != nil {
				return nil, "", err
			}
			if part.FormName() == "file" {
				return part, part.FileName(), nil
			}
			if part.FormName() == "files[]" {
				return part, part.FileName(), nil
			}
		}
	}

	if up.Body.Available == false {
		return nil, "", io.EOF
	}

	up.Body.Available = false

	return up.Body.Body, up.Meta.Filename, nil
}

// Returns a temporary file to download the file or resume chunk.
func (up *Uploader) TempFile() (*os.File, error) {
	if up.Meta.Range == nil {
		return TempFile()
	}
	return TempFileChunks(up.Meta.Range.Start, up.Root, up.Meta.UploadSid, up.Meta.Filename)
}

// Returns the newly created temporary file.
func TempFile() (*os.File, error) {
	return ioutil.TempFile(os.TempDir(), "pavo")
}

// Returns a temporary file to download chunk.
// To calculate a unique file name used cookie named pavo and the original file name.
// File located in the directory chunks storage root directory.
// Before returning the file pointer is shifted by the value of offset,
// in a situation where the pieces are loaded from the second to the last.
func TempFileChunks(offset int64, storage, upload_sid, user_filename string) (*os.File, error) {
	hasher := md5.New()
	hasher.Write([]byte(upload_sid + user_filename))
	filename := hex.EncodeToString(hasher.Sum(nil))

	path := filepath.Join(storage, "chunks")

	err := os.MkdirAll(path, 0755)
	if err != nil {
		return nil, err
	}

	file, err := os.OpenFile(filepath.Join(path, filename), os.O_CREATE|os.O_WRONLY, 0664)
	if err != nil {
		return nil, err
	}

	if _, err = file.Seek(offset, 0); err != nil {
		return nil, err
	}

	return file, nil
}

// The function writes a temporary file value from reader.
func (up *Uploader) Write(temp_file *os.File, body io.Reader) error {
	var err error
	if up.Meta.Range == nil {
		_, err = io.Copy(temp_file, body)
	} else {
		chunk_size := up.Meta.Range.End - up.Meta.Range.Start + 1
		_, err = io.CopyN(temp_file, body, chunk_size)
	}
	return err
}

// Get base mime type
//		$ file --mime-type pic.jpg
//		pic.jpg: image/jpeg
func IdentifyMime(file string) (string, error) {
	out, err := exec.Command("file", "--mime-type", file).CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("identify: err: %s; detail: %s", err, string(out))
	}

	mime := strings.Split(strings.Split(string(out), ": ")[1], "/")[0]

	return mime, nil
}
<|endoftext|>


<commit_before>package glock

import (
	"fmt"
	"testing"
	"time"
)

func TestLockUnlock(t *testing.T) {
	client := NewClient("localhost:45625")
	fmt.Println("1 getting lock")
	id1 := client.LockEx("x", 10000)
	fmt.Println("1 got lock")
	go func() {
		fmt.Println("2 getting lock")
		id2 := client.LockEx("x", 10000)
		fmt.Println("2 got lock")
		fmt.Println("2 releasing lock")
		client.UnLock("x", id2)
		fmt.Println("2 released lock")
	}()
	fmt.Println("sleeping")
	time.Sleep(5 * time.Second)
	fmt.Println("1 releasing lock")
	client.UnLock("x", id1)
	fmt.Println("1 released lock")
	time.Sleep(5 * time.Second)
}
<commit_msg>Some formatting<commit_after>package glock

import (
	"fmt"
	"testing"
	"time"
)

func TestLockUnlock(t *testing.T) {
	client := NewClient("localhost:45625")

	fmt.Println("1 getting lock")
	id1 := client.LockEx("x", 10000)
	fmt.Println("1 got lock")

	go func() {

		fmt.Println("2 getting lock")
		id2 := client.LockEx("x", 10000)
		fmt.Println("2 got lock")

		fmt.Println("2 releasing lock")
		client.UnLock("x", id2)
		fmt.Println("2 released lock")

	}()

	fmt.Println("sleeping")
	time.Sleep(5 * time.Second)
	fmt.Println("finished sleeping")

	fmt.Println("1 releasing lock")
	client.UnLock("x", id1)
	fmt.Println("1 released lock")

	time.Sleep(5 * time.Second)
}
<|endoftext|>


<commit_before>package config

import (
	"fmt"
	"github.com/mailgun/go-statsd-client/statsd"
	cfg "github.com/mailgun/gotools-config"
	log "github.com/mailgun/gotools-log"
	"github.com/mailgun/pong/model"
	"math/rand"
	"net/http"
	"strconv"
	"strings"
	"time"
)

type Response struct {
	Rate        string
	Code        int
	Body        string
	ContentType string
	Delay       string
	Drop        bool
}

type Handler struct {
	Id        string
	Path      string
	Responses []*Response
}

type Server struct {
	Addr         string
	Path         string
	Handlers     []*Handler
	ReadTimeout  string
	WriteTimeout string
}

type Config struct {
	Statsd struct {
		Url    string
		Prefix string
	}
	Servers []*Server
	Logging []*log.LogConfig
}

type HandlerFn func(http.ResponseWriter, *http.Request)

func ParseConfig(path string) ([]*model.Server, []*log.LogConfig, error) {
	config := Config{}
	if err := cfg.LoadConfig(path, &config); err != nil {
		return nil, nil, fmt.Errorf("Failed to load config file '%s' err:", path, err)
	}

	client, err := statsd.New(config.Statsd.Url, config.Statsd.Prefix)
	if err != nil {
		return nil, nil, err
	}
	builder := &Builder{
		client: client,
	}
	servers, err := builder.parseServers(config.Servers)
	if err != nil {
		return nil, nil, err
	}
	return servers, config.Logging, nil
}

type Builder struct {
	client *statsd.Client
}

func (b *Builder) parseServers(in []*Server) ([]*model.Server, error) {
	out := make([]*model.Server, len(in))
	for i, cserver := range in {
		serv, err := b.parseServer(cserver)
		if err != nil {
			return nil, err
		}
		out[i] = serv
	}
	return out, nil
}

func (b *Builder) parseServer(in *Server) (*model.Server, error) {
	readTimeout, err := time.ParseDuration(in.ReadTimeout)
	if err != nil {
		return nil, err
	}
	writeTimeout, err := time.ParseDuration(in.WriteTimeout)
	if err != nil {
		return nil, err
	}
	mux := http.NewServeMux()
	for _, h := range in.Handlers {
		handler, err := b.buildHandler(h.Id, h.Responses)
		if err != nil {
			return nil, err
		}
		mux.Handle(h.Path, handler)
	}
	return &model.Server{
		Addr:         in.Addr,
		ReadTimeout:  readTimeout,
		WriteTimeout: writeTimeout,
		Handler:      mux,
	}, nil
}

type Responder struct {
	id        string
	responses []*model.Response
	index     int
	client    *statsd.Client
}

func (d *Responder) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	r := d.responses[d.index]
	d.index = (d.index + 1) % len(d.responses)

	d.client.Inc(metric("requests"), 1, 1)
	d.client.Inc(metric(d.id, "requests"), 1, 1)

	if r.Delay > 0 {
		time.Sleep(r.Delay)
	}
	if r.Drop {
		h := w.(http.Hijacker)
		conn, _, _ := h.Hijack()
		conn.Close()
		return
	}
	w.WriteHeader(r.Code)
	w.Header().Set("Content-Type", r.ContentType)
	w.Write([]byte(r.Body))
}

func (b *Builder) buildHandler(id string, distribution []*Response) (http.Handler, error) {
	responses, err := b.buildResponses(distribution)
	if err != nil {
		return nil, err
	}
	return &Responder{
		id:        id,
		responses: responses,
		client:    b.client,
	}, nil
}

func (b *Builder) buildResponses(responses []*Response) ([]*model.Response, error) {
	out := []*model.Response{}
	total := 0
	for _, re := range responses {
		duration, err := time.ParseDuration(re.Delay)
		if err != nil {
			return nil, fmt.Errorf("Bad delday '%s', should me in form '1s'", re.Delay)
		}
		count, err := parsePercent(re.Rate)
		if err != nil {
			return nil, err
		}
		total += count
		for i := 0; i < count; i += 1 {
			log.Infof("Appending %s", re.Body)
			out = append(out, &model.Response{
				Drop:        re.Drop,
				Delay:       duration,
				Code:        re.Code,
				Body:        []byte(re.Body),
				ContentType: re.ContentType,
			})
		}
	}
	if total != 100 {
		return nil, fmt.Errorf("Percentages should form 100 in sum, got %d", total)
	}
	return shuffle(out), nil
}

func shuffle(in []*model.Response) []*model.Response {
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	indexes := r.Perm(len(in))
	out := make([]*model.Response, len(in))
	for i, v := range indexes {
		out[i] = in[v]
	}
	return out
}

func parsePercent(in string) (int, error) {
	num, err := strconv.Atoi(strings.TrimSuffix(in, "%"))
	if err != nil {
		return -1, fmt.Errorf("Use percentages, e.g 10%")
	}
	if num > 100 || num < 0 {
		return -1, fmt.Errorf("Percentage value should be withing 0% to 100%")
	}
	return num, nil
}

func metric(vals ...string) string {
	return strings.Join(vals, ".")
}
<commit_msg>Fix statsd change<commit_after>package config

import (
	"fmt"
	"github.com/mailgun/go-statsd-client/statsd"
	cfg "github.com/mailgun/gotools-config"
	log "github.com/mailgun/gotools-log"
	"github.com/mailgun/pong/model"
	"math/rand"
	"net/http"
	"strconv"
	"strings"
	"time"
)

type Response struct {
	Rate        string
	Code        int
	Body        string
	ContentType string
	Delay       string
	Drop        bool
}

type Handler struct {
	Id        string
	Path      string
	Responses []*Response
}

type Server struct {
	Addr         string
	Path         string
	Handlers     []*Handler
	ReadTimeout  string
	WriteTimeout string
}

type Config struct {
	Statsd struct {
		Url    string
		Prefix string
	}
	Servers []*Server
	Logging []*log.LogConfig
}

type HandlerFn func(http.ResponseWriter, *http.Request)

func ParseConfig(path string) ([]*model.Server, []*log.LogConfig, error) {
	config := Config{}
	if err := cfg.LoadConfig(path, &config); err != nil {
		return nil, nil, fmt.Errorf("Failed to load config file '%s' err:", path, err)
	}

	client, err := statsd.New(config.Statsd.Url, config.Statsd.Prefix)
	if err != nil {
		return nil, nil, err
	}
	builder := &Builder{
		client: client,
	}
	servers, err := builder.parseServers(config.Servers)
	if err != nil {
		return nil, nil, err
	}
	return servers, config.Logging, nil
}

type Builder struct {
	client statsd.Client
}

func (b *Builder) parseServers(in []*Server) ([]*model.Server, error) {
	out := make([]*model.Server, len(in))
	for i, cserver := range in {
		serv, err := b.parseServer(cserver)
		if err != nil {
			return nil, err
		}
		out[i] = serv
	}
	return out, nil
}

func (b *Builder) parseServer(in *Server) (*model.Server, error) {
	readTimeout, err := time.ParseDuration(in.ReadTimeout)
	if err != nil {
		return nil, err
	}
	writeTimeout, err := time.ParseDuration(in.WriteTimeout)
	if err != nil {
		return nil, err
	}
	mux := http.NewServeMux()
	for _, h := range in.Handlers {
		handler, err := b.buildHandler(h.Id, h.Responses)
		if err != nil {
			return nil, err
		}
		mux.Handle(h.Path, handler)
	}
	return &model.Server{
		Addr:         in.Addr,
		ReadTimeout:  readTimeout,
		WriteTimeout: writeTimeout,
		Handler:      mux,
	}, nil
}

type Responder struct {
	id        string
	responses []*model.Response
	index     int
	client    statsd.Client
}

func (d *Responder) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	r := d.responses[d.index]
	d.index = (d.index + 1) % len(d.responses)

	d.client.Inc(metric("requests"), 1, 1)
	d.client.Inc(metric(d.id, "requests"), 1, 1)

	if r.Delay > 0 {
		time.Sleep(r.Delay)
	}
	if r.Drop {
		h := w.(http.Hijacker)
		conn, _, _ := h.Hijack()
		conn.Close()
		return
	}
	w.WriteHeader(r.Code)
	w.Header().Set("Content-Type", r.ContentType)
	w.Write([]byte(r.Body))
}

func (b *Builder) buildHandler(id string, distribution []*Response) (http.Handler, error) {
	responses, err := b.buildResponses(distribution)
	if err != nil {
		return nil, err
	}
	return &Responder{
		id:        id,
		responses: responses,
		client:    b.client,
	}, nil
}

func (b *Builder) buildResponses(responses []*Response) ([]*model.Response, error) {
	out := []*model.Response{}
	total := 0
	for _, re := range responses {
		duration, err := time.ParseDuration(re.Delay)
		if err != nil {
			return nil, fmt.Errorf("Bad delday '%s', should me in form '1s'", re.Delay)
		}
		count, err := parsePercent(re.Rate)
		if err != nil {
			return nil, err
		}
		total += count
		for i := 0; i < count; i += 1 {
			log.Infof("Appending %s", re.Body)
			out = append(out, &model.Response{
				Drop:        re.Drop,
				Delay:       duration,
				Code:        re.Code,
				Body:        []byte(re.Body),
				ContentType: re.ContentType,
			})
		}
	}
	if total != 100 {
		return nil, fmt.Errorf("Percentages should form 100 in sum, got %d", total)
	}
	return shuffle(out), nil
}

func shuffle(in []*model.Response) []*model.Response {
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	indexes := r.Perm(len(in))
	out := make([]*model.Response, len(in))
	for i, v := range indexes {
		out[i] = in[v]
	}
	return out
}

func parsePercent(in string) (int, error) {
	num, err := strconv.Atoi(strings.TrimSuffix(in, "%"))
	if err != nil {
		return -1, fmt.Errorf("Use percentages, e.g 10%")
	}
	if num > 100 || num < 0 {
		return -1, fmt.Errorf("Percentage value should be withing 0% to 100%")
	}
	return num, nil
}

func metric(vals ...string) string {
	return strings.Join(vals, ".")
}
<|endoftext|>


<commit_before>package upload

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"

	"github.com/subutai-io/base/agent/log"
	"github.com/subutai-io/gorjun/config"
	"github.com/subutai-io/gorjun/db"
)

type share struct {
	Token  string   `json:"token"`
	Id     string   `json:"id"`
	Add    []string `json:"add"`
	Remove []string `json:"remove"`
}

//Handler function works with income upload requests, makes sanity checks, etc
func Handler(w http.ResponseWriter, r *http.Request) (hash, owner string) {
	r.ParseMultipartForm(32 << 20)
	if len(r.MultipartForm.Value["token"]) == 0 || len(db.CheckToken(r.MultipartForm.Value["token"][0])) == 0 {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("Not authorized"))
		log.Warn(r.RemoteAddr + " - rejecting unauthorized upload request")
		return
	}

	owner = db.CheckToken(r.MultipartForm.Value["token"][0])

	file, header, err := r.FormFile("file")
	defer file.Close()
	if log.Check(log.WarnLevel, "Failed to parse POST form", err) {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Cannot get file from request"))
		return
	}

	l := сheckLength(owner, r.Header.Get("Content-Length"))
	if l == 1 {
		w.WriteHeader(http.StatusNotAcceptable)
		w.Write([]byte("Storage quota exceeded"))
		log.Warn("User " + owner + " exceeded storage quota, rejecting upload")
		return
	}

	out, err := os.Create(config.Storage.Path + header.Filename)
	if log.Check(log.WarnLevel, "Unable to create the file for writing", err) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("Cannot create file"))
		return
	}
	defer out.Close()

	// write the content from POST to the file
	_, err = io.Copy(out, file)
	if log.Check(log.WarnLevel, "Writing file", err) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("Failed to write file"))
		return
	}

	if l == 2 {
		info, _ := out.Stat()
		log.Warn("Counting quota after upload!")
		log.Warn("Uploaded: " + strconv.Itoa(int(info.Size())) + ", quota left: " + strconv.Itoa(db.QuotaLeft(owner)))
		if int(info.Size()) > db.QuotaLeft(owner) {
			w.WriteHeader(http.StatusNotAcceptable)
			w.Write([]byte("Storage quota exceeded after upload"))
			os.Remove(config.Storage.Path + header.Filename)
			return
		}
	}

	hash = genHash(config.Storage.Path + header.Filename)
	if len(hash) == 0 {
		log.Warn("Failed to calculate hash for " + header.Filename)
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("Failed to calculate hash"))
		return
	}

	os.Rename(config.Storage.Path+header.Filename, config.Storage.Path+hash)
	log.Info("File uploaded successfully: " + header.Filename + "(" + hash + ")")

	return hash, owner
}

func genHash(file string) string {
	f, err := os.Open(file)
	log.Check(log.WarnLevel, "Opening file"+file, err)
	defer f.Close()

	hash := md5.New()
	if _, err := io.Copy(hash, f); err != nil {
		return ""
	}
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func Delete(w http.ResponseWriter, r *http.Request) string {
	hash := r.URL.Query().Get("id")
	token := r.URL.Query().Get("token")
	if len(hash) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Empty file id"))
		log.Warn(r.RemoteAddr + " - empty file id")
		return ""
	}
	if len(token) == 0 {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("Empty token"))
		log.Warn(r.RemoteAddr + " - empty token")
		return ""
	}
	user := db.CheckToken(token)
	info := db.Info(hash)
	if len(info) == 0 {
		log.Warn("File not found by hash")
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("File not found"))
		return ""
	}

	if !db.CheckOwner(user, hash) {
		log.Warn("File " + info["name"] + "(" + hash + ") is not owned by " + user + ", rejecting deletion request")
		w.WriteHeader(http.StatusForbidden)
		w.Write([]byte("File " + info["name"] + " is not owned by " + user))
		return ""
	}

	f, _ := os.Stat(config.Storage.Path + hash)
	db.QuotaUsageSet(user, -int(f.Size()))

	if db.Delete(user, hash) <= 0 {
		if log.Check(log.WarnLevel, "Removing "+info["name"]+"from disk", os.Remove(config.Storage.Path+hash)) {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("Failed to remove file"))
			return ""
		}
	}

	log.Info("Removing " + info["name"])
	return hash
}

func Share(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		if len(r.FormValue("json")) == 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Empty json"))
			log.Warn("Share request: empty json, nothing to do")
			return
		}
		var data share
		if log.Check(log.WarnLevel, "Parsing share request json", json.Unmarshal([]byte(r.FormValue("json")), &data)) {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Failed to parse json body"))
			return
		}
		if len(data.Token) == 0 || len(db.CheckToken(data.Token)) == 0 {
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte("Not authorized"))
			log.Warn("Empty or invalid token, rejecting share request")
			return
		}
		if len(data.Id) == 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Empty file id"))
			log.Warn("Empty file id, rejecting share request")
			return
		}
		owner := db.CheckToken(data.Token)
		if !db.CheckOwner(owner, data.Id) {
			w.WriteHeader(http.StatusForbidden)
			w.Write([]byte("File is not owned by authorized user"))
			log.Warn("User tried to share another's file, rejecting")
			return
		}
		for _, v := range data.Add {
			db.ShareWith(data.Id, owner, v)
		}
		for _, v := range data.Remove {
			db.UnshareWith(data.Id, owner, v)
		}
	} else if r.Method == "GET" {
		id := r.URL.Query().Get("id")
		if len(id) == 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Empty file id"))
			return
		}
		token := r.URL.Query().Get("token")
		if len(token) == 0 || len(db.CheckToken(token)) == 0 {
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte("Not authorized"))
			return
		}
		owner := db.CheckToken(token)
		if !db.CheckOwner(owner, id) {
			w.WriteHeader(http.StatusForbidden)
			w.Write([]byte("File is not owned by authorized user"))
			log.Warn("User tried to request scope of another's file, rejecting")
			return
		}
		js, _ := json.Marshal(db.GetScope(id, owner))
		w.Write(js)
	}
}

func сheckLength(user, length string) int {
	l, err := strconv.Atoi(length)
	if log.Check(log.WarnLevel, "Converting content length to int", err) || len(length) == 0 {
		return 2
	}

	if l > db.QuotaLeft(user) {
		return 1
	}
	db.QuotaUsageSet(user, l)
	return 0
}

func QuotaTest(w http.ResponseWriter, r *http.Request) {
	user := r.URL.Query().Get("user")
	if len(user) != 0 {
		w.Write([]byte("Quota: " + strconv.Itoa(db.QuotaGet(user)) + "b\n"))
		w.Write([]byte("Quota used: " + strconv.Itoa(db.QuotaUsageGet(user)) + "b\n"))
		w.Write([]byte("Quota left: " + strconv.Itoa(db.QuotaLeft(user)) + "b\n"))
	}

	quota := r.URL.Query().Get("quota")
	if len(quota) != 0 {
		db.QuotaSet(user, quota)
		w.Write([]byte("user: " + user + ", quota: " + quota + "b\n"))
	}
}
<commit_msg>Quota limitation on the go<commit_after>package upload

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"

	"github.com/subutai-io/base/agent/log"
	"github.com/subutai-io/gorjun/config"
	"github.com/subutai-io/gorjun/db"
)

type share struct {
	Token  string   `json:"token"`
	Id     string   `json:"id"`
	Add    []string `json:"add"`
	Remove []string `json:"remove"`
}

//Handler function works with income upload requests, makes sanity checks, etc
func Handler(w http.ResponseWriter, r *http.Request) (hash, owner string) {
	r.ParseMultipartForm(32 << 20)
	if len(r.MultipartForm.Value["token"]) == 0 || len(db.CheckToken(r.MultipartForm.Value["token"][0])) == 0 {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("Not authorized"))
		log.Warn(r.RemoteAddr + " - rejecting unauthorized upload request")
		return
	}

	owner = db.CheckToken(r.MultipartForm.Value["token"][0])

	file, header, err := r.FormFile("file")
	defer file.Close()
	if log.Check(log.WarnLevel, "Failed to parse POST form", err) {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Cannot get file from request"))
		return
	}

	l := сheckLength(owner, r.Header.Get("Content-Length"))
	if !l {
		w.WriteHeader(http.StatusNotAcceptable)
		w.Write([]byte("Storage quota exceeded"))
		log.Warn("User " + owner + " exceeded storage quota, rejecting upload")
		return
	}

	out, err := os.Create(config.Storage.Path + header.Filename)
	if log.Check(log.WarnLevel, "Unable to create the file for writing", err) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("Cannot create file"))
		return
	}
	defer out.Close()

	limit := int64(db.QuotaLeft(owner))
	f := io.LimitReader(file, limit)

	// write the content from POST to the file
	if copied, err := io.Copy(out, f); copied == limit || err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("Failed to write file or storage quota exceeded"))
		log.Warn("User " + owner + " exceeded storage quota, removing file")
		os.Remove(config.Storage.Path + header.Filename)
		return
	}

	hash = genHash(config.Storage.Path + header.Filename)
	if len(hash) == 0 {
		log.Warn("Failed to calculate hash for " + header.Filename)
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("Failed to calculate hash"))
		return
	}

	os.Rename(config.Storage.Path+header.Filename, config.Storage.Path+hash)
	log.Info("File uploaded successfully: " + header.Filename + "(" + hash + ")")

	return hash, owner
}

func genHash(file string) string {
	f, err := os.Open(file)
	log.Check(log.WarnLevel, "Opening file"+file, err)
	defer f.Close()

	hash := md5.New()
	if _, err := io.Copy(hash, f); err != nil {
		return ""
	}
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func Delete(w http.ResponseWriter, r *http.Request) string {
	hash := r.URL.Query().Get("id")
	token := r.URL.Query().Get("token")
	if len(hash) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("Empty file id"))
		log.Warn(r.RemoteAddr + " - empty file id")
		return ""
	}
	if len(token) == 0 {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("Empty token"))
		log.Warn(r.RemoteAddr + " - empty token")
		return ""
	}
	user := db.CheckToken(token)
	info := db.Info(hash)
	if len(info) == 0 {
		log.Warn("File not found by hash")
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("File not found"))
		return ""
	}

	if !db.CheckOwner(user, hash) {
		log.Warn("File " + info["name"] + "(" + hash + ") is not owned by " + user + ", rejecting deletion request")
		w.WriteHeader(http.StatusForbidden)
		w.Write([]byte("File " + info["name"] + " is not owned by " + user))
		return ""
	}

	f, _ := os.Stat(config.Storage.Path + hash)
	db.QuotaUsageSet(user, -int(f.Size()))

	if db.Delete(user, hash) <= 0 {
		if log.Check(log.WarnLevel, "Removing "+info["name"]+"from disk", os.Remove(config.Storage.Path+hash)) {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("Failed to remove file"))
			return ""
		}
	}

	log.Info("Removing " + info["name"])
	return hash
}

func Share(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		if len(r.FormValue("json")) == 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Empty json"))
			log.Warn("Share request: empty json, nothing to do")
			return
		}
		var data share
		if log.Check(log.WarnLevel, "Parsing share request json", json.Unmarshal([]byte(r.FormValue("json")), &data)) {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Failed to parse json body"))
			return
		}
		if len(data.Token) == 0 || len(db.CheckToken(data.Token)) == 0 {
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte("Not authorized"))
			log.Warn("Empty or invalid token, rejecting share request")
			return
		}
		if len(data.Id) == 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Empty file id"))
			log.Warn("Empty file id, rejecting share request")
			return
		}
		owner := db.CheckToken(data.Token)
		if !db.CheckOwner(owner, data.Id) {
			w.WriteHeader(http.StatusForbidden)
			w.Write([]byte("File is not owned by authorized user"))
			log.Warn("User tried to share another's file, rejecting")
			return
		}
		for _, v := range data.Add {
			db.ShareWith(data.Id, owner, v)
		}
		for _, v := range data.Remove {
			db.UnshareWith(data.Id, owner, v)
		}
	} else if r.Method == "GET" {
		id := r.URL.Query().Get("id")
		if len(id) == 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Empty file id"))
			return
		}
		token := r.URL.Query().Get("token")
		if len(token) == 0 || len(db.CheckToken(token)) == 0 {
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte("Not authorized"))
			return
		}
		owner := db.CheckToken(token)
		if !db.CheckOwner(owner, id) {
			w.WriteHeader(http.StatusForbidden)
			w.Write([]byte("File is not owned by authorized user"))
			log.Warn("User tried to request scope of another's file, rejecting")
			return
		}
		js, _ := json.Marshal(db.GetScope(id, owner))
		w.Write(js)
	}
}

func сheckLength(user, length string) bool {
	l, err := strconv.Atoi(length)
	if err != nil || len(length) == 0 {
		log.Warn("Empty or invalid content length")
		return true
	}

	if l > db.QuotaLeft(user) {
		return false
	}
	db.QuotaUsageSet(user, l)
	return true
}

func QuotaTest(w http.ResponseWriter, r *http.Request) {
	user := r.URL.Query().Get("user")
	if len(user) != 0 {
		w.Write([]byte("Quota: " + strconv.Itoa(db.QuotaGet(user)) + "b\n"))
		w.Write([]byte("Quota used: " + strconv.Itoa(db.QuotaUsageGet(user)) + "b\n"))
		w.Write([]byte("Quota left: " + strconv.Itoa(db.QuotaLeft(user)) + "b\n"))
	}

	quota := r.URL.Query().Get("quota")
	if len(quota) != 0 {
		db.QuotaSet(user, quota)
		w.Write([]byte("user: " + user + ", quota: " + quota + "b\n"))
	}
}
<|endoftext|>


<commit_before>package sourcemap

import (
	"bytes"
	"encoding/json"
	"io"
	"sort"
	"strings"
)

type Map struct {
	Version         int      `json:"version"`
	File            string   `json:"file,omitempty"`
	SourceRoot      string   `json:"sourceRoot,omitempty"`
	Sources         []string `json:"sources,omitempty"`
	Names           []string `json:"names,omitempty"`
	Mappings        string   `json:"mappings"`
	decodedMappings []*Mapping
}

type Mapping struct {
	GeneratedLine   int
	GeneratedColumn int
	OriginalFile    string
	OriginalLine    int
	OriginalColumn  int
	OriginalName    string
}

func ReadFrom(r io.Reader) (*Map, error) {
	d := json.NewDecoder(r)
	var m Map
	if err := d.Decode(&m); err != nil {
		return nil, err
	}
	return &m, nil
}

const base64encode = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"

var base64decode [256]int

func init() {
	for i := 0; i < len(base64decode); i++ {
		base64decode[i] = 0xff
	}
	for i := 0; i < len(base64encode); i++ {
		base64decode[base64encode[i]] = i
	}
}

func (m *Map) decodeMappings() {
	if m.decodedMappings != nil {
		return
	}

	r := strings.NewReader(m.Mappings)
	var generatedLine = 1
	var generatedColumn = 0
	var originalFile = 0
	var originalLine = 1
	var originalColumn = 0
	var originalName = 0
	for r.Len() != 0 {
		b, _ := r.ReadByte()
		if b == ',' {
			continue
		}
		if b == ';' {
			generatedLine++
			generatedColumn = 0
			continue
		}
		r.UnreadByte()

		count := 0
		readVLQ := func() int {
			v := 0
			s := uint(0)
			for {
				b, _ := r.ReadByte()
				o := base64decode[b]
				if o == 0xff {
					r.UnreadByte()
					return 0
				}
				v += (o &^ 32) << s
				if o&32 == 0 {
					break
				}
				s += 5
			}
			count++
			if v&1 != 0 {
				return -(v >> 1)
			}
			return v >> 1
		}
		generatedColumn += readVLQ()
		originalFile += readVLQ()
		originalLine += readVLQ()
		originalColumn += readVLQ()
		originalName += readVLQ()

		switch count {
		case 1:
			m.decodedMappings = append(m.decodedMappings, &Mapping{generatedLine, generatedColumn, "", 0, 0, ""})
		case 4:
			m.decodedMappings = append(m.decodedMappings, &Mapping{generatedLine, generatedColumn, m.Sources[originalFile], originalLine, originalColumn, ""})
		case 5:
			m.decodedMappings = append(m.decodedMappings, &Mapping{generatedLine, generatedColumn, m.Sources[originalFile], originalLine, originalColumn, m.Names[originalName]})
		}
	}
}

func (m *Map) DecodedMappings() []*Mapping {
	m.decodeMappings()
	return m.decodedMappings
}

func (m *Map) ClearMappings() {
	m.Mappings = ""
	m.decodedMappings = nil
}

func (m *Map) AddMapping(mapping *Mapping) {
	m.decodedMappings = append(m.decodedMappings, mapping)
}

func (m *Map) Len() int {
	m.decodeMappings()
	return len(m.DecodedMappings())
}

func (m *Map) Less(i, j int) bool {
	a := m.decodedMappings[i]
	b := m.decodedMappings[j]
	return a.GeneratedLine < b.GeneratedLine || (a.GeneratedLine == b.GeneratedLine && a.GeneratedColumn < b.GeneratedColumn)
}

func (m *Map) Swap(i, j int) {
	m.decodedMappings[i], m.decodedMappings[j] = m.decodedMappings[j], m.decodedMappings[i]
}

func (m *Map) EncodeMappings() {
	sort.Sort(m)
	m.Sources = nil
	fileIndexMap := make(map[string]int)
	m.Names = nil
	nameIndexMap := make(map[string]int)
	var generatedLine = 1
	var generatedColumn = 0
	var originalFile = 0
	var originalLine = 1
	var originalColumn = 0
	var originalName = 0
	buf := bytes.NewBuffer(nil)
	comma := false
	for _, mapping := range m.decodedMappings {
		for mapping.GeneratedLine > generatedLine {
			buf.WriteByte(';')
			generatedLine++
			generatedColumn = 0
			comma = false
		}
		if comma {
			buf.WriteByte(',')
		}

		writeVLQ := func(v int) {
			v <<= 1
			if v < 0 {
				v = -v
				v |= 1
			}
			for v >= 32 {
				buf.WriteByte(base64encode[32|(v&31)])
				v >>= 5
			}
			buf.WriteByte(base64encode[v])
		}

		writeVLQ(mapping.GeneratedColumn - generatedColumn)
		generatedColumn = mapping.GeneratedColumn

		if mapping.OriginalFile != "" {
			fileIndex, ok := fileIndexMap[mapping.OriginalFile]
			if !ok {
				fileIndex = len(m.Sources)
				fileIndexMap[mapping.OriginalFile] = fileIndex
				m.Sources = append(m.Sources, mapping.OriginalFile)
			}
			writeVLQ(fileIndex - originalFile)
			originalFile = fileIndex

			writeVLQ(mapping.OriginalLine - originalLine)
			originalLine = mapping.OriginalLine

			writeVLQ(mapping.OriginalColumn - originalColumn)
			originalColumn = mapping.OriginalColumn

			if mapping.OriginalName != "" {
				nameIndex, ok := nameIndexMap[mapping.OriginalName]
				if !ok {
					nameIndex = len(m.Names)
					nameIndexMap[mapping.OriginalName] = nameIndex
					m.Names = append(m.Names, mapping.OriginalName)
				}
				writeVLQ(nameIndex - originalName)
				originalName = nameIndex
			}
		}

		comma = true
	}
	m.Mappings = buf.String()
}

func (m *Map) WriteTo(w io.Writer) error {
	if m.Version == 0 {
		m.Version = 3
	}
	if m.decodedMappings != nil {
		m.EncodeMappings()
	}
	enc := json.NewEncoder(w)
	return enc.Encode(m)
}
<commit_msg>Include names and sources even when empty<commit_after>package sourcemap

import (
	"bytes"
	"encoding/json"
	"io"
	"sort"
	"strings"
)

type Map struct {
	Version         int      `json:"version"`
	File            string   `json:"file,omitempty"`
	SourceRoot      string   `json:"sourceRoot,omitempty"`
	Sources         []string `json:"sources"`
	Names           []string `json:"names"`
	Mappings        string   `json:"mappings"`
	decodedMappings []*Mapping
}

type Mapping struct {
	GeneratedLine   int
	GeneratedColumn int
	OriginalFile    string
	OriginalLine    int
	OriginalColumn  int
	OriginalName    string
}

func ReadFrom(r io.Reader) (*Map, error) {
	d := json.NewDecoder(r)
	var m Map
	if err := d.Decode(&m); err != nil {
		return nil, err
	}
	return &m, nil
}

const base64encode = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"

var base64decode [256]int

func init() {
	for i := 0; i < len(base64decode); i++ {
		base64decode[i] = 0xff
	}
	for i := 0; i < len(base64encode); i++ {
		base64decode[base64encode[i]] = i
	}
}

func (m *Map) decodeMappings() {
	if m.decodedMappings != nil {
		return
	}

	r := strings.NewReader(m.Mappings)
	var generatedLine = 1
	var generatedColumn = 0
	var originalFile = 0
	var originalLine = 1
	var originalColumn = 0
	var originalName = 0
	for r.Len() != 0 {
		b, _ := r.ReadByte()
		if b == ',' {
			continue
		}
		if b == ';' {
			generatedLine++
			generatedColumn = 0
			continue
		}
		r.UnreadByte()

		count := 0
		readVLQ := func() int {
			v := 0
			s := uint(0)
			for {
				b, _ := r.ReadByte()
				o := base64decode[b]
				if o == 0xff {
					r.UnreadByte()
					return 0
				}
				v += (o &^ 32) << s
				if o&32 == 0 {
					break
				}
				s += 5
			}
			count++
			if v&1 != 0 {
				return -(v >> 1)
			}
			return v >> 1
		}
		generatedColumn += readVLQ()
		originalFile += readVLQ()
		originalLine += readVLQ()
		originalColumn += readVLQ()
		originalName += readVLQ()

		switch count {
		case 1:
			m.decodedMappings = append(m.decodedMappings, &Mapping{generatedLine, generatedColumn, "", 0, 0, ""})
		case 4:
			m.decodedMappings = append(m.decodedMappings, &Mapping{generatedLine, generatedColumn, m.Sources[originalFile], originalLine, originalColumn, ""})
		case 5:
			m.decodedMappings = append(m.decodedMappings, &Mapping{generatedLine, generatedColumn, m.Sources[originalFile], originalLine, originalColumn, m.Names[originalName]})
		}
	}
}

func (m *Map) DecodedMappings() []*Mapping {
	m.decodeMappings()
	return m.decodedMappings
}

func (m *Map) ClearMappings() {
	m.Mappings = ""
	m.decodedMappings = nil
}

func (m *Map) AddMapping(mapping *Mapping) {
	m.decodedMappings = append(m.decodedMappings, mapping)
}

func (m *Map) Len() int {
	m.decodeMappings()
	return len(m.DecodedMappings())
}

func (m *Map) Less(i, j int) bool {
	a := m.decodedMappings[i]
	b := m.decodedMappings[j]
	return a.GeneratedLine < b.GeneratedLine || (a.GeneratedLine == b.GeneratedLine && a.GeneratedColumn < b.GeneratedColumn)
}

func (m *Map) Swap(i, j int) {
	m.decodedMappings[i], m.decodedMappings[j] = m.decodedMappings[j], m.decodedMappings[i]
}

func (m *Map) EncodeMappings() {
	sort.Sort(m)
	m.Sources = nil
	fileIndexMap := make(map[string]int)
	m.Names = nil
	nameIndexMap := make(map[string]int)
	var generatedLine = 1
	var generatedColumn = 0
	var originalFile = 0
	var originalLine = 1
	var originalColumn = 0
	var originalName = 0
	buf := bytes.NewBuffer(nil)
	comma := false
	for _, mapping := range m.decodedMappings {
		for mapping.GeneratedLine > generatedLine {
			buf.WriteByte(';')
			generatedLine++
			generatedColumn = 0
			comma = false
		}
		if comma {
			buf.WriteByte(',')
		}

		writeVLQ := func(v int) {
			v <<= 1
			if v < 0 {
				v = -v
				v |= 1
			}
			for v >= 32 {
				buf.WriteByte(base64encode[32|(v&31)])
				v >>= 5
			}
			buf.WriteByte(base64encode[v])
		}

		writeVLQ(mapping.GeneratedColumn - generatedColumn)
		generatedColumn = mapping.GeneratedColumn

		if mapping.OriginalFile != "" {
			fileIndex, ok := fileIndexMap[mapping.OriginalFile]
			if !ok {
				fileIndex = len(m.Sources)
				fileIndexMap[mapping.OriginalFile] = fileIndex
				m.Sources = append(m.Sources, mapping.OriginalFile)
			}
			writeVLQ(fileIndex - originalFile)
			originalFile = fileIndex

			writeVLQ(mapping.OriginalLine - originalLine)
			originalLine = mapping.OriginalLine

			writeVLQ(mapping.OriginalColumn - originalColumn)
			originalColumn = mapping.OriginalColumn

			if mapping.OriginalName != "" {
				nameIndex, ok := nameIndexMap[mapping.OriginalName]
				if !ok {
					nameIndex = len(m.Names)
					nameIndexMap[mapping.OriginalName] = nameIndex
					m.Names = append(m.Names, mapping.OriginalName)
				}
				writeVLQ(nameIndex - originalName)
				originalName = nameIndex
			}
		}

		comma = true
	}
	m.Mappings = buf.String()
}

func (m *Map) WriteTo(w io.Writer) error {
	if m.Version == 0 {
		m.Version = 3
	}
	if m.decodedMappings != nil {
		m.EncodeMappings()
	}
	enc := json.NewEncoder(w)
	return enc.Encode(m)
}
<|endoftext|>


<commit_before>package es_stdout

import (
	"github.com/mattn/go-colorable"
	"github.com/watermint/toolbox/essentials/runtime/es_env"
	"github.com/watermint/toolbox/essentials/terminal/es_terminfo"
	"github.com/watermint/toolbox/quality/infra/qt_secure"
	"io"
	"io/ioutil"
	"os"
	"sync"
)

type Feature interface {
	IsQuiet() bool
	IsTest() bool
}

func newDefaultOut(test, quiet bool) io.WriteCloser {
	if quiet {
		return &syncOut{co: ioutil.Discard}
	}
	if test {
		if qt_secure.IsSecureEndToEndTest() {
			return &syncOut{co: ioutil.Discard}
		} else {
			return &syncOut{co: os.Stdout}
		}
	} else {
		if es_terminfo.IsOutColorTerminal() {
			return newWriteCloser(colorable.NewColorableStdout())
		} else {
			return newWriteCloser(os.Stdout)
		}
	}
}

func NewDiscard() io.WriteCloser {
	return &syncOut{co: ioutil.Discard}
}

func NewDefaultOut(feature Feature) io.WriteCloser {
	return newDefaultOut(feature.IsTest(), feature.IsQuiet())
}

func NewTestOut() io.WriteCloser {
	return newDefaultOut(true, false)
}

func NewDirectOut() io.WriteCloser {
	return newDefaultOut(false, false)
}

func newWriteCloser(co io.Writer) io.WriteCloser {
	if es_env.IsEnabled("TOOLBOX_ASYNC_CONSOLE") {
		return newAsync(co)
	} else {
		return newSync(co)
	}
}

type AsyncMsg struct {
	Wr      io.Writer
	Payload []byte
}

var (
	asyncQueue         = make(chan *AsyncMsg, 100)
	discardQueue       = make(chan *AsyncMsg)
	asyncQueueLauncher sync.Once
)

func newSync(co io.Writer) io.WriteCloser {
	return &syncOut{
		co: co,
	}
}

type syncOut struct {
	co io.Writer
}

func (z syncOut) Write(p []byte) (n int, err error) {
	return z.co.Write(p)
}

func (z syncOut) Close() error {
	return nil
}

func asyncLoop() {
	for m := range asyncQueue {
		m.Wr.Write(m.Payload)
	}
}
func discardLoop() {
	for range discardQueue {
	}
}

func newAsync(co io.Writer) io.WriteCloser {
	asyncQueueLauncher.Do(func() {
		go asyncLoop()
		go discardLoop()
	})
	return &asyncOut{
		co: co,
	}
}

type asyncOut struct {
	co io.Writer
}

func (z *asyncOut) Write(p []byte) (n int, err error) {
	m := &AsyncMsg{
		Wr:      z.co,
		Payload: p,
	}
	select {
	case asyncQueue <- m:
	case discardQueue <- m:
	}
	return len(p), nil
}

func (z *asyncOut) Close() error {
	return nil
}
<commit_msg>fix #411 : recover on IO error<commit_after>package es_stdout

import (
	"errors"
	"github.com/mattn/go-colorable"
	"github.com/watermint/toolbox/essentials/log/esl"
	"github.com/watermint/toolbox/essentials/terminal/es_terminfo"
	"github.com/watermint/toolbox/quality/infra/qt_secure"
	"io"
	"io/ioutil"
	"os"
)

var (
	ErrorGeneralIOFailure = errors.New("general i/o error")
)

type Feature interface {
	IsQuiet() bool
	IsTest() bool
}

func newDefaultOut(test, quiet bool) io.WriteCloser {
	if quiet {
		return &syncOut{co: ioutil.Discard}
	}
	if test {
		if qt_secure.IsSecureEndToEndTest() {
			return &syncOut{co: ioutil.Discard}
		} else {
			return &syncOut{co: os.Stdout}
		}
	} else {
		if es_terminfo.IsOutColorTerminal() {
			return newWriteCloser(colorable.NewColorableStdout())
		} else {
			return newWriteCloser(os.Stdout)
		}
	}
}

func NewDiscard() io.WriteCloser {
	return &syncOut{co: ioutil.Discard}
}

func NewDefaultOut(feature Feature) io.WriteCloser {
	return newDefaultOut(feature.IsTest(), feature.IsQuiet())
}

func NewTestOut() io.WriteCloser {
	return newDefaultOut(true, false)
}

func NewDirectOut() io.WriteCloser {
	return newDefaultOut(false, false)
}

func newWriteCloser(co io.Writer) io.WriteCloser {
	return newSync(co)
}

func newSync(co io.Writer) io.WriteCloser {
	return &syncOut{
		co: co,
	}
}

type syncOut struct {
	co io.Writer
}

func (z syncOut) Write(p []byte) (n int, err error) {
	// Recovery option for I/O error
	// https://github.com/watermint/toolbox/issues/411
	defer func() {
		if r := recover(); err != nil {
			l := esl.Default()
			l.Debug("Recovery from the syncOut error", esl.Any("r", r))
			if errVal, ok := r.(error); ok {
				err = errVal
			} else {
				err = ErrorGeneralIOFailure
			}
		}
	}()
	return z.co.Write(p)
}

func (z syncOut) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package driver

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/nomad/structs"
)

var (
	reQemuVersion = regexp.MustCompile("QEMU emulator version ([\\d\\.]+),.+")
)

// QemuDriver is a driver for running images via Qemu
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type QemuDriver struct {
	logger *log.Logger
}

// qemuHandle is returned from Start/Open as a handle to the PID
type qemuHandle struct {
	proc   *os.Process
	vmID   string
	waitCh chan error
	doneCh chan struct{}
}

// qemuPID is a struct to map the pid running the process to the vm image on
// disk
type qemuPID struct {
	Pid  int
	VmID string
}

// NewQemuDriver is used to create a new exec driver
func NewQemuDriver(logger *log.Logger) Driver {
	d := &QemuDriver{
		logger: logger,
	}
	return d
}

func (d *QemuDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	outBytes, err := exec.Command("qemu-system-x86_64", "-version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	matches := reQemuVersion.FindStringSubmatch(out)
	if len(matches) != 2 {
		return false, fmt.Errorf("Unable to parse Qemu version string: %#v", matches)
	}

	node.Attributes["driver.qemu"] = "true"
	node.Attributes["driver.qemu.version"] = matches[1]

	return true, nil
}

// Run an existing Qemu image. Start() will pull down an existing, valid Qemu
// image and save it to the Drivers Allocation Dir
func (d *QemuDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	// Get the image source
	source, ok := task.Config["image_source"]
	if !ok || source == "" {
		return nil, fmt.Errorf("Missing source image Qemu driver")
	}

	// Qemu defaults to 128M of RAM for a given VM. Instead, we force users to
	// supply a memory size in the tasks resources
	if task.Resources == nil || task.Resources.MemoryMB == 0 {
		return nil, fmt.Errorf("Missing required Task Resource: Memory")
	}

	// Attempt to download the thing
	// Should be extracted to some kind of Http Fetcher
	// Right now, assume publicly accessible HTTP url
	resp, err := http.Get(source)
	if err != nil {
		return nil, fmt.Errorf("Error downloading source for Qemu driver: %s", err)
	}

	// Create a location in the AllocDir to download and store the image.
	// TODO: Caching
	vmID := fmt.Sprintf("qemu-vm-%s-%s", structs.GenerateUUID(), filepath.Base(source))
	fPath := filepath.Join(ctx.AllocDir, vmID)
	vmPath, err := os.OpenFile(fPath, os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		return nil, fmt.Errorf("Error opening file to download too: %s", err)
	}
	defer vmPath.Close()

	// Copy remote file to local AllocDir for execution
	// TODO: a retry of sort if io.Copy fails, for large binaries
	_, ioErr := io.Copy(vmPath, resp.Body)
	if ioErr != nil {
		return nil, fmt.Errorf("Error copying jar from source: %s", ioErr)
	}

	// compute and check checksum
	if check, ok := task.Config["checksum"]; ok {
		d.logger.Printf("[DEBUG] Running checksum on (%s)", vmID)
		hasher := sha256.New()
		file, err := os.Open(vmPath.Name())
		if err != nil {
			return nil, fmt.Errorf("Failed to open file for checksum")
		}

		io.Copy(hasher, file)

		sum := hex.EncodeToString(hasher.Sum(nil))
		if sum != check {
			return nil, fmt.Errorf(
				"Error in Qemu: checksums did not match.\nExpected (%s), got (%s)",
				check,
				sum)
		}
	}

	// Parse configuration arguments
	// Create the base arguments
	accelerator := "tcg"
	mem := "512M"
	if acc, ok := task.Config["accelerator"]; ok {
		accelerator = acc
	}
	if m, ok := task.Config["memory"]; ok {
		mem = m
	}

	args := []string{
		"qemu-system-x86_64",
		"-machine", "type=pc,accel=" + accelerator,
		"-name", vmID,
		"-m", mem,
		"-drive", "file=" + vmPath.Name(),
		"-nodefconfig",
		"-nodefaults",
		"-nographic",
	}

	// TODO: Consolidate these into map of host/guest port when we have HCL
	// Note: Host port must be open and available
	if task.Config["guest_port"] != "" && task.Config["host_port"] != "" {
		args = append(args,
			"-netdev",
			fmt.Sprintf("user,id=user.0,hostfwd=tcp::%s-:%s",
				task.Config["host_port"],
				task.Config["guest_port"]),
			"-device", "virtio-net,netdev=user.0",
		)
	}

	// If using KVM, add optimization args
	if accelerator == "kvm" {
		args = append(args,
			"-enable-kvm",
			"-cpu", "host",
			// Do we have cores information available to the Driver?
			// "-smp", fmt.Sprintf("%d", cores),
		)
	}

	// Start Qemu
	var outBuf, errBuf bytes.Buffer
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdout = &outBuf
	cmd.Stderr = &errBuf

	d.logger.Printf("[DEBUG] Starting QemuVM command: %q", strings.Join(args, " "))
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf(
			"Error running QEMU: %s\n\nOutput: %s\n\nError: %s",
			err, outBuf.String(), errBuf.String())
	}

	d.logger.Printf("[INFO] Started new QemuVM: %s", vmID)

	// Create and Return Handle
	h := &qemuHandle{
		proc:   cmd.Process,
		vmID:   vmPath.Name(),
		doneCh: make(chan struct{}),
		waitCh: make(chan error, 1),
	}

	go h.run()
	return h, nil
}

func (d *QemuDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	// Parse the handle
	pidBytes := []byte(strings.TrimPrefix(handleID, "QEMU:"))
	qpid := &qemuPID{}
	if err := json.Unmarshal(pidBytes, qpid); err != nil {
		return nil, fmt.Errorf("failed to parse Qemu handle '%s': %v", handleID, err)
	}

	// Find the process
	proc, err := os.FindProcess(qpid.Pid)
	if proc == nil || err != nil {
		return nil, fmt.Errorf("failed to find Qemu PID %d: %v", qpid.Pid, err)
	}

	// Return a driver handle
	h := &qemuHandle{
		proc:   proc,
		vmID:   qpid.VmID,
		doneCh: make(chan struct{}),
		waitCh: make(chan error, 1),
	}

	go h.run()
	return h, nil
}

func (h *qemuHandle) ID() string {
	// Return a handle to the PID
	pid := &qemuPID{
		Pid:  h.proc.Pid,
		VmID: h.vmID,
	}
	data, err := json.Marshal(pid)
	if err != nil {
		log.Printf("[ERROR] failed to marshal Qemu PID to JSON: %s", err)
	}
	return fmt.Sprintf("QEMU:%s", string(data))
}

func (h *qemuHandle) WaitCh() chan error {
	return h.waitCh
}

func (h *qemuHandle) Update(task *structs.Task) error {
	// Update is not possible
	return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
//
// TODO: allow a 'shutdown_command' that can be executed over a ssh connection
// to the VM
func (h *qemuHandle) Kill() error {
	h.proc.Signal(os.Interrupt)
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(5 * time.Second):
		return h.proc.Kill()
	}
}

func (h *qemuHandle) run() {
	ps, err := h.proc.Wait()
	close(h.doneCh)
	if err != nil {
		h.waitCh <- err
	} else if !ps.Success() {
		h.waitCh <- fmt.Errorf("task exited with error")
	}
	close(h.waitCh)
}
<commit_msg>driver/qemu: Actually use the requred Memory from the Task<commit_after>package driver

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/nomad/structs"
)

var (
	reQemuVersion = regexp.MustCompile("QEMU emulator version ([\\d\\.]+),.+")
)

// QemuDriver is a driver for running images via Qemu
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type QemuDriver struct {
	logger *log.Logger
}

// qemuHandle is returned from Start/Open as a handle to the PID
type qemuHandle struct {
	proc   *os.Process
	vmID   string
	waitCh chan error
	doneCh chan struct{}
}

// qemuPID is a struct to map the pid running the process to the vm image on
// disk
type qemuPID struct {
	Pid  int
	VmID string
}

// NewQemuDriver is used to create a new exec driver
func NewQemuDriver(logger *log.Logger) Driver {
	d := &QemuDriver{
		logger: logger,
	}
	return d
}

func (d *QemuDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	outBytes, err := exec.Command("qemu-system-x86_64", "-version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	matches := reQemuVersion.FindStringSubmatch(out)
	if len(matches) != 2 {
		return false, fmt.Errorf("Unable to parse Qemu version string: %#v", matches)
	}

	node.Attributes["driver.qemu"] = "true"
	node.Attributes["driver.qemu.version"] = matches[1]

	return true, nil
}

// Run an existing Qemu image. Start() will pull down an existing, valid Qemu
// image and save it to the Drivers Allocation Dir
func (d *QemuDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	// Get the image source
	source, ok := task.Config["image_source"]
	if !ok || source == "" {
		return nil, fmt.Errorf("Missing source image Qemu driver")
	}

	// Qemu defaults to 128M of RAM for a given VM. Instead, we force users to
	// supply a memory size in the tasks resources
	if task.Resources == nil || task.Resources.MemoryMB == 0 {
		return nil, fmt.Errorf("Missing required Task Resource: Memory")
	}

	// Attempt to download the thing
	// Should be extracted to some kind of Http Fetcher
	// Right now, assume publicly accessible HTTP url
	resp, err := http.Get(source)
	if err != nil {
		return nil, fmt.Errorf("Error downloading source for Qemu driver: %s", err)
	}

	// Create a location in the AllocDir to download and store the image.
	// TODO: Caching
	vmID := fmt.Sprintf("qemu-vm-%s-%s", structs.GenerateUUID(), filepath.Base(source))
	fPath := filepath.Join(ctx.AllocDir, vmID)
	vmPath, err := os.OpenFile(fPath, os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		return nil, fmt.Errorf("Error opening file to download too: %s", err)
	}
	defer vmPath.Close()

	// Copy remote file to local AllocDir for execution
	// TODO: a retry of sort if io.Copy fails, for large binaries
	_, ioErr := io.Copy(vmPath, resp.Body)
	if ioErr != nil {
		return nil, fmt.Errorf("Error copying jar from source: %s", ioErr)
	}

	// compute and check checksum
	if check, ok := task.Config["checksum"]; ok {
		d.logger.Printf("[DEBUG] Running checksum on (%s)", vmID)
		hasher := sha256.New()
		file, err := os.Open(vmPath.Name())
		if err != nil {
			return nil, fmt.Errorf("Failed to open file for checksum")
		}

		io.Copy(hasher, file)

		sum := hex.EncodeToString(hasher.Sum(nil))
		if sum != check {
			return nil, fmt.Errorf(
				"Error in Qemu: checksums did not match.\nExpected (%s), got (%s)",
				check,
				sum)
		}
	}

	// Parse configuration arguments
	// Create the base arguments
	accelerator := "tcg"
	if acc, ok := task.Config["accelerator"]; ok {
		accelerator = acc
	}
	// TODO: Check a lower bounds, e.g. the default 128 of Qemu
	mem := fmt.Sprintf("%dM", task.Resources.MemoryMB)

	args := []string{
		"qemu-system-x86_64",
		"-machine", "type=pc,accel=" + accelerator,
		"-name", vmID,
		"-m", mem,
		"-drive", "file=" + vmPath.Name(),
		"-nodefconfig",
		"-nodefaults",
		"-nographic",
	}

	// TODO: Consolidate these into map of host/guest port when we have HCL
	// Note: Host port must be open and available
	if task.Config["guest_port"] != "" && task.Config["host_port"] != "" {
		args = append(args,
			"-netdev",
			fmt.Sprintf("user,id=user.0,hostfwd=tcp::%s-:%s",
				task.Config["host_port"],
				task.Config["guest_port"]),
			"-device", "virtio-net,netdev=user.0",
		)
	}

	// If using KVM, add optimization args
	if accelerator == "kvm" {
		args = append(args,
			"-enable-kvm",
			"-cpu", "host",
			// Do we have cores information available to the Driver?
			// "-smp", fmt.Sprintf("%d", cores),
		)
	}

	// Start Qemu
	var outBuf, errBuf bytes.Buffer
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdout = &outBuf
	cmd.Stderr = &errBuf

	d.logger.Printf("[DEBUG] Starting QemuVM command: %q", strings.Join(args, " "))
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf(
			"Error running QEMU: %s\n\nOutput: %s\n\nError: %s",
			err, outBuf.String(), errBuf.String())
	}

	d.logger.Printf("[INFO] Started new QemuVM: %s", vmID)

	// Create and Return Handle
	h := &qemuHandle{
		proc:   cmd.Process,
		vmID:   vmPath.Name(),
		doneCh: make(chan struct{}),
		waitCh: make(chan error, 1),
	}

	go h.run()
	return h, nil
}

func (d *QemuDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	// Parse the handle
	pidBytes := []byte(strings.TrimPrefix(handleID, "QEMU:"))
	qpid := &qemuPID{}
	if err := json.Unmarshal(pidBytes, qpid); err != nil {
		return nil, fmt.Errorf("failed to parse Qemu handle '%s': %v", handleID, err)
	}

	// Find the process
	proc, err := os.FindProcess(qpid.Pid)
	if proc == nil || err != nil {
		return nil, fmt.Errorf("failed to find Qemu PID %d: %v", qpid.Pid, err)
	}

	// Return a driver handle
	h := &qemuHandle{
		proc:   proc,
		vmID:   qpid.VmID,
		doneCh: make(chan struct{}),
		waitCh: make(chan error, 1),
	}

	go h.run()
	return h, nil
}

func (h *qemuHandle) ID() string {
	// Return a handle to the PID
	pid := &qemuPID{
		Pid:  h.proc.Pid,
		VmID: h.vmID,
	}
	data, err := json.Marshal(pid)
	if err != nil {
		log.Printf("[ERROR] failed to marshal Qemu PID to JSON: %s", err)
	}
	return fmt.Sprintf("QEMU:%s", string(data))
}

func (h *qemuHandle) WaitCh() chan error {
	return h.waitCh
}

func (h *qemuHandle) Update(task *structs.Task) error {
	// Update is not possible
	return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
//
// TODO: allow a 'shutdown_command' that can be executed over a ssh connection
// to the VM
func (h *qemuHandle) Kill() error {
	h.proc.Signal(os.Interrupt)
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(5 * time.Second):
		return h.proc.Kill()
	}
}

func (h *qemuHandle) run() {
	ps, err := h.proc.Wait()
	close(h.doneCh)
	if err != nil {
		h.waitCh <- err
	} else if !ps.Success() {
		h.waitCh <- fmt.Errorf("task exited with error")
	}
	close(h.waitCh)
}
<|endoftext|>


<commit_before>package config

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/mackerelio/mackerel-agent/logging"
	"github.com/mackerelio/mackerel-agent/util"
)

var configLogger = logging.GetLogger("config")

// `apibase` and `agentName` are set from build flags
var apibase string

func getApibase() string {
	if apibase != "" {
		return apibase
	}
	return "https://mackerel.io"
}

var agentName string

func getAgentName() string {
	if agentName != "" {
		return agentName
	}
	return "mackerel-agent"
}

// Config represents mackerel-agent's configuration file.
type Config struct {
	Apibase     string
	Apikey      string
	Root        string
	Pidfile     string
	Conffile    string
	Roles       []string
	Verbose     bool
	Silent      bool
	Diagnostic  bool `toml:"diagnostic"`
	Connection  ConnectionConfig
	DisplayName string      `toml:"display_name"`
	HostStatus  HostStatus  `toml:"host_status"`
	Filesystems Filesystems `toml:"filesystems"`
	HTTPProxy   string      `toml:"http_proxy"`

	// Corresponds to the set of [plugin.<kind>.<name>] sections
	// the key of the map is <kind>, which should be one of "metrics" or "checks".
	Plugin map[string]PluginConfigs

	Include string

	// Cannot exist in configuration files
	HostIDStorage HostIDStorage
	MetricPlugins map[string]*MetricPlugin
	CheckPlugins  map[string]*CheckPlugin
}

// PluginConfigs represents a set of [plugin.<kind>.<name>] sections in the configuration file
// under a specific <kind>. The key of the map is <name>, for example "mysql" of "plugin.metrics.mysql".
type PluginConfigs map[string]*PluginConfig

// PluginConfig represents a section of [plugin.*].
// `MaxCheckAttempts`, `NotificationInterval` and `CheckInterval` options are used with check monitoring plugins. Custom metrics plugins ignore these options.
// `User` option is ignore in windows
type PluginConfig struct {
	CommandRaw           interface{} `toml:"command"`
	Command              string
	CommandArgs          []string
	User                 string
	NotificationInterval *int32  `toml:"notification_interval"`
	CheckInterval        *int32  `toml:"check_interval"`
	MaxCheckAttempts     *int32  `toml:"max_check_attempts"`
	CustomIdentifier     *string `toml:"custom_identifier"`
}

func makeMetricPlugin(pconf *PluginConfig) (*MetricPlugin, error) {
	err := pconf.prepareCommand()
	if err != nil {
		return nil, err
	}
	return &MetricPlugin{
		Command:          pconf.Command,
		CommandArgs:      pconf.CommandArgs,
		User:             pconf.User,
		CustomIdentifier: pconf.CustomIdentifier,
	}, nil
}

func makeCheckPlugin(pconf *PluginConfig) (*CheckPlugin, error) {
	err := pconf.prepareCommand()
	if err != nil {
		return nil, err
	}
	return &CheckPlugin{
		Command:              pconf.Command,
		CommandArgs:          pconf.CommandArgs,
		User:                 pconf.User,
		NotificationInterval: pconf.NotificationInterval,
		CheckInterval:        pconf.CheckInterval,
		MaxCheckAttempts:     pconf.MaxCheckAttempts,
	}, nil
}

// MetricPlugin represents the configuration of a metric plugin
// The `User` option is ignored in Windows
type MetricPlugin struct {
	Command          string
	CommandArgs      []string
	User             string
	CustomIdentifier *string
}

// Run the plugin
func (pconf *MetricPlugin) Run() (string, string, int, error) {
	if len(pconf.CommandArgs) > 0 {
		return util.RunCommandArgs(pconf.CommandArgs, pconf.User)
	}
	return util.RunCommand(pconf.Command, pconf.User)
}

// CommandString returns the command string for log messages
func (pconf *MetricPlugin) CommandString() string {
	if len(pconf.CommandArgs) > 0 {
		return strings.Join(pconf.CommandArgs, " ")
	}
	return pconf.Command
}

// CheckPlugin represents the configuration of a check plugin
// The `User` option is ignored in Windows
type CheckPlugin struct {
	Command              string
	CommandArgs          []string
	User                 string
	NotificationInterval *int32
	CheckInterval        *int32
	MaxCheckAttempts     *int32
}

// Run the plugin
func (pconf *CheckPlugin) Run() (string, string, int, error) {
	if len(pconf.CommandArgs) > 0 {
		return util.RunCommandArgs(pconf.CommandArgs, pconf.User)
	}
	return util.RunCommand(pconf.Command, pconf.User)
}

func (pconf *PluginConfig) prepareCommand() error {
	const errFmt = "failed to prepare plugin command. A configuration value of `command` should be string or string slice, but %T"
	v := pconf.CommandRaw
	switch t := v.(type) {
	case string:
		pconf.Command = t
	case []interface{}:
		if len(t) > 0 {
			for _, vv := range t {
				str, ok := vv.(string)
				if !ok {
					return fmt.Errorf(errFmt, v)
				}
				pconf.CommandArgs = append(pconf.CommandArgs, str)
			}
		} else {
			return fmt.Errorf(errFmt, v)
		}
	case []string:
		pconf.CommandArgs = t
	default:
		return fmt.Errorf(errFmt, v)
	}
	return nil
}

const postMetricsDequeueDelaySecondsMax = 59   // max delay seconds for dequeuing from buffer queue
const postMetricsRetryDelaySecondsMax = 3 * 60 // max delay seconds for retrying a request that caused errors

// PostMetricsInterval XXX
var PostMetricsInterval = 1 * time.Minute

// ConnectionConfig XXX
type ConnectionConfig struct {
	PostMetricsDequeueDelaySeconds int `toml:"post_metrics_dequeue_delay_seconds"` // delay for dequeuing from buffer queue
	PostMetricsRetryDelaySeconds   int `toml:"post_metrics_retry_delay_seconds"`   // delay for retrying a request that caused errors
	PostMetricsRetryMax            int `toml:"post_metrics_retry_max"`             // max numbers of retries for a request that causes errors
	PostMetricsBufferSize          int `toml:"post_metrics_buffer_size"`           // max numbers of requests stored in buffer queue.
}

// HostStatus configure host status on agent start/stop
type HostStatus struct {
	OnStart string `toml:"on_start"`
	OnStop  string `toml:"on_stop"`
}

// Filesystems configure filesystem related settings
type Filesystems struct {
	Ignore        Regexpwrapper `toml:"ignore"`
	UseMountpoint bool          `toml:"use_mountpoint"`
}

// Regexpwrapper is a wrapper type for marshalling string
type Regexpwrapper struct {
	*regexp.Regexp
}

// UnmarshalText for compiling regexp string while loading toml
func (r *Regexpwrapper) UnmarshalText(text []byte) error {
	var err error
	r.Regexp, err = regexp.Compile(string(text))
	return err
}

// CheckNames return list of plugin.checks._name_
func (conf *Config) CheckNames() []string {
	checks := []string{}
	for name := range conf.Plugin["checks"] {
		checks = append(checks, name)
	}
	return checks
}

// LoadConfig XXX
func LoadConfig(conffile string) (*Config, error) {
	config, err := loadConfigFile(conffile)
	if err != nil {
		return nil, err
	}

	// set default values if config does not have values
	if config.Apibase == "" {
		config.Apibase = DefaultConfig.Apibase
	}
	if config.Root == "" {
		config.Root = DefaultConfig.Root
	}
	if config.Pidfile == "" {
		config.Pidfile = DefaultConfig.Pidfile
	}
	if config.Verbose == false {
		config.Verbose = DefaultConfig.Verbose
	}
	if config.Diagnostic == false {
		config.Diagnostic = DefaultConfig.Diagnostic
	}
	if config.Connection.PostMetricsDequeueDelaySeconds == 0 {
		config.Connection.PostMetricsDequeueDelaySeconds = DefaultConfig.Connection.PostMetricsDequeueDelaySeconds
	}
	if config.Connection.PostMetricsDequeueDelaySeconds > postMetricsDequeueDelaySecondsMax {
		configLogger.Warningf("'post_metrics_dequese_delay_seconds' is set to %d (Maximum Value).", postMetricsDequeueDelaySecondsMax)
		config.Connection.PostMetricsDequeueDelaySeconds = postMetricsDequeueDelaySecondsMax
	}
	if config.Connection.PostMetricsRetryDelaySeconds == 0 {
		config.Connection.PostMetricsRetryDelaySeconds = DefaultConfig.Connection.PostMetricsRetryDelaySeconds
	}
	if config.Connection.PostMetricsRetryDelaySeconds > postMetricsRetryDelaySecondsMax {
		configLogger.Warningf("'post_metrics_retry_delay_seconds' is set to %d (Maximum Value).", postMetricsRetryDelaySecondsMax)
		config.Connection.PostMetricsRetryDelaySeconds = postMetricsRetryDelaySecondsMax
	}
	if config.Connection.PostMetricsRetryMax == 0 {
		config.Connection.PostMetricsRetryMax = DefaultConfig.Connection.PostMetricsRetryMax
	}
	if config.Connection.PostMetricsBufferSize == 0 {
		config.Connection.PostMetricsBufferSize = DefaultConfig.Connection.PostMetricsBufferSize
	}

	return config, err
}

func loadConfigFile(file string) (*Config, error) {
	config := &Config{}
	if _, err := toml.DecodeFile(file, config); err != nil {
		return config, err
	}

	if config.Include != "" {
		if err := includeConfigFile(config, config.Include); err != nil {
			return config, err
		}
	}

	config.MetricPlugins = make(map[string]*MetricPlugin)
	if pconfs, ok := config.Plugin["metrics"]; ok {
		var err error
		for name, pconf := range pconfs {
			config.MetricPlugins[name], err = makeMetricPlugin(pconf)
			if err != nil {
				return nil, err
			}
		}
	}

	config.CheckPlugins = make(map[string]*CheckPlugin)
	if pconfs, ok := config.Plugin["checks"]; ok {
		var err error
		for name, pconf := range pconfs {
			config.CheckPlugins[name], err = makeCheckPlugin(pconf)
			if err != nil {
				return nil, err
			}
		}
	}
	return config, nil
}

func includeConfigFile(config *Config, include string) error {
	files, err := filepath.Glob(include)
	if err != nil {
		return err
	}

	for _, file := range files {
		// Save current "roles" value and reset it
		// because toml.DecodeFile()-ing on a fulfilled struct
		// produces bizarre array values.
		rolesSaved := config.Roles
		config.Roles = nil

		// Also, save plugin values for later merging
		pluginSaved := map[string]PluginConfigs{}
		for kind, plugins := range config.Plugin {
			pluginSaved[kind] = plugins
		}

		meta, err := toml.DecodeFile(file, &config)
		if err != nil {
			return fmt.Errorf("while loading included config file %s: %s", file, err)
		}

		// If included config does not have "roles" key,
		// use the previous roles configuration value.
		if meta.IsDefined("roles") == false {
			config.Roles = rolesSaved
		}

		for kind, plugins := range config.Plugin {
			for key, conf := range plugins {
				if pluginSaved[kind] == nil {
					pluginSaved[kind] = PluginConfigs{}
				}
				pluginSaved[kind][key] = conf
			}
		}

		config.Plugin = pluginSaved
	}

	return nil
}

func (conf *Config) hostIDStorage() HostIDStorage {
	if conf.HostIDStorage == nil {
		conf.HostIDStorage = &FileSystemHostIDStorage{Root: conf.Root}
	}
	return conf.HostIDStorage
}

// LoadHostID loads the previously saved host id.
func (conf *Config) LoadHostID() (string, error) {
	return conf.hostIDStorage().LoadHostID()
}

// SaveHostID saves the host id, which may be restored by LoadHostID.
func (conf *Config) SaveHostID(id string) error {
	return conf.hostIDStorage().SaveHostID(id)
}

// DeleteSavedHostID deletes the host id saved by SaveHostID.
func (conf *Config) DeleteSavedHostID() error {
	return conf.hostIDStorage().DeleteSavedHostID()
}

// HostIDStorage is an interface which maintains persistency
// of the "Host ID" for the current host where the agent is running on.
// The ID is always generated and given by Mackerel (mackerel.io).
type HostIDStorage interface {
	LoadHostID() (string, error)
	SaveHostID(id string) error
	DeleteSavedHostID() error
}

// FileSystemHostIDStorage is the default HostIDStorage
// which saves/loads the host id using an id file on the local filesystem.
// The file will be located at /var/lib/mackerel-agent/id by default on linux.
type FileSystemHostIDStorage struct {
	Root string
}

const idFileName = "id"

// HostIDFile is the location of the host id file.
func (s FileSystemHostIDStorage) HostIDFile() string {
	return filepath.Join(s.Root, idFileName)
}

// LoadHostID loads the current host ID from the mackerel-agent's id file.
func (s FileSystemHostIDStorage) LoadHostID() (string, error) {
	content, err := ioutil.ReadFile(s.HostIDFile())
	if err != nil {
		return "", err
	}
	return strings.TrimRight(string(content), "\r\n"), nil
}

// SaveHostID saves the host ID to the mackerel-agent's id file.
func (s FileSystemHostIDStorage) SaveHostID(id string) error {
	err := os.MkdirAll(s.Root, 0755)
	if err != nil {
		return err
	}

	file, err := os.Create(s.HostIDFile())
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = file.Write([]byte(id))
	if err != nil {
		return err
	}

	return nil
}

// DeleteSavedHostID deletes the mackerel-agent's id file.
func (s FileSystemHostIDStorage) DeleteSavedHostID() error {
	return os.Remove(s.HostIDFile())
}
<commit_msg>use CheckPlugins in CheckNames<commit_after>package config

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/mackerelio/mackerel-agent/logging"
	"github.com/mackerelio/mackerel-agent/util"
)

var configLogger = logging.GetLogger("config")

// `apibase` and `agentName` are set from build flags
var apibase string

func getApibase() string {
	if apibase != "" {
		return apibase
	}
	return "https://mackerel.io"
}

var agentName string

func getAgentName() string {
	if agentName != "" {
		return agentName
	}
	return "mackerel-agent"
}

// Config represents mackerel-agent's configuration file.
type Config struct {
	Apibase     string
	Apikey      string
	Root        string
	Pidfile     string
	Conffile    string
	Roles       []string
	Verbose     bool
	Silent      bool
	Diagnostic  bool `toml:"diagnostic"`
	Connection  ConnectionConfig
	DisplayName string      `toml:"display_name"`
	HostStatus  HostStatus  `toml:"host_status"`
	Filesystems Filesystems `toml:"filesystems"`
	HTTPProxy   string      `toml:"http_proxy"`

	// Corresponds to the set of [plugin.<kind>.<name>] sections
	// the key of the map is <kind>, which should be one of "metrics" or "checks".
	Plugin map[string]PluginConfigs

	Include string

	// Cannot exist in configuration files
	HostIDStorage HostIDStorage
	MetricPlugins map[string]*MetricPlugin
	CheckPlugins  map[string]*CheckPlugin
}

// PluginConfigs represents a set of [plugin.<kind>.<name>] sections in the configuration file
// under a specific <kind>. The key of the map is <name>, for example "mysql" of "plugin.metrics.mysql".
type PluginConfigs map[string]*PluginConfig

// PluginConfig represents a section of [plugin.*].
// `MaxCheckAttempts`, `NotificationInterval` and `CheckInterval` options are used with check monitoring plugins. Custom metrics plugins ignore these options.
// `User` option is ignore in windows
type PluginConfig struct {
	CommandRaw           interface{} `toml:"command"`
	Command              string
	CommandArgs          []string
	User                 string
	NotificationInterval *int32  `toml:"notification_interval"`
	CheckInterval        *int32  `toml:"check_interval"`
	MaxCheckAttempts     *int32  `toml:"max_check_attempts"`
	CustomIdentifier     *string `toml:"custom_identifier"`
}

func makeMetricPlugin(pconf *PluginConfig) (*MetricPlugin, error) {
	err := pconf.prepareCommand()
	if err != nil {
		return nil, err
	}
	return &MetricPlugin{
		Command:          pconf.Command,
		CommandArgs:      pconf.CommandArgs,
		User:             pconf.User,
		CustomIdentifier: pconf.CustomIdentifier,
	}, nil
}

func makeCheckPlugin(pconf *PluginConfig) (*CheckPlugin, error) {
	err := pconf.prepareCommand()
	if err != nil {
		return nil, err
	}
	return &CheckPlugin{
		Command:              pconf.Command,
		CommandArgs:          pconf.CommandArgs,
		User:                 pconf.User,
		NotificationInterval: pconf.NotificationInterval,
		CheckInterval:        pconf.CheckInterval,
		MaxCheckAttempts:     pconf.MaxCheckAttempts,
	}, nil
}

// MetricPlugin represents the configuration of a metric plugin
// The `User` option is ignored in Windows
type MetricPlugin struct {
	Command          string
	CommandArgs      []string
	User             string
	CustomIdentifier *string
}

// Run the plugin
func (pconf *MetricPlugin) Run() (string, string, int, error) {
	if len(pconf.CommandArgs) > 0 {
		return util.RunCommandArgs(pconf.CommandArgs, pconf.User)
	}
	return util.RunCommand(pconf.Command, pconf.User)
}

// CommandString returns the command string for log messages
func (pconf *MetricPlugin) CommandString() string {
	if len(pconf.CommandArgs) > 0 {
		return strings.Join(pconf.CommandArgs, " ")
	}
	return pconf.Command
}

// CheckPlugin represents the configuration of a check plugin
// The `User` option is ignored in Windows
type CheckPlugin struct {
	Command              string
	CommandArgs          []string
	User                 string
	NotificationInterval *int32
	CheckInterval        *int32
	MaxCheckAttempts     *int32
}

// Run the plugin
func (pconf *CheckPlugin) Run() (string, string, int, error) {
	if len(pconf.CommandArgs) > 0 {
		return util.RunCommandArgs(pconf.CommandArgs, pconf.User)
	}
	return util.RunCommand(pconf.Command, pconf.User)
}

func (pconf *PluginConfig) prepareCommand() error {
	const errFmt = "failed to prepare plugin command. A configuration value of `command` should be string or string slice, but %T"
	v := pconf.CommandRaw
	switch t := v.(type) {
	case string:
		pconf.Command = t
	case []interface{}:
		if len(t) > 0 {
			for _, vv := range t {
				str, ok := vv.(string)
				if !ok {
					return fmt.Errorf(errFmt, v)
				}
				pconf.CommandArgs = append(pconf.CommandArgs, str)
			}
		} else {
			return fmt.Errorf(errFmt, v)
		}
	case []string:
		pconf.CommandArgs = t
	default:
		return fmt.Errorf(errFmt, v)
	}
	return nil
}

const postMetricsDequeueDelaySecondsMax = 59   // max delay seconds for dequeuing from buffer queue
const postMetricsRetryDelaySecondsMax = 3 * 60 // max delay seconds for retrying a request that caused errors

// PostMetricsInterval XXX
var PostMetricsInterval = 1 * time.Minute

// ConnectionConfig XXX
type ConnectionConfig struct {
	PostMetricsDequeueDelaySeconds int `toml:"post_metrics_dequeue_delay_seconds"` // delay for dequeuing from buffer queue
	PostMetricsRetryDelaySeconds   int `toml:"post_metrics_retry_delay_seconds"`   // delay for retrying a request that caused errors
	PostMetricsRetryMax            int `toml:"post_metrics_retry_max"`             // max numbers of retries for a request that causes errors
	PostMetricsBufferSize          int `toml:"post_metrics_buffer_size"`           // max numbers of requests stored in buffer queue.
}

// HostStatus configure host status on agent start/stop
type HostStatus struct {
	OnStart string `toml:"on_start"`
	OnStop  string `toml:"on_stop"`
}

// Filesystems configure filesystem related settings
type Filesystems struct {
	Ignore        Regexpwrapper `toml:"ignore"`
	UseMountpoint bool          `toml:"use_mountpoint"`
}

// Regexpwrapper is a wrapper type for marshalling string
type Regexpwrapper struct {
	*regexp.Regexp
}

// UnmarshalText for compiling regexp string while loading toml
func (r *Regexpwrapper) UnmarshalText(text []byte) error {
	var err error
	r.Regexp, err = regexp.Compile(string(text))
	return err
}

// CheckNames return list of plugin.checks._name_
func (conf *Config) CheckNames() []string {
	checks := []string{}
	for name := range conf.CheckPlugins {
		checks = append(checks, name)
	}
	return checks
}

// LoadConfig XXX
func LoadConfig(conffile string) (*Config, error) {
	config, err := loadConfigFile(conffile)
	if err != nil {
		return nil, err
	}

	// set default values if config does not have values
	if config.Apibase == "" {
		config.Apibase = DefaultConfig.Apibase
	}
	if config.Root == "" {
		config.Root = DefaultConfig.Root
	}
	if config.Pidfile == "" {
		config.Pidfile = DefaultConfig.Pidfile
	}
	if config.Verbose == false {
		config.Verbose = DefaultConfig.Verbose
	}
	if config.Diagnostic == false {
		config.Diagnostic = DefaultConfig.Diagnostic
	}
	if config.Connection.PostMetricsDequeueDelaySeconds == 0 {
		config.Connection.PostMetricsDequeueDelaySeconds = DefaultConfig.Connection.PostMetricsDequeueDelaySeconds
	}
	if config.Connection.PostMetricsDequeueDelaySeconds > postMetricsDequeueDelaySecondsMax {
		configLogger.Warningf("'post_metrics_dequese_delay_seconds' is set to %d (Maximum Value).", postMetricsDequeueDelaySecondsMax)
		config.Connection.PostMetricsDequeueDelaySeconds = postMetricsDequeueDelaySecondsMax
	}
	if config.Connection.PostMetricsRetryDelaySeconds == 0 {
		config.Connection.PostMetricsRetryDelaySeconds = DefaultConfig.Connection.PostMetricsRetryDelaySeconds
	}
	if config.Connection.PostMetricsRetryDelaySeconds > postMetricsRetryDelaySecondsMax {
		configLogger.Warningf("'post_metrics_retry_delay_seconds' is set to %d (Maximum Value).", postMetricsRetryDelaySecondsMax)
		config.Connection.PostMetricsRetryDelaySeconds = postMetricsRetryDelaySecondsMax
	}
	if config.Connection.PostMetricsRetryMax == 0 {
		config.Connection.PostMetricsRetryMax = DefaultConfig.Connection.PostMetricsRetryMax
	}
	if config.Connection.PostMetricsBufferSize == 0 {
		config.Connection.PostMetricsBufferSize = DefaultConfig.Connection.PostMetricsBufferSize
	}

	return config, err
}

func loadConfigFile(file string) (*Config, error) {
	config := &Config{}
	if _, err := toml.DecodeFile(file, config); err != nil {
		return config, err
	}

	if config.Include != "" {
		if err := includeConfigFile(config, config.Include); err != nil {
			return config, err
		}
	}

	config.MetricPlugins = make(map[string]*MetricPlugin)
	if pconfs, ok := config.Plugin["metrics"]; ok {
		var err error
		for name, pconf := range pconfs {
			config.MetricPlugins[name], err = makeMetricPlugin(pconf)
			if err != nil {
				return nil, err
			}
		}
	}

	config.CheckPlugins = make(map[string]*CheckPlugin)
	if pconfs, ok := config.Plugin["checks"]; ok {
		var err error
		for name, pconf := range pconfs {
			config.CheckPlugins[name], err = makeCheckPlugin(pconf)
			if err != nil {
				return nil, err
			}
		}
	}
	return config, nil
}

func includeConfigFile(config *Config, include string) error {
	files, err := filepath.Glob(include)
	if err != nil {
		return err
	}

	for _, file := range files {
		// Save current "roles" value and reset it
		// because toml.DecodeFile()-ing on a fulfilled struct
		// produces bizarre array values.
		rolesSaved := config.Roles
		config.Roles = nil

		// Also, save plugin values for later merging
		pluginSaved := map[string]PluginConfigs{}
		for kind, plugins := range config.Plugin {
			pluginSaved[kind] = plugins
		}

		meta, err := toml.DecodeFile(file, &config)
		if err != nil {
			return fmt.Errorf("while loading included config file %s: %s", file, err)
		}

		// If included config does not have "roles" key,
		// use the previous roles configuration value.
		if meta.IsDefined("roles") == false {
			config.Roles = rolesSaved
		}

		for kind, plugins := range config.Plugin {
			for key, conf := range plugins {
				if pluginSaved[kind] == nil {
					pluginSaved[kind] = PluginConfigs{}
				}
				pluginSaved[kind][key] = conf
			}
		}

		config.Plugin = pluginSaved
	}

	return nil
}

func (conf *Config) hostIDStorage() HostIDStorage {
	if conf.HostIDStorage == nil {
		conf.HostIDStorage = &FileSystemHostIDStorage{Root: conf.Root}
	}
	return conf.HostIDStorage
}

// LoadHostID loads the previously saved host id.
func (conf *Config) LoadHostID() (string, error) {
	return conf.hostIDStorage().LoadHostID()
}

// SaveHostID saves the host id, which may be restored by LoadHostID.
func (conf *Config) SaveHostID(id string) error {
	return conf.hostIDStorage().SaveHostID(id)
}

// DeleteSavedHostID deletes the host id saved by SaveHostID.
func (conf *Config) DeleteSavedHostID() error {
	return conf.hostIDStorage().DeleteSavedHostID()
}

// HostIDStorage is an interface which maintains persistency
// of the "Host ID" for the current host where the agent is running on.
// The ID is always generated and given by Mackerel (mackerel.io).
type HostIDStorage interface {
	LoadHostID() (string, error)
	SaveHostID(id string) error
	DeleteSavedHostID() error
}

// FileSystemHostIDStorage is the default HostIDStorage
// which saves/loads the host id using an id file on the local filesystem.
// The file will be located at /var/lib/mackerel-agent/id by default on linux.
type FileSystemHostIDStorage struct {
	Root string
}

const idFileName = "id"

// HostIDFile is the location of the host id file.
func (s FileSystemHostIDStorage) HostIDFile() string {
	return filepath.Join(s.Root, idFileName)
}

// LoadHostID loads the current host ID from the mackerel-agent's id file.
func (s FileSystemHostIDStorage) LoadHostID() (string, error) {
	content, err := ioutil.ReadFile(s.HostIDFile())
	if err != nil {
		return "", err
	}
	return strings.TrimRight(string(content), "\r\n"), nil
}

// SaveHostID saves the host ID to the mackerel-agent's id file.
func (s FileSystemHostIDStorage) SaveHostID(id string) error {
	err := os.MkdirAll(s.Root, 0755)
	if err != nil {
		return err
	}

	file, err := os.Create(s.HostIDFile())
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = file.Write([]byte(id))
	if err != nil {
		return err
	}

	return nil
}

// DeleteSavedHostID deletes the mackerel-agent's id file.
func (s FileSystemHostIDStorage) DeleteSavedHostID() error {
	return os.Remove(s.HostIDFile())
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"

	cmn "github.com/tendermint/go-common"
	logger "github.com/tendermint/go-logger"
)

var log = logger.New()

func main() {
	var listenAddr string
	var verbose bool

	flag.StringVar(&listenAddr, "-listen-addr", "tcp://0.0.0.0:46670", "HTTP and Websocket server listen address")
	flag.BoolVar(&verbose, "v", false, "verbose logging")

	flag.Usage = func() {
		fmt.Println(`Tendermint monitor watches over one or more Tendermint core
applications, collecting and providing various statistics to the user.

Usage:
	tm-monitor [-v] [--listen-addr="tcp://0.0.0.0:46670"] [endpoints]

Examples:
	# monitor single instance
	tm-monitor localhost:46657

	# monitor a few instances by providing comma-separated list of RPC endpoints
	tm-monitor host1:46657,host2:46657`)
		fmt.Println("Flags:")
		flag.PrintDefaults()
	}

	flag.Parse()

	if flag.NArg() == 0 {
		flag.Usage()
		os.Exit(1)
	}

	if verbose {
		log.SetHandler(logger.LvlFilterHandler(
			logger.LvlDebug,
			logger.BypassHandler(),
		))
	} else {
		log.SetHandler(logger.LvlFilterHandler(
			logger.LvlInfo,
			logger.BypassHandler(),
		))
	}

	m := startMonitor(flag.Arg(0))

	startRPC(listenAddr, m)

	ton := NewTon(m)
	ton.Start()

	cmn.TrapSignal(func() {
		ton.Stop()
		m.Stop()
	})
}

func startMonitor(endpoints string) *Monitor {
	m := NewMonitor()

	for _, e := range strings.Split(endpoints, ",") {
		if err := m.Monitor(NewNode(e)); err != nil {
			log.Crit(err.Error())
			os.Exit(1)
		}
	}

	if err := m.Start(); err != nil {
		log.Crit(err.Error())
		os.Exit(1)
	}

	return m
}
<commit_msg>no-ton flag to disable ton (default to false)<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"

	cmn "github.com/tendermint/go-common"
	logger "github.com/tendermint/go-logger"
	log15 "github.com/tendermint/log15"
)

var log = logger.New()

func main() {
	var listenAddr string
	var verbose, noton bool

	flag.StringVar(&listenAddr, "listen-addr", "tcp://0.0.0.0:46670", "HTTP and Websocket server listen address")
	flag.BoolVar(&verbose, "v", false, "verbose logging")
	flag.BoolVar(&noton, "no-ton", false, "Do not show ton (table of nodes)")

	flag.Usage = func() {
		fmt.Println(`Tendermint monitor watches over one or more Tendermint core
applications, collecting and providing various statistics to the user.

Usage:
	tm-monitor [-v] [-no-ton] [-listen-addr="tcp://0.0.0.0:46670"] [endpoints]

Examples:
	# monitor single instance
	tm-monitor localhost:46657

	# monitor a few instances by providing comma-separated list of RPC endpoints
	tm-monitor host1:46657,host2:46657`)
		fmt.Println("Flags:")
		flag.PrintDefaults()
	}

	flag.Parse()

	if flag.NArg() == 0 {
		flag.Usage()
		os.Exit(1)
	}

	m := startMonitor(flag.Arg(0))

	startRPC(listenAddr, m)

	var ton *Ton
	if !noton {
		logToFile("tm-monitor.log", verbose)
		ton = NewTon(m)
		ton.Start()
	} else {
		logToStdout(verbose)
	}

	cmn.TrapSignal(func() {
		if !noton {
			ton.Stop()
		}
		m.Stop()
	})
}

func startMonitor(endpoints string) *Monitor {
	m := NewMonitor()

	for _, e := range strings.Split(endpoints, ",") {
		if err := m.Monitor(NewNode(e)); err != nil {
			log.Crit(err.Error())
			os.Exit(1)
		}
	}

	if err := m.Start(); err != nil {
		log.Crit(err.Error())
		os.Exit(1)
	}

	return m
}

func logToStdout(verbose bool) {
	if verbose {
		log.SetHandler(logger.LvlFilterHandler(
			logger.LvlDebug,
			logger.BypassHandler(),
		))
	} else {
		log.SetHandler(logger.LvlFilterHandler(
			logger.LvlInfo,
			logger.BypassHandler(),
		))
	}
}

func logToFile(filename string, verbose bool) {
	if verbose {
		log.SetHandler(logger.LvlFilterHandler(
			logger.LvlDebug,
			log15.Must.FileHandler(filename, log15.LogfmtFormat()),
		))
	} else {
		log.SetHandler(logger.LvlFilterHandler(
			logger.LvlInfo,
			log15.Must.FileHandler(filename, log15.LogfmtFormat()),
		))
	}
}
<|endoftext|>