<commit_before>package drive

// This file contains test helper functions.  It is recommended to call all the
// public functions in this file in tests specific to each implementation of
// drive.Client.  This helps reduce duplication, and ensures they have
// uniform behavior.

import (
	"bytes"
	"crypto/rand"
	"sync"
	"testing"

	"github.com/asjoyner/shade"
)

// TestFileRoundTrip is a helper function,
func TestFileRoundTrip(t *testing.T, c Client) {
	testFiles := randChunk(100)

	// Populate testFiles into the client
	for stringSum, file := range testFiles {
		if err := c.PutFile([]byte(stringSum), []byte(file)); err != nil {
			t.Fatal("Failed to put test file: ", err)
		}
	}

	// Populate them all again, which should not return an error.
	for stringSum, file := range testFiles {
		if err := c.PutFile([]byte(stringSum), []byte(file)); err != nil {
			t.Fatal("Failed to put test file a second time: ", err)
		}
	}

	// Get all the files which were populated
	files, err := c.ListFiles()
	if err != nil {
		t.Fatalf("Failed to retrieve file map: %s", err)
	}
	for stringSum := range testFiles {
		var found bool
		for _, returnedSum := range files {
			if bytes.Equal([]byte(stringSum), returnedSum) {
				found = true
			}
		}
		if !found {
			t.Errorf("test file not returned: %x", stringSum)
			t.Logf("%+v\n", files)
		}
	}
}

// TestChunkRoundTrip allocates 100 random chunks of data, stores them in the
// client, then retrieves each one by the its Sum and compares the bytes that
// are returned.
func TestChunkRoundTrip(t *testing.T, c Client) {
	testChunks := randChunk(100)

	// Populate test chunks into the client
	for stringSum, chunk := range testChunks {
		err := c.PutChunk([]byte(stringSum), chunk)
		if err != nil {
			t.Fatalf("Failed to put chunk \"%x\": %s", stringSum, err)
		}
	}

	// Populate them all again, which should not return an error.
	for stringSum, chunk := range testChunks {
		err := c.PutChunk([]byte(stringSum), chunk)
		if err != nil {
			t.Fatalf("Failed to put test chunk a second time \"%x\": %s", stringSum, err)
		}
	}

	// Get each chunk by its Sum
	for stringSum, chunk := range testChunks {
		returnedChunk, err := c.GetChunk([]byte(stringSum))
		if err != nil {
			t.Errorf("Failed to retrieve chunk \"%x\": %s", stringSum, err)
			continue
		}
		if !bytes.Equal(returnedChunk, chunk) {
			t.Errorf("returned chunk for \"%x\"does not match", stringSum)
			// t.Errorf("got %q, want: %q", string(returnedChunk), string(chunk))
		}
	}
}

// TestParallelRoundTrip calls 10 copies of both test functions in parallel, to
// try to tickle race conditions in the implementation.
func TestParallelRoundTrip(t *testing.T, c Client) {
	var wg sync.WaitGroup
	for i := 0; i < 3; i++ {
		wg.Add(1)
		go runAndDone(TestFileRoundTrip, t, c, &wg)
		wg.Add(1)
		go runAndDone(TestChunkRoundTrip, t, c, &wg)
	}
	wg.Wait()
	return
}

func runAndDone(f func(*testing.T, Client), t *testing.T, c Client, wg *sync.WaitGroup) {
	defer wg.Done()
	f(t, c)
}

// Generate some random test chunks
func randChunk(n int) map[string][]byte {
	testChunks := make(map[string][]byte, n)
	for i := 0; i < n; i++ {
		c := make([]byte, 100*256)
		rand.Read(c)
		testChunks[string(shade.Sum(c))] = c
	}
	return testChunks
}
<commit_msg>Tweak the comment in the test helper.<commit_after>package drive

// This file contains test helper functions.  It is recommended to call all the
// public functions in this file in tests specific to each implementation of
// drive.Client.  This helps reduce duplication, and ensures they have
// uniform behavior.

import (
	"bytes"
	"crypto/rand"
	"sync"
	"testing"

	"github.com/asjoyner/shade"
)

// TestFileRoundTrip is a helper function, it allocates 100 random []byte,
// stores them in the provided client as files, retrieves them, and ensures all
// of the files were returned.
func TestFileRoundTrip(t *testing.T, c Client) {
	testFiles := randChunk(100)

	// Populate testFiles into the client
	for stringSum, file := range testFiles {
		if err := c.PutFile([]byte(stringSum), []byte(file)); err != nil {
			t.Fatal("Failed to put test file: ", err)
		}
	}

	// Populate them all again, which should not return an error.
	for stringSum, file := range testFiles {
		if err := c.PutFile([]byte(stringSum), []byte(file)); err != nil {
			t.Fatal("Failed to put test file a second time: ", err)
		}
	}

	// Get all the files which were populated
	files, err := c.ListFiles()
	if err != nil {
		t.Fatalf("Failed to retrieve file map: %s", err)
	}
	for stringSum := range testFiles {
		var found bool
		for _, returnedSum := range files {
			if bytes.Equal([]byte(stringSum), returnedSum) {
				found = true
			}
		}
		if !found {
			t.Errorf("test file not returned: %x", stringSum)
			t.Logf("%+v\n", files)
		}
	}
}

// TestChunkRoundTrip allocates 100 random chunks of data, stores them in the
// client, then retrieves each one by the its Sum and compares the bytes that
// are returned.
func TestChunkRoundTrip(t *testing.T, c Client) {
	testChunks := randChunk(100)

	// Populate test chunks into the client
	for stringSum, chunk := range testChunks {
		err := c.PutChunk([]byte(stringSum), chunk)
		if err != nil {
			t.Fatalf("Failed to put chunk \"%x\": %s", stringSum, err)
		}
	}

	// Populate them all again, which should not return an error.
	for stringSum, chunk := range testChunks {
		err := c.PutChunk([]byte(stringSum), chunk)
		if err != nil {
			t.Fatalf("Failed to put test chunk a second time \"%x\": %s", stringSum, err)
		}
	}

	// Get each chunk by its Sum
	for stringSum, chunk := range testChunks {
		returnedChunk, err := c.GetChunk([]byte(stringSum))
		if err != nil {
			t.Errorf("Failed to retrieve chunk \"%x\": %s", stringSum, err)
			continue
		}
		if !bytes.Equal(returnedChunk, chunk) {
			t.Errorf("returned chunk for \"%x\"does not match", stringSum)
			// t.Errorf("got %q, want: %q", string(returnedChunk), string(chunk))
		}
	}
}

// TestParallelRoundTrip calls 10 copies of both test functions in parallel, to
// try to tickle race conditions in the implementation.
func TestParallelRoundTrip(t *testing.T, c Client) {
	var wg sync.WaitGroup
	for i := 0; i < 3; i++ {
		wg.Add(1)
		go runAndDone(TestFileRoundTrip, t, c, &wg)
		wg.Add(1)
		go runAndDone(TestChunkRoundTrip, t, c, &wg)
	}
	wg.Wait()
	return
}

func runAndDone(f func(*testing.T, Client), t *testing.T, c Client, wg *sync.WaitGroup) {
	defer wg.Done()
	f(t, c)
}

// Generate some random test chunks
func randChunk(n int) map[string][]byte {
	testChunks := make(map[string][]byte, n)
	for i := 0; i < n; i++ {
		c := make([]byte, 100*256)
		rand.Read(c)
		testChunks[string(shade.Sum(c))] = c
	}
	return testChunks
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Monetas.
// Copyright 2016 Daniel Krawisz.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package email

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/jordwest/imap-server/mailstore"
	"github.com/jordwest/imap-server/types"
	"github.com/DanielKrawisz/bmutil/identity"
	"github.com/DanielKrawisz/bmutil/wire"
	"github.com/DanielKrawisz/bmagent/keymgr"
)

// User implements the mailstore.User interface and represents
// a collection of imap folders belonging to a single user.
type User struct {
	username string
	boxes    map[string]*mailbox
	keys     *keymgr.Manager
	server   ServerOps
}

// NewUser creates a User object from the store.
func NewUser(username string, server ServerOps, keys *keymgr.Manager) (*User, error) {

	mboxes := server.Folders()

	if mboxes == nil {
		return nil, errors.New("Invalid user.")
	}

	u := &User{
		username : username,
		boxes:  make(map[string]*mailbox),
		server: server,
	}

	// The user is allowed to save in some mailboxes but not others.
	for _, mbox := range mboxes {
		var name = mbox.Name()
		var mb *mailbox
		var err error
		switch name {
			case DraftsFolderName:
			mb, err = NewDrafts(mbox, keys.Tags())
			default:
			mb, err = NewMailbox(mbox, keys.Tags())
		}
		if err != nil {
			return nil, err
		}
		u.boxes[name] = mb
	}

	return u, nil
}

// NewMailbox adds a new mailbox.
func (u *User) NewMailbox(name string) (Mailbox, error) {
	return nil, errors.New("Not yet implemented.")
}

// Mailboxes returns all the mailboxes. It is part of the IMAPMailbox interface.
func (u *User) Mailboxes() []mailstore.Mailbox {
	mboxes := make([]mailstore.Mailbox, 0, len(u.boxes))
	for _, mbox := range u.boxes {
		mboxes = append(mboxes, mbox)
	}
	return mboxes
}

// MailboxByName returns a mailbox by its name. It is part of the IMAPMailbox
// interface.
func (u *User) MailboxByName(name string) (mailstore.Mailbox, error) {
	// Enforce the case insensitivity of Inbox.
	if strings.ToLower(name) == strings.ToLower(InboxFolderName) {
		name = InboxFolderName
	}

	mbox, ok := u.boxes[name]
	if !ok {
		return nil, errors.New("Not found")
	}
	return mbox, nil
}

// DeliverFromBMNet adds a message received from bmd into the appropriate
// folder.
func (u *User) DeliverFromBMNet(bm *Bitmessage) error {
	// Put message in the right folder.
	return u.boxes[InboxFolderName].AddNew(bm, types.FlagRecent)
}

// DeliverFromSMTP adds a message received via SMTP to the POW queue, if needed,
// and the outbox.
func (u *User) DeliverFromSMTP(bm *Bitmessage) error {
	smtpLog.Trace("Bitmessage received from SMTP")

	// Check for command.
	if commandRegex.Match([]byte(bm.To)) {
		return errors.New("Commands not yet supported.")
	} else {

		// Attempt to run pow on the message and send it off on the network.
		// This will only happen if the pubkey can be found. An error is only
		// returned if the message could not be generated and the pubkey request
		// could not be sent.
		_, err := bm.SubmitPow(u.server)
		if err != nil {
			smtpLog.Error("Unable to submit for proof-of-work: ", err)
			return err
		}

	}

	// Put message in the right folder.
	return u.boxes[OutboxFolderName].AddNew(bm, types.FlagSeen)
}

// DeliverPublicKey takes a public key and attempts to match it with a message.
// If a matching message is found, the message is encoded to the wire format
// and sent to the pow queue.
func (u *User) DeliverPublicKey(address string, public *identity.Public) error {
	outbox := u.boxes[OutboxFolderName]
	var ids []uint64

	// Go through all messages in the Outbox and get IDs of all the matches.
	err := outbox.mbox.ForEachMessage(0, 0, 2, func(id, _ uint64, msg []byte) error {
		bmsg, err := DecodeBitmessage(msg)
		if err != nil { // (Almost) impossible error.
			return err
		}
		// We have a match!
		if bmsg.state.PubkeyRequested == true && bmsg.To == address {
			ids = append(ids, id)
		}
		return nil
	})
	if err != nil {
		return err
	}

	outbox.Lock()
	for _, id := range ids {
		bmsg := outbox.BitmessageByUID(id)
		bmsg.state.PubkeyRequested = false

		outbox.Lock()
		err = outbox.saveBitmessage(bmsg)
		outbox.Unlock()

		if err != nil {
			return err
		}

		if bmsg.state.AckExpected {
			// TODO generate the ack and send it to the pow queue.
		}

		// Add the message to the pow queue.
		_, err := bmsg.SubmitPow(u.server)
		if err != nil {
			return errors.New("Unable to add message to pow queue.")
		}

		// Save Bitmessage with pow index.
		outbox.Lock()
		err = outbox.saveBitmessage(bmsg)
		outbox.Unlock()
		if err != nil {
			return err
		}
	}
	outbox.Unlock()

	return nil
}

// DeliverPow delivers an object that has had pow done on it.
func (u *User) DeliverPow(index uint64, obj *wire.MsgObject) error {
	outbox := u.boxes[OutboxFolderName]

	var idMsg uint64

	// Go through all messages in the Outbox and get IDs of all the matches.
	err := outbox.mbox.ForEachMessage(0, 0, 2, func(id, _ uint64, msg []byte) error {
		bmsg, err := DecodeBitmessage(msg)
		if err != nil { // (Almost) impossible error.
			return err
		}
		// We have a match!
		if bmsg.state.PowIndex == index {
			idMsg = id
			return errors.New("Message found.")
		}
		return nil
	})
	if err == nil {
		return fmt.Errorf("Unable to find message in outbox with POW index %d",
			index)
	}

	bmsg := outbox.BitmessageByUID(idMsg)

	// Select new box for the message.
	var newBoxName string
	if bmsg.state.AckExpected {
		newBoxName = LimboFolderName
	} else {
		newBoxName = SentFolderName
	}
	newBox := u.boxes[newBoxName]

	bmsg.state.PowIndex = 0
	bmsg.state.SendTries++
	bmsg.state.LastSend = time.Now()

	// Move message from Outbox to the new mailbox.
	err = outbox.DeleteBitmessageByUID(idMsg)
	if err != nil {
		return err
	}

	bmsg.ImapData = nil
	return newBox.AddNew(bmsg, types.FlagSeen)
}

// DeliverPowAck delivers an ack message that was generated by the pow queue.
func (u *User) DeliverPowAck() {
	// TODO
}

// DeliverAckReply takes a message ack and marks a message as having been
// received by the recipient.
func (u *User) DeliverAckReply() {
	// TODO
}
<commit_msg>Improve faulty logic in DeliverPow<commit_after>// Copyright (c) 2015 Monetas.
// Copyright 2016 Daniel Krawisz.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package email

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/jordwest/imap-server/mailstore"
	"github.com/jordwest/imap-server/types"
	"github.com/DanielKrawisz/bmutil/identity"
	"github.com/DanielKrawisz/bmutil/wire"
	"github.com/DanielKrawisz/bmagent/keymgr"
)

// User implements the mailstore.User interface and represents
// a collection of imap folders belonging to a single user.
type User struct {
	username string
	boxes    map[string]*mailbox
	keys     *keymgr.Manager
	server   ServerOps
}

// NewUser creates a User object from the store.
func NewUser(username string, server ServerOps, keys *keymgr.Manager) (*User, error) {

	mboxes := server.Folders()

	if mboxes == nil {
		return nil, errors.New("Invalid user.")
	}

	u := &User{
		username : username,
		boxes:  make(map[string]*mailbox),
		server: server,
	}

	// The user is allowed to save in some mailboxes but not others.
	for _, mbox := range mboxes {
		var name = mbox.Name()
		var mb *mailbox
		var err error
		switch name {
			case DraftsFolderName:
			mb, err = NewDrafts(mbox, keys.Tags())
			default:
			mb, err = NewMailbox(mbox, keys.Tags())
		}
		if err != nil {
			return nil, err
		}
		u.boxes[name] = mb
	}

	return u, nil
}

// NewMailbox adds a new mailbox.
func (u *User) NewMailbox(name string) (Mailbox, error) {
	return nil, errors.New("Not yet implemented.")
}

// Mailboxes returns all the mailboxes. It is part of the IMAPMailbox interface.
func (u *User) Mailboxes() []mailstore.Mailbox {
	mboxes := make([]mailstore.Mailbox, 0, len(u.boxes))
	for _, mbox := range u.boxes {
		mboxes = append(mboxes, mbox)
	}
	return mboxes
}

// MailboxByName returns a mailbox by its name. It is part of the IMAPMailbox
// interface.
func (u *User) MailboxByName(name string) (mailstore.Mailbox, error) {
	// Enforce the case insensitivity of Inbox.
	if strings.ToLower(name) == strings.ToLower(InboxFolderName) {
		name = InboxFolderName
	}

	mbox, ok := u.boxes[name]
	if !ok {
		return nil, errors.New("Not found")
	}
	return mbox, nil
}

// DeliverFromBMNet adds a message received from bmd into the appropriate
// folder.
func (u *User) DeliverFromBMNet(bm *Bitmessage) error {
	// Put message in the right folder.
	return u.boxes[InboxFolderName].AddNew(bm, types.FlagRecent)
}

// DeliverFromSMTP adds a message received via SMTP to the POW queue, if needed,
// and the outbox.
func (u *User) DeliverFromSMTP(bm *Bitmessage) error {
	smtpLog.Trace("Bitmessage received from SMTP")

	// Check for command.
	if commandRegex.Match([]byte(bm.To)) {
		return errors.New("Commands not yet supported.")
	} else {

		// Attempt to run pow on the message and send it off on the network.
		// This will only happen if the pubkey can be found. An error is only
		// returned if the message could not be generated and the pubkey request
		// could not be sent.
		_, err := bm.SubmitPow(u.server)
		if err != nil {
			smtpLog.Error("Unable to submit for proof-of-work: ", err)
			return err
		}

	}

	// Put message in the right folder.
	return u.boxes[OutboxFolderName].AddNew(bm, types.FlagSeen)
}

// DeliverPublicKey takes a public key and attempts to match it with a message.
// If a matching message is found, the message is encoded to the wire format
// and sent to the pow queue.
func (u *User) DeliverPublicKey(address string, public *identity.Public) error {
	outbox := u.boxes[OutboxFolderName]
	var ids []uint64

	// Go through all messages in the Outbox and get IDs of all the matches.
	err := outbox.mbox.ForEachMessage(0, 0, 2, func(id, _ uint64, msg []byte) error {
		bmsg, err := DecodeBitmessage(msg)
		if err != nil { // (Almost) impossible error.
			return err
		}
		// We have a match!
		if bmsg.state.PubkeyRequested == true && bmsg.To == address {
			ids = append(ids, id)
		}
		return nil
	})
	if err != nil {
		return err
	}

	outbox.Lock()
	for _, id := range ids {
		bmsg := outbox.BitmessageByUID(id)
		bmsg.state.PubkeyRequested = false

		outbox.Lock()
		err = outbox.saveBitmessage(bmsg)
		outbox.Unlock()

		if err != nil {
			return err
		}

		if bmsg.state.AckExpected {
			// TODO generate the ack and send it to the pow queue.
		}

		// Add the message to the pow queue.
		_, err := bmsg.SubmitPow(u.server)
		if err != nil {
			return errors.New("Unable to add message to pow queue.")
		}

		// Save Bitmessage with pow index.
		outbox.Lock()
		err = outbox.saveBitmessage(bmsg)
		outbox.Unlock()
		if err != nil {
			return err
		}
	}
	outbox.Unlock()

	return nil
}

// DeliverPow delivers an object that has had pow done on it.
func (u *User) DeliverPow(index uint64, obj *wire.MsgObject) error {
	outbox := u.boxes[OutboxFolderName]

	var bmsg *Bitmessage
	var idMsg uint64

	// The error returned by the inner func which signifies that
	// a message was found successfully. This error is used to represent
	// a success case for the query.
	var errMessageFound = errors.New("Message found.")

	// Go through all messages in the Outbox and get IDs of all the matches.
	err := outbox.mbox.ForEachMessage(0, 0, 2, func(id, _ uint64, msg []byte) error {
		var dbErr error
		bmsg, dbErr = DecodeBitmessage(msg)
		if dbErr != nil { // (Almost) impossible error.
			return dbErr
		}
		// We have a match!
		if bmsg.state.PowIndex == index {
			idMsg = id
			return errMessageFound
		}
		return nil
	})
	if err == nil {
		return fmt.Errorf("Unable to find message in outbox with POW index %d",
			index)
	}
	if err != errMessageFound {
		return err
	}

	smtpLog.Trace("pow delivered for messege from " + bmsg.From + " to " + bmsg.To)

	// Select new box for the message.
	var newBoxName string
	if bmsg.state.AckExpected {
		newBoxName = LimboFolderName
	} else {
		newBoxName = SentFolderName
	}
	newBox := u.boxes[newBoxName]

	bmsg.state.PowIndex = 0
	bmsg.state.SendTries++
	bmsg.state.LastSend = time.Now()

	// Move message from Outbox to the new mailbox.
	err = outbox.DeleteBitmessageByUID(idMsg)
	if err != nil {
		return err
	}

	bmsg.ImapData = nil
	return newBox.AddNew(bmsg, types.FlagSeen)
}

// DeliverPowAck delivers an ack message that was generated by the pow queue.
func (u *User) DeliverPowAck() {
	// TODO
}

// DeliverAckReply takes a message ack and marks a message as having been
// received by the recipient.
func (u *User) DeliverAckReply() {
	// TODO
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"

	pubsub "google.golang.org/api/pubsub/v1"

	"github.com/urfave/cli"
)

func main() {
	app := cli.NewApp()
	app.Name = "blocks-gcs-proxy"
	app.Usage = "github.com/groovenauts/blocks-gcs-proxy"
	app.Version = VERSION

	configFlag := cli.StringFlag{
		Name:  "config, c",
		Usage: "Load configuration from `FILE`",
	}
	app.Flags = []cli.Flag{
		configFlag,
	}

	app.Commands = []cli.Command{
		{
			Name:  "check",
			Usage: "Check config file is valid",
			Action: func(c *cli.Context) error {
				LoadAndSetupProcessConfig(c)
				fmt.Println("OK")
				return nil
			},
			Flags: []cli.Flag{
				configFlag,
			},
		},
		{
			Name:  "download",
			Usage: "Download the files from GCS to downloads directory",
			Action: func(c *cli.Context) error {
				config := &ProcessConfig{}
				config.Log = &LogConfig{Level: "debug"}
				config.setup([]string{})
				config.Command.Downloaders = c.Int("downloaders")
				config.Job.Sustainer = &JobSustainerConfig{
					Disabled: true,
				}
				p := setupProcess(config)
				p.setup()
				files := []interface{}{}
				for _, arg := range c.Args() {
					files = append(files, arg)
				}
				job := &Job{
					config:              config.Command,
					downloads_dir:       c.String("downloads_dir"),
					remoteDownloadFiles: files,
					storage:             p.storage,
				}
				err := job.setupDownloadFiles()
				if err != nil {
					return err
				}
				err = job.downloadFiles()
				return err
			},
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "downloads_dir, d",
					Usage: "Path to the directory which has bucket_name/path/to/file",
				},
				cli.IntFlag{
					Name:  "downloaders, n",
					Usage: "Number of downloaders",
					Value: 6,
				},
			},
		},
		{
			Name:  "upload",
			Usage: "Upload the files under uploads directory",
			Action: func(c *cli.Context) error {
				fmt.Printf("Uploading files\n")
				config := &ProcessConfig{}
				config.Log = &LogConfig{Level: "debug"}
				config.setup([]string{})
				config.Command.Uploaders = c.Int("uploaders")
				config.Job.Sustainer = &JobSustainerConfig{
					Disabled: true,
				}
				p := setupProcess(config)
				p.setup()
				job := &Job{
					config:      config.Command,
					uploads_dir: c.String("uploads_dir"),
					storage:     p.storage,
				}
				fmt.Printf("Uploading files under %v\n", job.uploads_dir)
				err := job.uploadFiles()
				return err
			},
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "uploads_dir, d",
					Usage: "Path to the directory which has bucket_name/path/to/file",
				},
				cli.IntFlag{
					Name:  "uploaders, n",
					Usage: "Number of uploaders",
					Value: 6,
				},
			},
		},

		{
			Name:  "exec",
			Usage: "Execute job without download nor upload",
			Action: func(c *cli.Context) error {
				config := LoadAndSetupProcessConfig(c)

				msg_file := c.String("message")
				workspace := c.String("workspace")

				type Msg struct {
					Attributes  map[string]string `json:"attributes"`
					Data        string            `json:"data"`
					MessageId   string            `json:"messageId"`
					PublishTime string            `json:"publishTime"`
					AckId       string            `json:"ackId"`
				}
				var msg Msg

				data, err := ioutil.ReadFile(msg_file)
				if err != nil {
					fmt.Printf("Error to read file %v because of %v\n", msg_file, err)
					os.Exit(1)
				}

				err = json.Unmarshal(data, &msg)
				if err != nil {
					fmt.Printf("Error to parse json file %v because of %v\n", msg_file, err)
					os.Exit(1)
				}

				job := &Job{
					workspace: workspace,
					config:    config.Command,
					message: &JobMessage{
						raw: &pubsub.ReceivedMessage{
							AckId: msg.AckId,
							Message: &pubsub.PubsubMessage{
								Attributes: msg.Attributes,
								Data:       msg.Data,
								MessageId:  msg.MessageId,
								// PublishTime: time.Now().Format(time.RFC3339),
								PublishTime: msg.PublishTime,
							},
						},
					},
				}
				fmt.Printf("Executing job %v\n", job.workspace)
				err = job.uploadFiles()
				return err
			},
			Flags: []cli.Flag{
				configFlag,
				cli.StringFlag{
					Name:  "message, m",
					Usage: "Path to the message json file which has attributes and data",
				},
				cli.StringFlag{
					Name:  "workspace, w",
					Usage: "Path to workspace directory which has downloads and uploads",
				},
			},
		},
	}

	app.Action = run

	app.Run(os.Args)
}

func run(c *cli.Context) error {
	config := LoadAndSetupProcessConfig(c)
	p := setupProcess(config)

	err := p.run()
	if err != nil {
		fmt.Printf("Error to run cause of %v\n", err)
		os.Exit(1)
	}
	return nil
}

func setupProcess(config *ProcessConfig) *Process {
	p := &Process{config: config}
	err := p.setup()
	if err != nil {
		fmt.Printf("Error to setup Process cause of %v\n", err)
		os.Exit(1)
	}
	return p
}

func LoadAndSetupProcessConfig(c *cli.Context) *ProcessConfig {
	path := configPath(c)
	config, err := LoadProcessConfig(path)
	if err != nil {
		fmt.Printf("Error to load %v cause of %v\n", path, err)
		os.Exit(1)
	}
	err = config.setup(os.Args[1:])
	if err != nil {
		fmt.Printf("Error to setup %v cause of %v\n", path, err)
		os.Exit(1)
	}
	return config
}

func configPath(c *cli.Context) string {
	r := c.String("config")
	if r == "" {
		r = "./config.json"
	}
	return r
}
<commit_msg>:+1: Prepare and execute in execute sub command<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"

	pubsub "google.golang.org/api/pubsub/v1"

	"github.com/urfave/cli"
)

func main() {
	app := cli.NewApp()
	app.Name = "blocks-gcs-proxy"
	app.Usage = "github.com/groovenauts/blocks-gcs-proxy"
	app.Version = VERSION

	configFlag := cli.StringFlag{
		Name:  "config, c",
		Usage: "Load configuration from `FILE`",
	}
	app.Flags = []cli.Flag{
		configFlag,
	}

	app.Commands = []cli.Command{
		{
			Name:  "check",
			Usage: "Check config file is valid",
			Action: func(c *cli.Context) error {
				LoadAndSetupProcessConfig(c)
				fmt.Println("OK")
				return nil
			},
			Flags: []cli.Flag{
				configFlag,
			},
		},
		{
			Name:  "download",
			Usage: "Download the files from GCS to downloads directory",
			Action: func(c *cli.Context) error {
				config := &ProcessConfig{}
				config.Log = &LogConfig{Level: "debug"}
				config.setup([]string{})
				config.Command.Downloaders = c.Int("downloaders")
				config.Job.Sustainer = &JobSustainerConfig{
					Disabled: true,
				}
				p := setupProcess(config)
				p.setup()
				files := []interface{}{}
				for _, arg := range c.Args() {
					files = append(files, arg)
				}
				job := &Job{
					config:              config.Command,
					downloads_dir:       c.String("downloads_dir"),
					remoteDownloadFiles: files,
					storage:             p.storage,
				}
				err := job.setupDownloadFiles()
				if err != nil {
					return err
				}
				err = job.downloadFiles()
				return err
			},
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "downloads_dir, d",
					Usage: "Path to the directory which has bucket_name/path/to/file",
				},
				cli.IntFlag{
					Name:  "downloaders, n",
					Usage: "Number of downloaders",
					Value: 6,
				},
			},
		},
		{
			Name:  "upload",
			Usage: "Upload the files under uploads directory",
			Action: func(c *cli.Context) error {
				fmt.Printf("Uploading files\n")
				config := &ProcessConfig{}
				config.Log = &LogConfig{Level: "debug"}
				config.setup([]string{})
				config.Command.Uploaders = c.Int("uploaders")
				config.Job.Sustainer = &JobSustainerConfig{
					Disabled: true,
				}
				p := setupProcess(config)
				p.setup()
				job := &Job{
					config:      config.Command,
					uploads_dir: c.String("uploads_dir"),
					storage:     p.storage,
				}
				fmt.Printf("Uploading files under %v\n", job.uploads_dir)
				err := job.uploadFiles()
				return err
			},
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "uploads_dir, d",
					Usage: "Path to the directory which has bucket_name/path/to/file",
				},
				cli.IntFlag{
					Name:  "uploaders, n",
					Usage: "Number of uploaders",
					Value: 6,
				},
			},
		},

		{
			Name:  "exec",
			Usage: "Execute job without download nor upload",
			Action: func(c *cli.Context) error {
				config := LoadAndSetupProcessConfig(c)

				msg_file := c.String("message")
				workspace := c.String("workspace")

				type Msg struct {
					Attributes  map[string]string `json:"attributes"`
					Data        string            `json:"data"`
					MessageId   string            `json:"messageId"`
					PublishTime string            `json:"publishTime"`
					AckId       string            `json:"ackId"`
				}
				var msg Msg

				data, err := ioutil.ReadFile(msg_file)
				if err != nil {
					fmt.Printf("Error to read file %v because of %v\n", msg_file, err)
					os.Exit(1)
				}

				err = json.Unmarshal(data, &msg)
				if err != nil {
					fmt.Printf("Error to parse json file %v because of %v\n", msg_file, err)
					os.Exit(1)
				}

				job := &Job{
					workspace: workspace,
					config:    config.Command,
					message: &JobMessage{
						raw: &pubsub.ReceivedMessage{
							AckId: msg.AckId,
							Message: &pubsub.PubsubMessage{
								Attributes: msg.Attributes,
								Data:       msg.Data,
								MessageId:  msg.MessageId,
								// PublishTime: time.Now().Format(time.RFC3339),
								PublishTime: msg.PublishTime,
							},
						},
					},
				}
				fmt.Printf("Preparing job\n")
				err = job.prepare()
				if err != nil {
					return err
				}
				fmt.Printf("Executing job\n")
				err = job.execute()
				return err
			},
			Flags: []cli.Flag{
				configFlag,
				cli.StringFlag{
					Name:  "message, m",
					Usage: "Path to the message json file which has attributes and data",
				},
				cli.StringFlag{
					Name:  "workspace, w",
					Usage: "Path to workspace directory which has downloads and uploads",
				},
			},
		},
	}

	app.Action = run

	app.Run(os.Args)
}

func run(c *cli.Context) error {
	config := LoadAndSetupProcessConfig(c)
	p := setupProcess(config)

	err := p.run()
	if err != nil {
		fmt.Printf("Error to run cause of %v\n", err)
		os.Exit(1)
	}
	return nil
}

func setupProcess(config *ProcessConfig) *Process {
	p := &Process{config: config}
	err := p.setup()
	if err != nil {
		fmt.Printf("Error to setup Process cause of %v\n", err)
		os.Exit(1)
	}
	return p
}

func LoadAndSetupProcessConfig(c *cli.Context) *ProcessConfig {
	path := configPath(c)
	config, err := LoadProcessConfig(path)
	if err != nil {
		fmt.Printf("Error to load %v cause of %v\n", path, err)
		os.Exit(1)
	}
	err = config.setup(os.Args[1:])
	if err != nil {
		fmt.Printf("Error to setup %v cause of %v\n", path, err)
		os.Exit(1)
	}
	return config
}

func configPath(c *cli.Context) string {
	r := c.String("config")
	if r == "" {
		r = "./config.json"
	}
	return r
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Qian Qiao
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package cli provides base constructs to build complex command line
// applications.
//
// By nesting components, one can build command line applications like: the
// go command.
//
package cli

import (
	"bytes"
	"context"
	"flag"
	"io"
	"strings"
	"text/template"
)

// Component represents a command line component
type Component struct {
	// Components are the sub-components of the current component
	Components []*Component

	// Run runs the component
	// args are the arguments after the component name
	Run func(ctx context.Context, comp *Component, args []string)

	// UsageLine is the one-line usage message.
	// The first word in the line is taken to be the component name
	UsageLine string

	// Short is the short description of the component
	Short string

	// Long is the longer more detailed description of the component
	Long string

	// flagSet is a set of flags specific to this component
	flagSet *flag.FlagSet
}

// FlagSet returns the set of command line flags
func (c *Component) FlagSet() *flag.FlagSet {
	if nil == c.flagSet {
		c.flagSet = flag.NewFlagSet(c.Name(), flag.ExitOnError)
		c.flagSet.Usage = c.Usage
	}

	return c.flagSet
}

// Name returns the name of the component: the first word in the UsageLine
func (c *Component) Name() string {
	name := c.UsageLine
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

// Runnable returns whether this component is runnable or pure informational
func (c *Component) Runnable() bool {
	return nil != c.Run
}

// SetOutput sets the destination for usage messages.
// If output is nil, stderr is used
func (c *Component) SetOutput(output io.Writer) {
	c.FlagSet().SetOutput(output)
}

var usageTemplate = `
{{- if .component.Runnable -}}
Usage: {{.component.UsageLine}}
{{end}}
{{- if ne (len .component.Long) 0 -}}
{{.component.Long | trim}}
{{end}}
{{- if ne (len .component.Components) 0}}
The components are:
{{- range .component.Components}}
{{- if .Runnable}}
  {{.Name | printf "%-11s"}} {{.Short -}}
{{end -}}
{{end}}
{{end}}
{{- if ne (len .flags) 0}}
The flags are:
{{.flags -}}
{{end}}`

// Usage prints out the usage information
func (c *Component) Usage() {
	output := c.flagSet.Output()

	var buf bytes.Buffer
	c.flagSet.SetOutput(&buf)
	c.flagSet.PrintDefaults()

	c.flagSet.SetOutput(output)

	tmpl(output, usageTemplate, map[string]interface{}{
		"component": c,
		"flags":     buf.String(),
	})
}

// Passthrough is a implementation of the Run function that passes the
// execution through the sub commands
func Passthrough(ctx context.Context, component *Component, args []string) {
	if flag.ErrHelp == component.FlagSet().Parse(args) {
		return
	}

	if component.FlagSet().NArg() < 1 {
		component.FlagSet().Usage()
		return
	}

	name := component.FlagSet().Arg(0)

	for _, comp := range component.Components {
		if name == comp.Name() {
			if comp.Runnable() {
				comp.Run(ctx, comp, component.FlagSet().Args()[1:])
				return
			}
		}
	}
	component.FlagSet().Usage()
}

func tmpl(w io.Writer, text string, data interface{}) {
	t := template.New("top")
	t.Funcs(template.FuncMap{
		"trim": strings.TrimSpace,
	})
	template.Must(t.Parse(text))
	t.Execute(w, data)
}
<commit_msg>Making sure that we always lazy create the flagset, and also making sure that we don't go through FlagSet() call unnecessarily.<commit_after>// Copyright (c) 2017 Qian Qiao
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package cli provides base constructs to build complex command line
// applications.
//
// By nesting components, one can build command line applications like: the
// go command.
//
package cli

import (
	"bytes"
	"context"
	"flag"
	"io"
	"strings"
	"text/template"
)

// Component represents a command line component
type Component struct {
	// Components are the sub-components of the current component
	Components []*Component

	// Run runs the component
	// args are the arguments after the component name
	Run func(ctx context.Context, comp *Component, args []string)

	// UsageLine is the one-line usage message.
	// The first word in the line is taken to be the component name
	UsageLine string

	// Short is the short description of the component
	Short string

	// Long is the longer more detailed description of the component
	Long string

	// flagSet is a set of flags specific to this component
	flagSet *flag.FlagSet
}

// FlagSet returns the set of command line flags
func (c *Component) FlagSet() *flag.FlagSet {
	if nil == c.flagSet {
		c.flagSet = flag.NewFlagSet(c.Name(), flag.ExitOnError)
		c.flagSet.Usage = c.Usage
	}

	return c.flagSet
}

// Name returns the name of the component: the first word in the UsageLine
func (c *Component) Name() string {
	name := c.UsageLine
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

// Runnable returns whether this component is runnable or pure informational
func (c *Component) Runnable() bool {
	return nil != c.Run
}

// SetOutput sets the destination for usage messages.
// If output is nil, stderr is used
func (c *Component) SetOutput(output io.Writer) {
	c.FlagSet().SetOutput(output)
}

var usageTemplate = `
{{- if .component.Runnable -}}
Usage: {{.component.UsageLine}}
{{end}}
{{- if ne (len .component.Long) 0 -}}
{{.component.Long | trim}}
{{end}}
{{- if ne (len .component.Components) 0}}
The components are:
{{- range .component.Components}}
{{- if .Runnable}}
  {{.Name | printf "%-11s"}} {{.Short -}}
{{end -}}
{{end}}
{{end}}
{{- if ne (len .flags) 0}}
The flags are:
{{.flags -}}
{{end}}`

// Usage prints out the usage information
func (c *Component) Usage() {
	flagSet := c.FlagSet()
	output := flagSet.Output()

	// Capture the output of the flagset so that it can be merged with the rest
	// of the message
	var buf bytes.Buffer
	flagSet.SetOutput(&buf)
	flagSet.PrintDefaults()

	flagSet.SetOutput(output)

	tmpl(output, usageTemplate, map[string]interface{}{
		"component": c,
		"flags":     buf.String(),
	})
}

// Passthrough is a implementation of the Run function that passes the
// execution through the sub commands
func Passthrough(ctx context.Context, component *Component, args []string) {
	flagSet := component.FlagSet()

	if flag.ErrHelp == flagSet.Parse(args) {
		return
	}

	if flagSet.NArg() < 1 {
		flagSet.Usage()
		return
	}

	name := flagSet.Arg(0)

	for _, comp := range component.Components {
		if name == comp.Name() {
			if comp.Runnable() {
				comp.Run(ctx, comp, flagSet.Args()[1:])
				return
			}
		}
	}
	flagSet.Usage()
}

func tmpl(w io.Writer, text string, data interface{}) {
	t := template.New("top")
	t.Funcs(template.FuncMap{
		"trim": strings.TrimSpace,
	})
	template.Must(t.Parse(text))
	t.Execute(w, data)
}
<|endoftext|>


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

// +build !nocli

package main

import (
	"fmt"
	"io"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"syscall"

	"github.com/tgrennan/go-gnureadline"
)

func (adm *Adm) CLI() (err error) {
	var prompt, line string
	var inrl bool
	adm.rxq = make(chan *PDU, 16)
	home := os.Getenv("HOME")
	term := os.Getenv("TERM")
	history := filepath.Join(home, ".asnadm_history")
	rc := filepath.Join(home, ".asnadmrc")
	if s := adm.asn.Debug.String(); s == "" {
		prompt = "asnadm: "
	} else {
		prompt = s + "# "
	}
	if _, err = os.Stat(rc); err == nil {
		if err = gnureadline.ReadInitFile(rc); err != nil {
			return err
		}
	}
	if _, err = os.Stat(history); err == nil {
		gnureadline.ReadHistory(history)
	} else {
		err = nil
	}
	gnureadline.StifleHistory(16)
	winch := make(chan os.Signal, 1)
	done := make(Done, 1)
	signal.Notify(winch, syscall.SIGWINCH)
	go func() {
		for {
			select {
			case pdu := <-adm.rxq:
				if inrl {
					println()
					adm.ObjDump(pdu)
					gnureadline.Rl_resize_terminal()
				} else {
					adm.ObjDump(pdu)
				}
			case sig := <-winch:
				if sig == syscall.SIGWINCH {
					gnureadline.Rl_resize_terminal()
				} else if sig == syscall.SIGTERM {
					signal.Stop(winch)
					gnureadline.Rl_reset_terminal(term)
					done <- nil
					runtime.Goexit()
				}
			}
		}
	}()
	for err == nil {
		inrl = true
		line, err = gnureadline.Readline(prompt, true)
		inrl = false
		if err == nil {
			err = adm.cmdLine(line)
			if err != nil && err != io.EOF {
				fmt.Println(err)
				err = nil
			}
		}
	}
	winch <- syscall.SIGTERM
	<-done
	return
}
<commit_msg>Deprecate go-gnureadline fork with upstream fix<commit_after>// Copyright 2014-2015 Apptimist, Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !nocli

package main

import (
	"fmt"
	"io"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"syscall"

	"github.com/rocky/go-gnureadline"
)

func (adm *Adm) CLI() (err error) {
	var prompt, line string
	var inrl bool
	adm.rxq = make(chan *PDU, 16)
	home := os.Getenv("HOME")
	term := os.Getenv("TERM")
	history := filepath.Join(home, ".asnadm_history")
	rc := filepath.Join(home, ".asnadmrc")
	if s := adm.asn.Debug.String(); s == "" {
		prompt = "asnadm: "
	} else {
		prompt = s + "# "
	}
	if _, err = os.Stat(rc); err == nil {
		if err = gnureadline.ReadInitFile(rc); err != nil {
			return err
		}
	}
	if _, err = os.Stat(history); err == nil {
		gnureadline.ReadHistory(history)
	} else {
		err = nil
	}
	gnureadline.StifleHistory(16)
	winch := make(chan os.Signal, 1)
	done := make(Done, 1)
	signal.Notify(winch, syscall.SIGWINCH)
	go func() {
		for {
			select {
			case pdu := <-adm.rxq:
				if inrl {
					println()
					adm.ObjDump(pdu)
					gnureadline.Rl_resize_terminal()
				} else {
					adm.ObjDump(pdu)
				}
			case sig := <-winch:
				if sig == syscall.SIGWINCH {
					gnureadline.Rl_resize_terminal()
				} else if sig == syscall.SIGTERM {
					signal.Stop(winch)
					gnureadline.Rl_reset_terminal(term)
					done <- nil
					runtime.Goexit()
				}
			}
		}
	}()
	for err == nil {
		inrl = true
		line, err = gnureadline.Readline(prompt, true)
		inrl = false
		if err == nil {
			err = adm.cmdLine(line)
			if err != nil && err != io.EOF {
				fmt.Println(err)
				err = nil
			}
		}
	}
	winch <- syscall.SIGTERM
	<-done
	return
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"sync"
	"time"

	"github.com/yosisa/fluxion/event"
	"github.com/yosisa/fluxion/parser"
	"github.com/yosisa/fluxion/plugin"
	"gopkg.in/fsnotify.v1"
)

var posFiles = make(map[string]*PositionFile)

type Config struct {
	Tag          string `codec:"tag"`
	Path         string `codec:"path"`
	PosFile      string `codec:"pos_file"`
	Format       string `codec:"format"`
	TimeKey      string `codec:"time_key"`
	TimeFormat   string `codec:"time_format"`
	TimeZone     string `codec:"timezone"`
	ReadFromHead bool   `codec:"read_from_head"`
}

type TailInput struct {
	conf       Config
	parser     parser.Parser
	timeParser *parser.TimeParser
	pf         *PositionFile
}

func (i *TailInput) Init(f plugin.ConfigFeeder) (err error) {
	if err = f(&i.conf); err != nil {
		return
	}
	if i.conf.TimeKey == "" {
		i.conf.TimeKey = "time"
	}
	i.parser, i.timeParser, err = parser.Get(i.conf.Format, i.conf.TimeFormat, i.conf.TimeZone)
	if err != nil {
		return
	}

	pf, ok := posFiles[i.conf.PosFile]
	if !ok {
		if pf, err = NewPositionFile(i.conf.PosFile); err != nil {
			return
		}
		posFiles[i.conf.PosFile] = pf
	}
	i.pf = pf
	return
}

func (i *TailInput) Start() error {
	pe := i.pf.Get(i.conf.Path)
	pe.ReadFromHead = i.conf.ReadFromHead
	NewWatcher(pe, i.parseLine)
	return nil
}

func (i *TailInput) parseLine(line []byte) {
	v, err := i.parser.Parse(string(line))
	if err != nil {
		return
	}

	var record *event.Record
	if i.conf.TimeKey != "" && i.timeParser != nil {
		if s, ok := v[i.conf.TimeKey].(string); ok {
			t, err := i.timeParser.Parse(s)
			if err == nil {
				delete(v, i.conf.TimeKey)
				record = event.NewRecordWithTime(i.conf.Tag, t, v)
			}
		}
	}
	if record == nil {
		record = event.NewRecord(i.conf.Tag, v)
	}
	plugin.Emit(record)
}

type TailHandler func([]byte)

type Watcher struct {
	pe       *PositionEntry
	fsw      *fsnotify.Watcher
	r        *PositionReader
	handler  TailHandler
	rotating bool
	m        sync.Mutex
}

func NewWatcher(pe *PositionEntry, h TailHandler) *Watcher {
	fsw, err := fsnotify.NewWatcher()
	if err != nil {
		panic(err)
	}
	w := &Watcher{
		pe:      pe,
		fsw:     fsw,
		handler: h,
	}
	w.open()
	go w.eventLoop()
	return w
}

func (w *Watcher) open() {
	w.m.Lock()
	defer w.m.Unlock()

	w.rotating = false
	if w.r != nil {
		w.r.Close()
	}

	r, err := NewPositionReader(w.pe)
	if err != nil {
		plugin.Log.Warning(err, ", wait for creation")
	} else {
		w.r = r
		w.fsw.Add(w.pe.Path)
	}
}

func (w *Watcher) eventLoop() {
	tick := time.Tick(10 * time.Second)
	for {
		select {
		case ev := <-w.fsw.Events:
			plugin.Log.Debug(ev)
			if err := w.Scan(); err != nil {
				plugin.Log.Warning(err)
			}
		case err := <-w.fsw.Errors:
			plugin.Log.Warning(err)
		case <-tick:
			if err := w.Scan(); err != nil {
				plugin.Log.Warning(err)
			}
		}
	}
}

func (w *Watcher) Scan() error {
	// To make Scan run only one thread at a time.
	// Also used to block rotation until current scanning completed.
	w.m.Lock()
	defer w.m.Unlock()

	if !w.rotating && w.pe.IsRotated() {
		plugin.Log.Infof("Rotation detected: %s", w.pe.Path)
		var wait time.Duration
		if w.r != nil {
			wait = 5 * time.Second
		}
		w.rotating = true
		time.AfterFunc(wait, w.open)
	}

	if w.r == nil {
		return nil
	}

	for {
		line, _, err := w.r.ReadLine()
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
		w.handler(line)
	}
	return nil
}

func main() {
	plugin.New(func() plugin.Plugin {
		return &TailInput{}
	}).Run()
}
<commit_msg>Support wildcard path<commit_after>package main

import (
	"io"
	"path/filepath"
	"sync"
	"time"

	"github.com/yosisa/fluxion/event"
	"github.com/yosisa/fluxion/parser"
	"github.com/yosisa/fluxion/plugin"
	"gopkg.in/fsnotify.v1"
)

var posFiles = make(map[string]*PositionFile)

type Config struct {
	Tag          string `codec:"tag"`
	Path         string `codec:"path"`
	PosFile      string `codec:"pos_file"`
	Format       string `codec:"format"`
	TimeKey      string `codec:"time_key"`
	TimeFormat   string `codec:"time_format"`
	TimeZone     string `codec:"timezone"`
	ReadFromHead bool   `codec:"read_from_head"`
}

type TailInput struct {
	conf       Config
	parser     parser.Parser
	timeParser *parser.TimeParser
	pf         *PositionFile
	watchers   map[string]*Watcher
}

func (i *TailInput) Init(f plugin.ConfigFeeder) (err error) {
	i.watchers = make(map[string]*Watcher)
	if err = f(&i.conf); err != nil {
		return
	}
	if i.conf.TimeKey == "" {
		i.conf.TimeKey = "time"
	}
	i.parser, i.timeParser, err = parser.Get(i.conf.Format, i.conf.TimeFormat, i.conf.TimeZone)
	if err != nil {
		return
	}

	pf, ok := posFiles[i.conf.PosFile]
	if !ok {
		if pf, err = NewPositionFile(i.conf.PosFile); err != nil {
			return
		}
		posFiles[i.conf.PosFile] = pf
	}
	i.pf = pf
	return
}

func (i *TailInput) Start() error {
	go i.pathWatcher()
	return nil
}

func (i *TailInput) pathWatcher() {
	tick := time.Tick(1 * time.Minute)
	for {
		files, err := filepath.Glob(i.conf.Path)
		if err != nil {
			plugin.Log.Error(err)
			return
		}

		changes := make(map[string]bool)
		for f := range i.watchers {
			changes[f] = false
		}
		for _, f := range files {
			if _, ok := changes[f]; ok {
				delete(changes, f)
			} else {
				changes[f] = true
			}
		}

		for f, added := range changes {
			if added {
				plugin.Log.Info("Start watching file: ", f)
				pe := i.pf.Get(f)
				pe.ReadFromHead = i.conf.ReadFromHead
				i.watchers[f] = NewWatcher(pe, i.parseLine)
			} else {
				plugin.Log.Info("Stop watching file: ", f)
				i.watchers[f].Close()
				delete(i.watchers, f)
			}
		}

		<-tick
	}
}

func (i *TailInput) parseLine(line []byte) {
	v, err := i.parser.Parse(string(line))
	if err != nil {
		return
	}

	var record *event.Record
	if i.conf.TimeKey != "" && i.timeParser != nil {
		if s, ok := v[i.conf.TimeKey].(string); ok {
			t, err := i.timeParser.Parse(s)
			if err == nil {
				delete(v, i.conf.TimeKey)
				record = event.NewRecordWithTime(i.conf.Tag, t, v)
			}
		}
	}
	if record == nil {
		record = event.NewRecord(i.conf.Tag, v)
	}
	plugin.Emit(record)
}

type TailHandler func([]byte)

type Watcher struct {
	pe       *PositionEntry
	fsw      *fsnotify.Watcher
	r        *PositionReader
	handler  TailHandler
	rotating bool
	m        sync.Mutex
	closeCh  chan struct{}
}

func NewWatcher(pe *PositionEntry, h TailHandler) *Watcher {
	fsw, err := fsnotify.NewWatcher()
	if err != nil {
		panic(err)
	}
	w := &Watcher{
		pe:      pe,
		fsw:     fsw,
		handler: h,
		closeCh: make(chan struct{}),
	}
	w.open()
	go w.eventLoop()
	return w
}

func (w *Watcher) Close() {
	w.fsw.Close()
	close(w.closeCh)
}

func (w *Watcher) open() {
	w.m.Lock()
	defer w.m.Unlock()

	w.rotating = false
	if w.r != nil {
		w.r.Close()
	}

	r, err := NewPositionReader(w.pe)
	if err != nil {
		plugin.Log.Warning(err, ", wait for creation")
		w.fsw.Remove(w.pe.Path)
	} else {
		w.r = r
		w.fsw.Add(w.pe.Path)
	}
}

func (w *Watcher) eventLoop() {
	tick := time.Tick(10 * time.Second)
	for {
		select {
		case <-w.closeCh:
			return
		case ev := <-w.fsw.Events:
			plugin.Log.Debug(ev)
			if err := w.Scan(); err != nil {
				plugin.Log.Warning(err)
			}
		case err := <-w.fsw.Errors:
			plugin.Log.Warning(err)
		case <-tick:
			if err := w.Scan(); err != nil {
				plugin.Log.Warning(err)
			}
		}
	}
}

func (w *Watcher) Scan() error {
	// To make Scan run only one thread at a time.
	// Also used to block rotation until current scanning completed.
	w.m.Lock()
	defer w.m.Unlock()

	if !w.rotating && w.pe.IsRotated() {
		plugin.Log.Infof("Rotation detected: %s", w.pe.Path)
		var wait time.Duration
		if w.r != nil {
			wait = 5 * time.Second
		}
		w.rotating = true
		time.AfterFunc(wait, w.open)
	}

	if w.r == nil {
		return nil
	}

	for {
		line, _, err := w.r.ReadLine()
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
		w.handler(line)
	}
	return nil
}

func main() {
	plugin.New(func() plugin.Plugin {
		return &TailInput{}
	}).Run()
}
<|endoftext|>


<commit_before>package hook

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"testing"

	"github.com/buildkite/agent/v3/bootstrap/shell"
	"github.com/buildkite/agent/v3/env"
	"github.com/stretchr/testify/assert"
)

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

	ctx := context.Background()
	var script []string

	if runtime.GOOS != "windows" {
		script = []string{
			"#!/bin/bash",
			"export LLAMAS=rock",
			"export Alpacas=\"are ok\"",
			"echo hello world",
		}
	} else {
		script = []string{
			"@echo off",
			"set LLAMAS=rock",
			"set Alpacas=are ok",
			"echo hello world",
		}
	}

	wrapper := newTestScriptWrapper(t, script)
	defer os.Remove(wrapper.Path())

	sh := shell.NewTestShell(t)

	if err := sh.RunScript(ctx, wrapper.Path(), nil); err != nil {
		t.Fatal(err)
	}

	changes, err := wrapper.Changes()
	if err != nil {
		t.Fatal(err)
	}

	// Windows’ batch 'SET >' normalises environment variables case so we apply
	// the 'expected' and 'actual' diffs to a blank Environment which handles
	// case normalisation for us
	expected := (&env.Environment{}).Apply(env.Diff{
		Added: map[string]string{
			"LLAMAS":  "rock",
			"Alpacas": "are ok",
		},
		Changed: map[string]env.DiffPair{},
		Removed: map[string]struct{}{},
	})

	actual := (&env.Environment{}).Apply(changes.Diff)

	// The strict equals check here also ensures we aren't bubbling up the
	// internal BUILDKITE_HOOK_EXIT_STATUS and BUILDKITE_HOOK_WORKING_DIR
	// environment variables
	assert.Equal(t, expected, actual)
}

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

	tempDir, err := ioutil.TempDir("", "hookwrapperdir")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tempDir)

	ctx := context.Background()
	var script []string

	if runtime.GOOS != "windows" {
		script = []string{
			"#!/bin/bash",
			"mkdir mysubdir",
			"cd mysubdir",
			"echo hello world",
		}
	} else {
		script = []string{
			"@echo off",
			"mkdir mysubdir",
			"cd mysubdir",
			"echo hello world",
		}
	}

	wrapper := newTestScriptWrapper(t, script)
	defer os.Remove(wrapper.Path())

	sh := shell.NewTestShell(t)
	if err := sh.Chdir(tempDir); err != nil {
		t.Fatal(err)
	}

	if err := sh.RunScript(ctx, wrapper.Path(), nil); err != nil {
		t.Fatal(err)
	}

	changes, err := wrapper.Changes()
	if err != nil {
		t.Fatal(err)
	}

	expected, err := filepath.EvalSymlinks(filepath.Join(tempDir, "mysubdir"))
	if err != nil {
		t.Fatal(err)
	}

	afterWd, err := changes.GetAfterWd()
	if err != nil {
		t.Fatal(err)
	}

	changesDir, err := filepath.EvalSymlinks(afterWd)
	if err != nil {
		t.Fatal(err)
	}

	if changesDir != expected {
		t.Fatalf("Expected working dir of %q, got %q", expected, changesDir)
	}
}

func newTestScriptWrapper(t *testing.T, script []string) *ScriptWrapper {
	hookName := "hookwrapper"
	if runtime.GOOS == "windows" {
		hookName += ".bat"
	}

	hookFile, err := shell.TempFileWithExtension(hookName)
	if err != nil {
		t.Fatal(err)
	}

	for _, line := range script {
		if _, err = fmt.Fprintln(hookFile, line); err != nil {
			t.Fatal(err)
		}
	}

	hookFile.Close()

	wrapper, err := NewScriptWrapper(WithHookPath(hookFile.Name()))
	if err != nil {
		t.Fatal(err)
	}

	return wrapper
}
<commit_msg>Add tests for scriptwrapper script generation<commit_after>package hook

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"testing"

	"github.com/buildkite/agent/v3/bootstrap/shell"
	"github.com/buildkite/agent/v3/env"
	"github.com/stretchr/testify/assert"
)

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

	ctx := context.Background()
	var script []string

	if runtime.GOOS != "windows" {
		script = []string{
			"#!/bin/bash",
			"export LLAMAS=rock",
			"export Alpacas=\"are ok\"",
			"echo hello world",
		}
	} else {
		script = []string{
			"@echo off",
			"set LLAMAS=rock",
			"set Alpacas=are ok",
			"echo hello world",
		}
	}

	wrapper := newTestScriptWrapper(t, script)
	defer os.Remove(wrapper.Path())

	sh := shell.NewTestShell(t)

	if err := sh.RunScript(ctx, wrapper.Path(), nil); err != nil {
		t.Fatal(err)
	}

	changes, err := wrapper.Changes()
	if err != nil {
		t.Fatal(err)
	}

	// Windows’ batch 'SET >' normalises environment variables case so we apply
	// the 'expected' and 'actual' diffs to a blank Environment which handles
	// case normalisation for us
	expected := (&env.Environment{}).Apply(env.Diff{
		Added: map[string]string{
			"LLAMAS":  "rock",
			"Alpacas": "are ok",
		},
		Changed: map[string]env.DiffPair{},
		Removed: map[string]struct{}{},
	})

	actual := (&env.Environment{}).Apply(changes.Diff)

	// The strict equals check here also ensures we aren't bubbling up the
	// internal BUILDKITE_HOOK_EXIT_STATUS and BUILDKITE_HOOK_WORKING_DIR
	// environment variables
	assert.Equal(t, expected, actual)
}

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

	hookFile, err := shell.TempFileWithExtension("hookName.bat")
	assert.NoError(t, err)

	_, err = fmt.Fprintln(hookFile, `echo Hello There!`)
	assert.NoError(t, err)

	hookFile.Close()

	wrapper, err := NewScriptWrapper(
		WithHookPath(hookFile.Name()),
		WithOS("windows"),
	)
	assert.NoError(t, err)

	defer wrapper.Close()

	// The double percent signs %% are sprintf-escaped literal percent signs. Escaping hell is impossible to get out of.
	// See: https://pkg.go.dev/fmt > ctrl-f for "%%"
	scriptTemplate := `@echo off
SETLOCAL ENABLEDELAYEDEXPANSION
SET > "%s"
CALL "%s"
SET BUILDKITE_HOOK_EXIT_STATUS=!ERRORLEVEL!
SET BUILDKITE_HOOK_WORKING_DIR=%%CD%%
SET > "%s"
EXIT %%BUILDKITE_HOOK_EXIT_STATUS%%`

	assertScriptLike(t, scriptTemplate, hookFile.Name(), wrapper)
}

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

	hookFile, err := shell.TempFileWithExtension("hookName.ps1")
	assert.NoError(t, err)

	_, err = fmt.Fprintln(hookFile, `Write-Output "Hello There!"`)
	assert.NoError(t, err)

	hookFile.Close()

	wrapper, err := NewScriptWrapper(
		WithHookPath(hookFile.Name()),
		WithOS("windows"),
	)
	assert.NoError(t, err)

	defer wrapper.Close()

	scriptTemplate := `$ErrorActionPreference = "STOP"
Get-ChildItem Env: | Foreach-Object {"$($_.Name)=$($_.Value)"} | Set-Content "%s"
%s
if ($LASTEXITCODE -eq $null) {$Env:BUILDKITE_HOOK_EXIT_STATUS = 0} else {$Env:BUILDKITE_HOOK_EXIT_STATUS = $LASTEXITCODE}
$Env:BUILDKITE_HOOK_WORKING_DIR = $PWD | Select-Object -ExpandProperty Path
Get-ChildItem Env: | Foreach-Object {"$($_.Name)=$($_.Value)"} | Set-Content "%s"
exit $Env:BUILDKITE_HOOK_EXIT_STATUS`

	assertScriptLike(t, scriptTemplate, hookFile.Name(), wrapper)
}

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

	hookFile, err := shell.TempFileWithExtension("hookName")
	assert.NoError(t, err)

	_, err = fmt.Fprintln(hookFile, `echo "Hello There!"`)
	assert.NoError(t, err)

	hookFile.Close()

	wrapper, err := NewScriptWrapper(
		WithHookPath(hookFile.Name()),
		WithOS("linux"),
	)
	assert.NoError(t, err)

	defer wrapper.Close()

	scriptTemplate := `export -p > "%s"
. "%s"
export BUILDKITE_HOOK_EXIT_STATUS=$?
export BUILDKITE_HOOK_WORKING_DIR=$PWD
export -p > "%s"
exit $BUILDKITE_HOOK_EXIT_STATUS`

	assertScriptLike(t, scriptTemplate, hookFile.Name(), wrapper)
}

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

	tempDir, err := ioutil.TempDir("", "hookwrapperdir")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tempDir)

	ctx := context.Background()
	var script []string

	if runtime.GOOS != "windows" {
		script = []string{
			"#!/bin/bash",
			"mkdir mysubdir",
			"cd mysubdir",
			"echo hello world",
		}
	} else {
		script = []string{
			"@echo off",
			"mkdir mysubdir",
			"cd mysubdir",
			"echo hello world",
		}
	}

	wrapper := newTestScriptWrapper(t, script)
	defer os.Remove(wrapper.Path())

	sh := shell.NewTestShell(t)
	if err := sh.Chdir(tempDir); err != nil {
		t.Fatal(err)
	}

	if err := sh.RunScript(ctx, wrapper.Path(), nil); err != nil {
		t.Fatal(err)
	}

	changes, err := wrapper.Changes()
	if err != nil {
		t.Fatal(err)
	}

	expected, err := filepath.EvalSymlinks(filepath.Join(tempDir, "mysubdir"))
	if err != nil {
		t.Fatal(err)
	}

	afterWd, err := changes.GetAfterWd()
	if err != nil {
		t.Fatal(err)
	}

	changesDir, err := filepath.EvalSymlinks(afterWd)
	if err != nil {
		t.Fatal(err)
	}

	if changesDir != expected {
		t.Fatalf("Expected working dir of %q, got %q", expected, changesDir)
	}
}

func newTestScriptWrapper(t *testing.T, script []string) *ScriptWrapper {
	hookName := "hookwrapper"
	if runtime.GOOS == "windows" {
		hookName += ".bat"
	}

	hookFile, err := shell.TempFileWithExtension(hookName)
	assert.NoError(t, err)

	for _, line := range script {
		_, err = fmt.Fprintln(hookFile, line)
		assert.NoError(t, err)
	}

	hookFile.Close()

	wrapper, err := NewScriptWrapper(WithHookPath(hookFile.Name()))
	assert.NoError(t, err)

	return wrapper
}

func assertScriptLike(t *testing.T, scriptTemplate, hookFileName string, wrapper *ScriptWrapper) {
	file, err := os.Open(wrapper.scriptFile.Name())
	assert.NoError(t, err)

	defer file.Close()

	wrapperScriptContents, err := ioutil.ReadAll(file)
	assert.NoError(t, err)

	expected := fmt.Sprintf(scriptTemplate, wrapper.beforeEnvFile.Name(), hookFileName, wrapper.afterEnvFile.Name())

	assert.Equal(t, expected, string(wrapperScriptContents))
}
<|endoftext|>


<commit_before>package process

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/buildkite/agent/v3/logger"
	"github.com/pkg/errors"
)

const (
	termType = `xterm-256color`
)

type Signal int

const (
	SIGHUP  Signal = 1
	SIGINT  Signal = 2
	SIGQUIT Signal = 3
	SIGUSR1 Signal = 10
	SIGUSR2 Signal = 12
	SIGTERM Signal = 15
)

var signalMap = map[string]Signal{
	`SIGHUP`:  SIGHUP,
	`SIGINT`:  SIGINT,
	`SIGQUIT`: SIGQUIT,
	`SIGUSR1`: SIGUSR1,
	`SIGUSR2`: SIGUSR2,
	`SIGTERM`: SIGTERM,
}

func (s Signal) String() string {
	for k, sig := range signalMap {
		if sig == s {
			return k
		}
	}
	return strconv.FormatInt(int64(s), 10)
}

func ParseSignal(sig string) (Signal, error) {
	s, ok := signalMap[strings.ToUpper(sig)]
	if !ok {
		return Signal(0), fmt.Errorf("Unknown signal %q", sig)
	}
	return s, nil
}

// Configuration for a Process
type Config struct {
	PTY             bool
	Timestamp       bool
	Path            string
	Args            []string
	Env             []string
	Stdout          io.Writer
	Stderr          io.Writer
	Dir             string
	Context         context.Context
	InterruptSignal Signal
}

// Process is an operating system level process
type Process struct {
	waitResult    error
	status        syscall.WaitStatus
	pid           int
	conf          Config
	logger        logger.Logger
	command       *exec.Cmd
	mu            sync.Mutex
	started, done chan struct{}
}

// New returns a new instance of Process
func New(l logger.Logger, c Config) *Process {
	return &Process{
		logger: l,
		conf:   c,
	}
}

// Pid is the pid of the running process
func (p *Process) Pid() int {
	return p.pid
}

// WaitResult returns the raw error returned by Wait()
func (p *Process) WaitResult() error {
	return p.waitResult
}

// WaitStatus returns the status of the Wait() call
func (p *Process) WaitStatus() syscall.WaitStatus {
	return p.status
}

// Run the command and block until it finishes
func (p *Process) Run() error {
	if p.command != nil {
		return fmt.Errorf("Process is already running")
	}

	// Create a command
	p.command = exec.Command(p.conf.Path, p.conf.Args...)

	// Setup the process to create a process group if supported
	// See https://github.com/kr/pty/issues/35 for context
	if !p.conf.PTY {
		p.setupProcessGroup()
	}

	// Configure working dir and fail if it doesn't exist, otherwise
	// we get confusing errors about fork/exec failing because the file
	// doesn't exist
	if p.conf.Dir != "" {
		if _, err := os.Stat(p.conf.Dir); os.IsNotExist(err) {
			return fmt.Errorf("Process working directory %q doesn't exist", p.conf.Dir)
		}
		p.command.Dir = p.conf.Dir
	}

	// Create channels for signalling started and done
	p.mu.Lock()
	if p.done == nil {
		p.done = make(chan struct{})
	}
	if p.started == nil {
		p.started = make(chan struct{})
	}
	p.mu.Unlock()

	// Copy the current processes ENV and merge in the new ones. We do this
	// so the sub process gets PATH and stuff. We merge our path in over
	// the top of the current one so the ENV from Buildkite and the agent
	// take precedence over the agent
	currentEnv := os.Environ()
	p.command.Env = append(currentEnv, p.conf.Env...)

	var waitGroup sync.WaitGroup

	// Toggle between running in a pty
	if p.conf.PTY {
		// Commands like tput expect a TERM value for a PTY
		p.command.Env = append(p.command.Env, `TERM=`+termType)

		pty, err := StartPTY(p.command)
		if err != nil {
			return err
		}

		// Make sure to close the pty at the end.
		defer func() { _ = pty.Close() }()

		p.pid = p.command.Process.Pid

		// Signal waiting consumers in Started() by closing the started channel
		close(p.started)

		waitGroup.Add(1)

		go func() {
			p.logger.Debug("[Process] Starting to copy PTY to the buffer")

			// Copy the pty to our writer. This will block until it
			// EOF's or something breaks.
			_, err = io.Copy(p.conf.Stdout, pty)
			if e, ok := err.(*os.PathError); ok && e.Err == syscall.EIO {
				// We can safely ignore this error, because
				// it's just the PTY telling us that it closed
				// successfully.  See:
				// https://github.com/buildkite/agent/v3/pull/34#issuecomment-46080419
				err = nil
			}

			if err != nil {
				p.logger.Error("[Process] PTY output copy failed with error: %T: %v", err, err)
			} else {
				p.logger.Debug("[Process] PTY has finished being copied to the buffer")
			}

			waitGroup.Done()
		}()
	} else {
		p.command.Stdout = p.conf.Stdout
		p.command.Stderr = p.conf.Stderr
		p.command.Stdin = nil

		err := p.command.Start()
		if err != nil {
			return err
		}

		p.pid = p.command.Process.Pid

		// Signal waiting consumers in Started() by closing the started channel
		close(p.started)
	}

	// When the context finishes, terminate the process
	if p.conf.Context != nil {
		go func() {
			select {
			case <-p.conf.Context.Done():
				p.logger.Debug("[Process] Context done, terminating")
				if err := p.Terminate(); err != nil {
					p.logger.Debug("[Process] Failed terminate: %v", err)
				}
				return

			case <-p.Done():
				return
			}
		}()
	}

	p.logger.Info("[Process] Process is running with PID: %d", p.pid)

	// Wait until the process has finished. The returned error is nil if the
	// command runs, has no problems copying stdin, stdout, and stderr, and
	// exits with a zero exit status.
	p.waitResult = p.command.Wait()

	// Signal waiting consumers in Done() by closing the done channel
	close(p.done)

	// Convert the wait result into a native WaitStatus
	if p.waitResult != nil {
		if err, ok := p.waitResult.(*exec.ExitError); ok {
			if s, ok := err.Sys().(syscall.WaitStatus); ok {
				p.status = s
			} else {
				return fmt.Errorf("Unimplemented for system where exec.ExitError.Sys() is not syscall.WaitStatus.")
			}
		} else {
			return fmt.Errorf("Unexpected error type %T", p.waitResult)
		}
	}

	// Find the exit status of the script
	p.logger.Info("Process with PID: %d finished with Exit Status: %d",
		p.pid, p.status.ExitStatus())

	// Sometimes (in docker containers) io.Copy never seems to finish. This is a mega
	// hack around it. If it doesn't finish after 1 second, just continue.
	p.logger.Debug("[Process] Waiting for routines to finish")
	err := timeoutWait(&waitGroup)
	if err != nil {
		p.logger.Debug("[Process] Timed out waiting for wait group: (%T: %v)", err, err)
	}

	return nil
}

// Done returns a channel that is closed when the process finishes
func (p *Process) Done() <-chan struct{} {
	p.mu.Lock()
	// We create this here in case this is called before Start()
	if p.done == nil {
		p.done = make(chan struct{})
	}
	d := p.done
	p.mu.Unlock()
	return d
}

// Started returns a channel that is closed when the process is started
func (p *Process) Started() <-chan struct{} {
	p.mu.Lock()
	// We create this here in case this is called before Start()
	if p.started == nil {
		p.started = make(chan struct{})
	}
	d := p.started
	p.mu.Unlock()
	return d
}

// Interrupt the process on platforms that support it, terminate otherwise
func (p *Process) Interrupt() error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.command == nil || p.command.Process == nil {
		p.logger.Debug("[Process] No process to interrupt yet")
		return nil
	}

	// interrupt the process (ctrl-c or SIGINT)
	if err := p.interruptProcessGroup(); err != nil {
		p.logger.Error("[Process] Failed to interrupt process %d: %v", p.pid, err)

		// Fallback to terminating if we get an error
		if termErr := p.terminateProcessGroup(); termErr != nil {
			return termErr
		}
	}

	return nil
}

// Terminate the process
func (p *Process) Terminate() error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.command == nil || p.command.Process == nil {
		p.logger.Debug("[Process] No process to terminate yet")
		return nil
	}

	return p.terminateProcessGroup()
}

func timeoutWait(waitGroup *sync.WaitGroup) error {
	// Make a chanel that we'll use as a timeout
	c := make(chan int, 1)

	// Start waiting for the routines to finish
	go func() {
		waitGroup.Wait()
		c <- 1
	}()

	select {
	case _ = <-c:
		return nil
	case <-time.After(10 * time.Second):
		return errors.New("Timeout")
	}
}
<commit_msg>Report child process exit signals in log output<commit_after>package process

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/buildkite/agent/v3/logger"
	"github.com/pkg/errors"
)

const (
	termType = `xterm-256color`
)

type Signal int

const (
	SIGHUP  Signal = 1
	SIGINT  Signal = 2
	SIGQUIT Signal = 3
	SIGUSR1 Signal = 10
	SIGUSR2 Signal = 12
	SIGTERM Signal = 15
)

var signalMap = map[string]Signal{
	`SIGHUP`:  SIGHUP,
	`SIGINT`:  SIGINT,
	`SIGQUIT`: SIGQUIT,
	`SIGUSR1`: SIGUSR1,
	`SIGUSR2`: SIGUSR2,
	`SIGTERM`: SIGTERM,
}

func (s Signal) String() string {
	for k, sig := range signalMap {
		if sig == s {
			return k
		}
	}
	return strconv.FormatInt(int64(s), 10)
}

func ParseSignal(sig string) (Signal, error) {
	s, ok := signalMap[strings.ToUpper(sig)]
	if !ok {
		return Signal(0), fmt.Errorf("Unknown signal %q", sig)
	}
	return s, nil
}

// Configuration for a Process
type Config struct {
	PTY             bool
	Timestamp       bool
	Path            string
	Args            []string
	Env             []string
	Stdout          io.Writer
	Stderr          io.Writer
	Dir             string
	Context         context.Context
	InterruptSignal Signal
}

// Process is an operating system level process
type Process struct {
	waitResult    error
	status        syscall.WaitStatus
	pid           int
	conf          Config
	logger        logger.Logger
	command       *exec.Cmd
	mu            sync.Mutex
	started, done chan struct{}
}

// New returns a new instance of Process
func New(l logger.Logger, c Config) *Process {
	return &Process{
		logger: l,
		conf:   c,
	}
}

// Pid is the pid of the running process
func (p *Process) Pid() int {
	return p.pid
}

// WaitResult returns the raw error returned by Wait()
func (p *Process) WaitResult() error {
	return p.waitResult
}

// WaitStatus returns the status of the Wait() call
func (p *Process) WaitStatus() syscall.WaitStatus {
	return p.status
}

// Run the command and block until it finishes
func (p *Process) Run() error {
	if p.command != nil {
		return fmt.Errorf("Process is already running")
	}

	// Create a command
	p.command = exec.Command(p.conf.Path, p.conf.Args...)

	// Setup the process to create a process group if supported
	// See https://github.com/kr/pty/issues/35 for context
	if !p.conf.PTY {
		p.setupProcessGroup()
	}

	// Configure working dir and fail if it doesn't exist, otherwise
	// we get confusing errors about fork/exec failing because the file
	// doesn't exist
	if p.conf.Dir != "" {
		if _, err := os.Stat(p.conf.Dir); os.IsNotExist(err) {
			return fmt.Errorf("Process working directory %q doesn't exist", p.conf.Dir)
		}
		p.command.Dir = p.conf.Dir
	}

	// Create channels for signalling started and done
	p.mu.Lock()
	if p.done == nil {
		p.done = make(chan struct{})
	}
	if p.started == nil {
		p.started = make(chan struct{})
	}
	p.mu.Unlock()

	// Copy the current processes ENV and merge in the new ones. We do this
	// so the sub process gets PATH and stuff. We merge our path in over
	// the top of the current one so the ENV from Buildkite and the agent
	// take precedence over the agent
	currentEnv := os.Environ()
	p.command.Env = append(currentEnv, p.conf.Env...)

	var waitGroup sync.WaitGroup

	// Toggle between running in a pty
	if p.conf.PTY {
		// Commands like tput expect a TERM value for a PTY
		p.command.Env = append(p.command.Env, `TERM=`+termType)

		pty, err := StartPTY(p.command)
		if err != nil {
			return err
		}

		// Make sure to close the pty at the end.
		defer func() { _ = pty.Close() }()

		p.pid = p.command.Process.Pid

		// Signal waiting consumers in Started() by closing the started channel
		close(p.started)

		waitGroup.Add(1)

		go func() {
			p.logger.Debug("[Process] Starting to copy PTY to the buffer")

			// Copy the pty to our writer. This will block until it
			// EOF's or something breaks.
			_, err = io.Copy(p.conf.Stdout, pty)
			if e, ok := err.(*os.PathError); ok && e.Err == syscall.EIO {
				// We can safely ignore this error, because
				// it's just the PTY telling us that it closed
				// successfully.  See:
				// https://github.com/buildkite/agent/v3/pull/34#issuecomment-46080419
				err = nil
			}

			if err != nil {
				p.logger.Error("[Process] PTY output copy failed with error: %T: %v", err, err)
			} else {
				p.logger.Debug("[Process] PTY has finished being copied to the buffer")
			}

			waitGroup.Done()
		}()
	} else {
		p.command.Stdout = p.conf.Stdout
		p.command.Stderr = p.conf.Stderr
		p.command.Stdin = nil

		err := p.command.Start()
		if err != nil {
			return err
		}

		p.pid = p.command.Process.Pid

		// Signal waiting consumers in Started() by closing the started channel
		close(p.started)
	}

	// When the context finishes, terminate the process
	if p.conf.Context != nil {
		go func() {
			select {
			case <-p.conf.Context.Done():
				p.logger.Debug("[Process] Context done, terminating")
				if err := p.Terminate(); err != nil {
					p.logger.Debug("[Process] Failed terminate: %v", err)
				}
				return

			case <-p.Done():
				return
			}
		}()
	}

	p.logger.Info("[Process] Process is running with PID: %d", p.pid)

	// Wait until the process has finished. The returned error is nil if the
	// command runs, has no problems copying stdin, stdout, and stderr, and
	// exits with a zero exit status.
	p.waitResult = p.command.Wait()

	// Signal waiting consumers in Done() by closing the done channel
	close(p.done)

	// Convert the wait result into a native WaitStatus
	if p.waitResult != nil {
		if err, ok := p.waitResult.(*exec.ExitError); ok {
			if s, ok := err.Sys().(syscall.WaitStatus); ok {
				p.status = s
			} else {
				return fmt.Errorf("Unimplemented for system where exec.ExitError.Sys() is not syscall.WaitStatus.")
			}
		} else {
			return fmt.Errorf("Unexpected error type %T", p.waitResult)
		}
	}

	exitSignal := "nil"
	if p.status.Signaled() {
		exitSignal = SignalString(p.status.Signal())
	}

	// Find the exit status of the script
	p.logger.Info("Process with PID: %d finished with Exit Status: %d, Signal: %s",
		p.pid, p.status.ExitStatus(), exitSignal)

	// Sometimes (in docker containers) io.Copy never seems to finish. This is a mega
	// hack around it. If it doesn't finish after 1 second, just continue.
	p.logger.Debug("[Process] Waiting for routines to finish")
	err := timeoutWait(&waitGroup)
	if err != nil {
		p.logger.Debug("[Process] Timed out waiting for wait group: (%T: %v)", err, err)
	}

	return nil
}

// Done returns a channel that is closed when the process finishes
func (p *Process) Done() <-chan struct{} {
	p.mu.Lock()
	// We create this here in case this is called before Start()
	if p.done == nil {
		p.done = make(chan struct{})
	}
	d := p.done
	p.mu.Unlock()
	return d
}

// Started returns a channel that is closed when the process is started
func (p *Process) Started() <-chan struct{} {
	p.mu.Lock()
	// We create this here in case this is called before Start()
	if p.started == nil {
		p.started = make(chan struct{})
	}
	d := p.started
	p.mu.Unlock()
	return d
}

// Interrupt the process on platforms that support it, terminate otherwise
func (p *Process) Interrupt() error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.command == nil || p.command.Process == nil {
		p.logger.Debug("[Process] No process to interrupt yet")
		return nil
	}

	// interrupt the process (ctrl-c or SIGINT)
	if err := p.interruptProcessGroup(); err != nil {
		p.logger.Error("[Process] Failed to interrupt process %d: %v", p.pid, err)

		// Fallback to terminating if we get an error
		if termErr := p.terminateProcessGroup(); termErr != nil {
			return termErr
		}
	}

	return nil
}

// Terminate the process
func (p *Process) Terminate() error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.command == nil || p.command.Process == nil {
		p.logger.Debug("[Process] No process to terminate yet")
		return nil
	}

	return p.terminateProcessGroup()
}

func timeoutWait(waitGroup *sync.WaitGroup) error {
	// Make a chanel that we'll use as a timeout
	c := make(chan int, 1)

	// Start waiting for the routines to finish
	go func() {
		waitGroup.Wait()
		c <- 1
	}()

	select {
	case _ = <-c:
		return nil
	case <-time.After(10 * time.Second):
		return errors.New("Timeout")
	}
}
<|endoftext|>


<commit_before>package ps

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	sh "github.com/codeskyblue/go-sh"
	"github.com/dokku/dokku/plugins/common"
	"github.com/dokku/dokku/plugins/config"
	dockeroptions "github.com/dokku/dokku/plugins/docker-options"
)

// TriggerAppRestart restarts an app
func TriggerAppRestart(appName string) error {
	return Restart(appName)
}

// TriggerCorePostDeploy removes extracted procfiles
// and sets a property to allow the app to be restored on boot
func TriggerCorePostDeploy(appName string) error {
	if err := removeProcfile(appName); err != nil {
		return err
	}

	entries := map[string]string{
		"DOKKU_APP_RESTORE": "1",
	}

	return common.SuppressOutput(func() error {
		return config.SetMany(appName, entries, false)
	})
}

// TriggerInstall initializes app restart policies
func TriggerInstall() error {
	directory := filepath.Join(common.MustGetEnv("DOKKU_LIB_ROOT"), "data", "ps")
	if err := os.MkdirAll(directory, 0755); err != nil {
		return err
	}

	if err := common.SetPermissions(directory, 0755); err != nil {
		return err
	}

	apps, err := common.DokkuApps()
	if err != nil {
		return nil
	}

	for _, appName := range apps {
		policies, err := getRestartPolicy(appName)
		if err != nil {
			return err
		}

		if len(policies) != 0 {
			continue
		}

		if err := dockeroptions.AddDockerOptionToPhases(appName, []string{"deploy"}, "--restart=on-failure:10"); err != nil {
			common.LogWarn(err.Error())
		}
	}

	return nil
}

// TriggerPostAppClone rebuilds the new app
func TriggerPostAppClone(oldAppName string, newAppName string) error {
	if os.Getenv("SKIP_REBUILD") == "true" {
		return nil
	}

	return Rebuild(newAppName)
}

// TriggerPostAppRename rebuilds the renamed app
func TriggerPostAppRename(oldAppName string, newAppName string) error {
	if os.Getenv("SKIP_REBUILD") == "true" {
		return nil
	}

	return Rebuild(newAppName)
}

// TriggerPostCreate ensures apps have a default restart policy
// and scale value for web
func TriggerPostCreate(appName string) error {
	if err := dockeroptions.AddDockerOptionToPhases(appName, []string{"deploy"}, "--restart=on-failure:10"); err != nil {
		return err
	}

	directory := filepath.Join(common.MustGetEnv("DOKKU_LIB_ROOT"), "data", "ps", appName)
	if err := os.MkdirAll(directory, 0755); err != nil {
		return err
	}

	if err := common.SetPermissions(directory, 0755); err != nil {
		return err
	}

	return updateScalefile(appName, make(map[string]int))
}

// TriggerPostDelete destroys the ps properties for a given app container
func TriggerPostDelete(appName string) error {
	return common.PropertyDestroy("ps", appName)
}

// TriggerPostExtract validates a procfile
func TriggerPostExtract(appName string, tempWorkDir string) error {
	procfile := filepath.Join(tempWorkDir, "Procfile")
	if !common.FileExists(procfile) {
		return nil
	}

	b, err := sh.Command("procfile-util", "check", "-P", procfile).CombinedOutput()
	if err != nil {
		return fmt.Errorf(strings.TrimSpace(string(b[:])))
	}
	return nil
}

// TriggerPostStop sets the restore property to false
func TriggerPostStop(appName string) error {
	entries := map[string]string{
		"DOKKU_APP_RESTORE": "0",
	}

	return common.SuppressOutput(func() error {
		return config.SetMany(appName, entries, false)
	})
}

// TriggerPreDeploy ensures an app has an up to date scale file
func TriggerPreDeploy(appName string, imageTag string) error {
	image := common.GetAppImageRepo(appName)
	removeProcfile(appName)

	procfilePath := getProcfilePath(appName)
	if err := extractProcfile(appName, image, procfilePath); err != nil {
		return err
	}

	if err := extractOrGenerateScalefile(appName, image); err != nil {
		return err
	}

	return nil
}

// TriggerProcfileExtract extracted the procfile
func TriggerProcfileExtract(appName string, image string) error {
	directory := filepath.Join(common.MustGetEnv("DOKKU_LIB_ROOT"), "data", "ps", appName)
	if err := os.MkdirAll(directory, 0755); err != nil {
		return err
	}

	if err := common.SetPermissions(directory, 0755); err != nil {
		return err
	}

	procfilePath := getProcfilePath(appName)

	if common.FileExists(procfilePath) {
		if err := common.PlugnTrigger("procfile-remove", []string{appName, procfilePath}...); err != nil {
			return err
		}
	}

	return extractProcfile(appName, image, procfilePath)
}

// TriggerProcfileGetCommand fetches a command from the procfile
func TriggerProcfileGetCommand(appName string, processType string, port int) error {
	procfilePath := getProcfilePath(appName)
	if !common.FileExists(procfilePath) {
		image := common.GetDeployingAppImageName(appName, "", "")
		err := common.SuppressOutput(func() error {
			return common.PlugnTrigger("procfile-extract", []string{appName, image}...)
		})

		if err != nil {
			return err
		}
	}
	command, err := getProcfileCommand(procfilePath, processType, port)
	if err != nil {
		return err
	}

	if command != "" {
		fmt.Printf("%s\n", command)
	}

	return nil
}

// TriggerProcfileRemove removes the procfile if it exists
func TriggerProcfileRemove(appName string, procfilePath string) error {
	if procfilePath == "" {
		procfilePath = getProcfilePath(appName)
	}

	if !common.FileExists(procfilePath) {
		return nil
	}

	os.Remove(procfilePath)
	return nil
}
<commit_msg>fix: correctly handle property and data perms<commit_after>package ps

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	sh "github.com/codeskyblue/go-sh"
	"github.com/dokku/dokku/plugins/common"
	"github.com/dokku/dokku/plugins/config"
	dockeroptions "github.com/dokku/dokku/plugins/docker-options"
)

// TriggerAppRestart restarts an app
func TriggerAppRestart(appName string) error {
	return Restart(appName)
}

// TriggerCorePostDeploy removes extracted procfiles
// and sets a property to allow the app to be restored on boot
func TriggerCorePostDeploy(appName string) error {
	if err := removeProcfile(appName); err != nil {
		return err
	}

	entries := map[string]string{
		"DOKKU_APP_RESTORE": "1",
	}

	return common.SuppressOutput(func() error {
		return config.SetMany(appName, entries, false)
	})
}

// TriggerInstall initializes app restart policies
func TriggerInstall() error {
	if err := common.PropertySetup("buildpacks"); err != nil {
		return fmt.Errorf("Unable to install the buildpacks plugin: %s", err.Error())
	}

	directory := filepath.Join(common.MustGetEnv("DOKKU_LIB_ROOT"), "data", "ps")
	if err := os.MkdirAll(directory, 0755); err != nil {
		return err
	}

	if err := common.SetPermissions(directory, 0755); err != nil {
		return err
	}

	apps, err := common.DokkuApps()
	if err != nil {
		return nil
	}

	for _, appName := range apps {
		policies, err := getRestartPolicy(appName)
		if err != nil {
			return err
		}

		if len(policies) != 0 {
			continue
		}

		if err := dockeroptions.AddDockerOptionToPhases(appName, []string{"deploy"}, "--restart=on-failure:10"); err != nil {
			common.LogWarn(err.Error())
		}
	}

	return nil
}

// TriggerPostAppClone rebuilds the new app
func TriggerPostAppClone(oldAppName string, newAppName string) error {
	if os.Getenv("SKIP_REBUILD") == "true" {
		return nil
	}

	return Rebuild(newAppName)
}

// TriggerPostAppRename rebuilds the renamed app
func TriggerPostAppRename(oldAppName string, newAppName string) error {
	if os.Getenv("SKIP_REBUILD") == "true" {
		return nil
	}

	return Rebuild(newAppName)
}

// TriggerPostCreate ensures apps have a default restart policy
// and scale value for web
func TriggerPostCreate(appName string) error {
	if err := dockeroptions.AddDockerOptionToPhases(appName, []string{"deploy"}, "--restart=on-failure:10"); err != nil {
		return err
	}

	directory := filepath.Join(common.MustGetEnv("DOKKU_LIB_ROOT"), "data", "ps", appName)
	if err := os.MkdirAll(directory, 0755); err != nil {
		return err
	}

	if err := common.SetPermissions(directory, 0755); err != nil {
		return err
	}

	return updateScalefile(appName, make(map[string]int))
}

// TriggerPostDelete destroys the ps properties for a given app container
func TriggerPostDelete(appName string) error {
	directory := filepath.Join(common.MustGetEnv("DOKKU_LIB_ROOT"), "data", "ps", appName)
	dataErr := os.RemoveAll(directory)
	propertyErr := common.PropertyDestroy("ps", appName)

	if dataErr != nil {
		return dataErr
	}

	return propertyErr
}

// TriggerPostExtract validates a procfile
func TriggerPostExtract(appName string, tempWorkDir string) error {
	procfile := filepath.Join(tempWorkDir, "Procfile")
	if !common.FileExists(procfile) {
		return nil
	}

	b, err := sh.Command("procfile-util", "check", "-P", procfile).CombinedOutput()
	if err != nil {
		return fmt.Errorf(strings.TrimSpace(string(b[:])))
	}
	return nil
}

// TriggerPostStop sets the restore property to false
func TriggerPostStop(appName string) error {
	entries := map[string]string{
		"DOKKU_APP_RESTORE": "0",
	}

	return common.SuppressOutput(func() error {
		return config.SetMany(appName, entries, false)
	})
}

// TriggerPreDeploy ensures an app has an up to date scale file
func TriggerPreDeploy(appName string, imageTag string) error {
	image := common.GetAppImageRepo(appName)
	removeProcfile(appName)

	procfilePath := getProcfilePath(appName)
	if err := extractProcfile(appName, image, procfilePath); err != nil {
		return err
	}

	if err := extractOrGenerateScalefile(appName, image); err != nil {
		return err
	}

	return nil
}

// TriggerProcfileExtract extracted the procfile
func TriggerProcfileExtract(appName string, image string) error {
	directory := filepath.Join(common.MustGetEnv("DOKKU_LIB_ROOT"), "data", "ps", appName)
	if err := os.MkdirAll(directory, 0755); err != nil {
		return err
	}

	if err := common.SetPermissions(directory, 0755); err != nil {
		return err
	}

	procfilePath := getProcfilePath(appName)

	if common.FileExists(procfilePath) {
		if err := common.PlugnTrigger("procfile-remove", []string{appName, procfilePath}...); err != nil {
			return err
		}
	}

	return extractProcfile(appName, image, procfilePath)
}

// TriggerProcfileGetCommand fetches a command from the procfile
func TriggerProcfileGetCommand(appName string, processType string, port int) error {
	procfilePath := getProcfilePath(appName)
	if !common.FileExists(procfilePath) {
		image := common.GetDeployingAppImageName(appName, "", "")
		err := common.SuppressOutput(func() error {
			return common.PlugnTrigger("procfile-extract", []string{appName, image}...)
		})

		if err != nil {
			return err
		}
	}
	command, err := getProcfileCommand(procfilePath, processType, port)
	if err != nil {
		return err
	}

	if command != "" {
		fmt.Printf("%s\n", command)
	}

	return nil
}

// TriggerProcfileRemove removes the procfile if it exists
func TriggerProcfileRemove(appName string, procfilePath string) error {
	if procfilePath == "" {
		procfilePath = getProcfilePath(appName)
	}

	if !common.FileExists(procfilePath) {
		return nil
	}

	os.Remove(procfilePath)
	return nil
}
<|endoftext|>


<commit_before>/*
*
 * Copyright 2014, Google Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
*/

package grpc

import (
	"math"
	"sync"

	"github.com/golang/protobuf/proto"
)

// Codec defines the interface gRPC uses to encode and decode messages.
// Note that implementations of this interface must be thread safe;
// a Codec's methods can be called from concurrent goroutines.
type Codec interface {
	// Marshal returns the wire format of v.
	Marshal(v interface{}) ([]byte, error)
	// Unmarshal parses the wire format into v.
	Unmarshal(data []byte, v interface{}) error
	// String returns the name of the Codec implementation. The returned
	// string will be used as part of content type in transmission.
	String() string
}

// protoCodec is a Codec implementation with protobuf. It is the default codec for gRPC.
type protoCodec struct {
}

type cachedProtoBuffer struct {
	lastMarshaledSize uint32
	proto.Buffer
}

func capToMaxInt32(val int) uint32 {
	if val > math.MaxInt32 {
		return uint32(math.MaxInt32)
	}
	return uint32(val)
}

func (p protoCodec) marshal(v interface{}, cb *cachedProtoBuffer) ([]byte, error) {
	protoMsg := v.(proto.Message)
	newSlice := make([]byte, 0, cb.lastMarshaledSize)

	cb.SetBuf(newSlice)
	cb.Reset()
	if err := cb.Marshal(protoMsg); err != nil {
		return nil, err
	}
	out := cb.Bytes()
	cb.lastMarshaledSize = capToMaxInt32(len(out))
	return out, nil
}

func (p protoCodec) Marshal(v interface{}) ([]byte, error) {
	cb := protoBufferPool.Get().(*cachedProtoBuffer)
	out, err := p.marshal(v, cb)

	// put back buffer and lose the ref to the slice
	cb.SetBuf(nil)
	protoBufferPool.Put(cb)
	return out, err
}

func (p protoCodec) Unmarshal(data []byte, v interface{}) error {
	cb := protoBufferPool.Get().(*cachedProtoBuffer)
	cb.SetBuf(data)
	err := cb.Unmarshal(v.(proto.Message))
	cb.SetBuf(nil)
	protoBufferPool.Put(cb)
	return err
}

func (protoCodec) String() string {
	return "proto"
}

var (
	protoBufferPool = &sync.Pool{
		New: func() interface{} {
			return &cachedProtoBuffer{
				Buffer:            proto.Buffer{},
				lastMarshaledSize: 16,
			}
		},
	}
)
<commit_msg>Reset proto before unmarshalling (#1222)<commit_after>/*
*
 * Copyright 2014, Google Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
*/

package grpc

import (
	"math"
	"sync"

	"github.com/golang/protobuf/proto"
)

// Codec defines the interface gRPC uses to encode and decode messages.
// Note that implementations of this interface must be thread safe;
// a Codec's methods can be called from concurrent goroutines.
type Codec interface {
	// Marshal returns the wire format of v.
	Marshal(v interface{}) ([]byte, error)
	// Unmarshal parses the wire format into v.
	Unmarshal(data []byte, v interface{}) error
	// String returns the name of the Codec implementation. The returned
	// string will be used as part of content type in transmission.
	String() string
}

// protoCodec is a Codec implementation with protobuf. It is the default codec for gRPC.
type protoCodec struct {
}

type cachedProtoBuffer struct {
	lastMarshaledSize uint32
	proto.Buffer
}

func capToMaxInt32(val int) uint32 {
	if val > math.MaxInt32 {
		return uint32(math.MaxInt32)
	}
	return uint32(val)
}

func (p protoCodec) marshal(v interface{}, cb *cachedProtoBuffer) ([]byte, error) {
	protoMsg := v.(proto.Message)
	newSlice := make([]byte, 0, cb.lastMarshaledSize)

	cb.SetBuf(newSlice)
	cb.Reset()
	if err := cb.Marshal(protoMsg); err != nil {
		return nil, err
	}
	out := cb.Bytes()
	cb.lastMarshaledSize = capToMaxInt32(len(out))
	return out, nil
}

func (p protoCodec) Marshal(v interface{}) ([]byte, error) {
	cb := protoBufferPool.Get().(*cachedProtoBuffer)
	out, err := p.marshal(v, cb)

	// put back buffer and lose the ref to the slice
	cb.SetBuf(nil)
	protoBufferPool.Put(cb)
	return out, err
}

func (p protoCodec) Unmarshal(data []byte, v interface{}) error {
	cb := protoBufferPool.Get().(*cachedProtoBuffer)
	cb.SetBuf(data)
	v.(proto.Message).Reset()
	err := cb.Unmarshal(v.(proto.Message))
	cb.SetBuf(nil)
	protoBufferPool.Put(cb)
	return err
}

func (protoCodec) String() string {
	return "proto"
}

var (
	protoBufferPool = &sync.Pool{
		New: func() interface{} {
			return &cachedProtoBuffer{
				Buffer:            proto.Buffer{},
				lastMarshaledSize: 16,
			}
		},
	}
)
<|endoftext|>


<commit_before>// Package stats add some statistics-related commands to the command loop.
//
// The new commands are in the form:
//
// stats {type} values...
//
package stats

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/gobs/args"
	"github.com/gobs/cmd"
	"github.com/gobs/cmd/internal"
	"github.com/montanaflynn/stats"
)

type statsPlugin struct {
	cmd.Plugin
}

var (
	Plugin = &statsPlugin{}
)

func floatString(v float64) string {
	s := strconv.FormatFloat(v, 'f', 3, 64)
	return strings.TrimSuffix(s, ".000")
}

//
// PluginInit initialize this plugin
//
func (p *statsPlugin) PluginInit(commander *cmd.Cmd, _ *internal.Context) error {

	commander.Add(cmd.Command{"stats",
		`
                stats {min|max|mean|median|sum|variance|pN} value...
                `,
		func(line string) (stop bool) {
			var res float64
			var err error

			parts := args.GetArgs(line) // [ type, value, ... ]
			if len(parts) == 0 {
				fmt.Println("usage: stats {min|max|mean|median|sum|variance|pN} value...")
				return
			}

			if len(parts) == 1 {
				res = 0.0
			} else {
				cmd := parts[0]
				data := stats.LoadRawData(parts[1:])

				switch cmd {
				case "min":
					res, err = data.Min()
				case "max":
					res, err = data.Max()
				case "mean":
					res, err = data.Mean()
				case "median":
					res, err = data.Median()
				case "sum":
					res, err = data.Sum()
				case "variance":
					res, err = data.Variance()
				default:
					fmt.Println("usage: stats {min|max|mean|median|sum|variance|pN} value...")
					return
				}
			}

			if err != nil {
				commander.SetVar("error", err, true)
				commander.SetVar("result", "0", true)
				fmt.Println(err)
			} else {
				sres := floatString(res)
				if !commander.Silent() {
					fmt.Println(sres)
				}

				commander.SetVar("error", "", true)
				commander.SetVar("result", sres, true)
			}

			return
		},
		nil})

	return nil
}
<commit_msg>Added standard deviation (std)<commit_after>// Package stats add some statistics-related commands to the command loop.
//
// The new commands are in the form:
//
// stats {type} values...
//
package stats

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/gobs/args"
	"github.com/gobs/cmd"
	"github.com/gobs/cmd/internal"
	"github.com/montanaflynn/stats"
)

type statsPlugin struct {
	cmd.Plugin
}

var (
	Plugin = &statsPlugin{}
)

func floatString(v float64) string {
	s := strconv.FormatFloat(v, 'f', 3, 64)
	return strings.TrimSuffix(s, ".000")
}

//
// PluginInit initialize this plugin
//
func (p *statsPlugin) PluginInit(commander *cmd.Cmd, _ *internal.Context) error {

	commander.Add(cmd.Command{"stats",
		`
                stats {min|max|mean|median|sum|variance|std|pN} value...
                `,
		func(line string) (stop bool) {
			var res float64
			var err error

			parts := args.GetArgs(line) // [ type, value, ... ]
			if len(parts) == 0 {
				fmt.Println("usage: stats {min|max|mean|median|sum|variance|std|pN} value...")
				return
			}

			if len(parts) == 1 {
				res = 0.0
			} else {
				cmd := parts[0]
				data := stats.LoadRawData(parts[1:])

				switch cmd {
				case "min":
					res, err = data.Min()
				case "max":
					res, err = data.Max()
				case "mean":
					res, err = data.Mean()
				case "median":
					res, err = data.Median()
				//case "mode":
				//	res, err = data.Mode()
				case "sum":
					res, err = data.Sum()
				case "variance":
					res, err = data.Variance()
				case "std":
					res, err = data.StandardDeviation()
				default:
					fmt.Println("usage: stats {min|max|mean|median|sum|variance|std|pN} value...")
					return
				}
			}

			if err != nil {
				commander.SetVar("error", err, true)
				commander.SetVar("result", "0", true)
				fmt.Println(err)
			} else {
				sres := floatString(res)
				if !commander.Silent() {
					fmt.Println(sres)
				}

				commander.SetVar("error", "", true)
				commander.SetVar("result", sres, true)
			}

			return
		},
		nil})

	return nil
}
<|endoftext|>