<commit_before>// Copyright 2017 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package api

import (
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/tsuru/tsuru/app"
	"github.com/tsuru/tsuru/auth"
	tsuruErrors "github.com/tsuru/tsuru/errors"
	"github.com/tsuru/tsuru/event"
	tsuruIo "github.com/tsuru/tsuru/io"
	"github.com/tsuru/tsuru/permission"
)

// title: app build
// path: /apps/{appname}/build
// method: POST
// consume: application/x-www-form-urlencoded
// responses:
//   200: OK
//   400: Invalid data
//   403: Forbidden
//   404: Not found
func build(w http.ResponseWriter, r *http.Request, t auth.Token) (err error) {
	tag := r.FormValue("tag")
	if tag == "" {
		return &tsuruErrors.HTTP{
			Code:    http.StatusBadRequest,
			Message: "you must specify the image tag.",
		}
	}
	opts, err := prepareToBuild(r)
	if err != nil {
		return err
	}
	if opts.File != nil {
		defer opts.File.Close()
	}
	w.Header().Set("Content-Type", "text")
	appName := r.URL.Query().Get(":appname")
	var userName string
	if t.IsAppToken() {
		if t.GetAppName() != appName && t.GetAppName() != app.InternalAppName {
			return &tsuruErrors.HTTP{Code: http.StatusUnauthorized, Message: "invalid app token"}
		}
		userName = r.FormValue("user")
	} else {
		userName = t.GetUserName()
	}
	instance, err := app.GetByName(appName)
	if err != nil {
		return &tsuruErrors.HTTP{Code: http.StatusNotFound, Message: err.Error()}
	}
	opts.App = instance
	opts.BuildTag = tag
	opts.User = userName
	opts.GetKind()
	if t.GetAppName() != app.InternalAppName {
		canBuild := permission.Check(t, permission.PermAppBuild, contextsForApp(instance)...)
		if !canBuild {
			return &tsuruErrors.HTTP{Code: http.StatusForbidden, Message: "User does not have permission to do this action in this app"}
		}
	}
	evt, err := event.New(&event.Opts{
		Target:        appTarget(appName),
		Kind:          permission.PermAppBuild,
		RawOwner:      event.Owner{Type: event.OwnerTypeUser, Name: userName},
		CustomData:    opts,
		Allowed:       event.Allowed(permission.PermAppReadEvents, contextsForApp(instance)...),
		AllowedCancel: event.Allowed(permission.PermAppUpdateEvents, contextsForApp(instance)...),
		Cancelable:    true,
	})
	if err != nil {
		return err
	}
	var imageID string
	defer func() { evt.DoneCustomData(err, map[string]string{"image": imageID}) }()
	opts.Event = evt
	writer := tsuruIo.NewKeepAliveWriter(w, 30*time.Second, "please wait...")
	defer writer.Stop()
	opts.OutputStream = writer
	imageID, err = app.Build(opts)
	if err == nil {
		fmt.Fprintln(w, imageID)
		fmt.Fprintln(w, "\nOK")
	}
	return err
}

func prepareToBuild(r *http.Request) (opts app.DeployOptions, err error) {
	var file multipart.File
	var fileSize int64
	if strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/") {
		file, _, err = r.FormFile("file")
		if err != nil {
			return opts, &tsuruErrors.HTTP{
				Code:    http.StatusBadRequest,
				Message: err.Error(),
			}
		}
		fileSize, err = file.Seek(0, io.SeekEnd)
		if err != nil {
			return opts, errors.Wrap(err, "unable to find uploaded file size")
		}
		file.Seek(0, io.SeekStart)
	}
	archiveURL := r.FormValue("archive-url")
	image := r.FormValue("image")
	if image == "" && archiveURL == "" && file == nil {
		return opts, &tsuruErrors.HTTP{
			Code:    http.StatusBadRequest,
			Message: "you must specify either the archive-url, a image url or upload a file.",
		}
	}
	var build bool
	buildString := r.FormValue("build")
	if buildString != "" {
		build, err = strconv.ParseBool(buildString)
		if err != nil {
			return opts, &tsuruErrors.HTTP{
				Code:    http.StatusBadRequest,
				Message: err.Error(),
			}
		}
	}
	opts.FileSize = fileSize
	opts.File = file
	opts.ArchiveURL = archiveURL
	opts.Image = image
	opts.Build = build
	return
}
<commit_msg>Update api/build.go<commit_after>// Copyright 2017 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package api

import (
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/tsuru/tsuru/app"
	"github.com/tsuru/tsuru/auth"
	tsuruErrors "github.com/tsuru/tsuru/errors"
	"github.com/tsuru/tsuru/event"
	tsuruIo "github.com/tsuru/tsuru/io"
	"github.com/tsuru/tsuru/permission"
)

// title: app build
// path: /apps/{appname}/build
// method: POST
// consume: application/x-www-form-urlencoded
// responses:
//   200: OK
//   400: Invalid data
//   403: Forbidden
//   404: Not found
func build(w http.ResponseWriter, r *http.Request, t auth.Token) (err error) {
	tag := r.FormValue("tag")
	if tag == "" {
		return &tsuruErrors.HTTP{
			Code:    http.StatusBadRequest,
			Message: "you must specify the image tag.",
		}
	}
	opts, err := prepareToBuild(r)
	if err != nil {
		return err
	}
	if opts.File != nil {
		defer opts.File.Close()
	}
	w.Header().Set("Content-Type", "text")
	appName := r.URL.Query().Get(":appname")
	var userName string
	if t.IsAppToken() {
		if t.GetAppName() != appName && t.GetAppName() != app.InternalAppName {
			return &tsuruErrors.HTTP{Code: http.StatusUnauthorized, Message: "invalid app token"}
		}
		userName = r.FormValue("user")
	} else {
		userName = t.GetUserName()
	}
	instance, err := app.GetByName(appName)
	if err != nil {
		return &tsuruErrors.HTTP{Code: http.StatusNotFound, Message: err.Error()}
	}
	opts.App = instance
	opts.BuildTag = tag
	opts.User = userName
	opts.GetKind()
	if t.GetAppName() != app.InternalAppName {
		canBuild := permission.Check(t, permission.PermAppBuild, contextsForApp(instance)...)
		if !canBuild {
			return &tsuruErrors.HTTP{Code: http.StatusForbidden, Message: "User does not have permission to do this action in this app"}
		}
	}
	evt, err := event.New(&event.Opts{
		Target:        appTarget(appName),
		Kind:          permission.PermAppBuild,
		RawOwner:      event.Owner{Type: event.OwnerTypeUser, Name: userName},
		CustomData:    opts,
		Allowed:       event.Allowed(permission.PermAppReadEvents, contextsForApp(instance)...),
		AllowedCancel: event.Allowed(permission.PermAppUpdateEvents, contextsForApp(instance)...),
		Cancelable:    true,
	})
	if err != nil {
		return err
	}
	var imageID string
	defer func() { evt.DoneCustomData(err, map[string]string{"image": imageID}) }()
	opts.Event = evt
	writer := tsuruIo.NewKeepAliveWriter(w, 30*time.Second, "please wait...")
	defer writer.Stop()
	opts.OutputStream = writer
	imageID, err = app.Build(opts)
	if err == nil {
		fmt.Fprintln(w, imageID)
		fmt.Fprintln(w, "OK")
	}
	return err
}

func prepareToBuild(r *http.Request) (opts app.DeployOptions, err error) {
	var file multipart.File
	var fileSize int64
	if strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/") {
		file, _, err = r.FormFile("file")
		if err != nil {
			return opts, &tsuruErrors.HTTP{
				Code:    http.StatusBadRequest,
				Message: err.Error(),
			}
		}
		fileSize, err = file.Seek(0, io.SeekEnd)
		if err != nil {
			return opts, errors.Wrap(err, "unable to find uploaded file size")
		}
		file.Seek(0, io.SeekStart)
	}
	archiveURL := r.FormValue("archive-url")
	image := r.FormValue("image")
	if image == "" && archiveURL == "" && file == nil {
		return opts, &tsuruErrors.HTTP{
			Code:    http.StatusBadRequest,
			Message: "you must specify either the archive-url, a image url or upload a file.",
		}
	}
	var build bool
	buildString := r.FormValue("build")
	if buildString != "" {
		build, err = strconv.ParseBool(buildString)
		if err != nil {
			return opts, &tsuruErrors.HTTP{
				Code:    http.StatusBadRequest,
				Message: err.Error(),
			}
		}
	}
	opts.FileSize = fileSize
	opts.File = file
	opts.ArchiveURL = archiveURL
	opts.Image = image
	opts.Build = build
	return
}
<|endoftext|>


<commit_before>package models

import "time"

// GetUnusedAccounts returns all user not seen for the duration passed in arguments
func GetUnusedAccounts(db XODB, unseenFor time.Duration) ([]*User, error) {
	var err error

	limitTime := time.Now().Add(-unseenFor)

	// sql query
	const sqlstr = `SELECT ` +
		`id, created, modified, email, auth, next_external, parent_user_id, aws_customer_identifier, aws_customer_entitlement, next_update_entitlement, anomalies_filters, last_seen, last_unused_reminder ` +
		`FROM trackit.user ` +
		`WHERE last_seen < ?`

	// run query
	XOLog(sqlstr, limitTime.Format(time.RFC3339))
	q, err := db.Query(sqlstr, limitTime.Format(time.RFC3339))
	if err != nil {
		return nil, err
	}
	defer q.Close()

	// load results
	res := []*User{}
	for q.Next() {
		u := User{
			_exists: true,
		}

		// scan
		err = q.Scan(&u.ID, &u.Created, &u.Modified, &u.Email, &u.Auth, &u.NextExternal, &u.ParentUserID, &u.AwsCustomerIdentifier, &u.AwsCustomerEntitlement, &u.NextUpdateEntitlement, &u.AnomaliesFilters, &u.LastSeen, &u.LastUnusedReminder)
		if err != nil {
			return nil, err
		}

		res = append(res, &u)
	}

	return res, nil
}
<commit_msg>Made same changes as previous commit on user.go<commit_after>package models

import "time"

// GetUnusedAccounts returns all user not seen for the duration passed in arguments
func GetUnusedAccounts(db XODB, unseenFor time.Duration) ([]*User, error) {
	var err error

	limitTime := time.Now().Add(-unseenFor)

	// sql query
	const sqlstr = `SELECT ` +
		`id, email, auth, next_external, parent_user_id, aws_customer_identifier, aws_customer_entitlement, next_update_entitlement, anomalies_filters, last_seen, last_unused_reminder ` +
		`FROM trackit.user ` +
		`WHERE last_seen < ?`

	// run query
	XOLog(sqlstr, limitTime.Format(time.RFC3339))
	q, err := db.Query(sqlstr, limitTime.Format(time.RFC3339))
	if err != nil {
		return nil, err
	}
	defer q.Close()

	// load results
	res := []*User{}
	for q.Next() {
		u := User{
			_exists: true,
		}

		// scan
		err = q.Scan(&u.ID, &u.Email, &u.Auth, &u.NextExternal, &u.ParentUserID, &u.AwsCustomerIdentifier, &u.AwsCustomerEntitlement, &u.NextUpdateEntitlement, &u.AnomaliesFilters, &u.LastSeen, &u.LastUnusedReminder)
		if err != nil {
			return nil, err
		}

		res = append(res, &u)
	}

	return res, nil
}
<|endoftext|>


<commit_before>package dispatch

import (
    "github.com/robmcl4/Mycroft-Core-Go/mycroft/cmd"
)

var commandChan chan *cmd.Command = make(chan *cmd.Command, 100)
var preemptChan chan *cmd.Command = make(chan *cmd.Command, 100)

// add a command to the dispatcher
func Enqueue(c *cmd.Command) {
    commandChan <- c
}

// skip to the front of the queue
func PreemptQueue(c *cmd.Command) {
    preemptChan <- c
}

// Listen for and execute commands
func Dispatch() {
    for {
        select {
        case toExec := <- preemptChan:
            toExec.Execute()
        default:
            toExec := <- commandChan
            toExec.Execute()
        }
    }
}
<commit_msg>remove dispatch module<commit_after><|endoftext|>


<commit_before>package vaulted

import (
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"strings"
	"syscall"
	"time"

	"github.com/aws/aws-sdk-go/aws/arn"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

var (
	SessionVersion = "2"
)

type Session struct {
	SessionVersion string `json:"version"`

	Name       string    `json:"name"`
	Expiration time.Time `json:"expiration"`

	ActiveRole string `json:"active_role,omitempty"`

	AWSCreds *AWSCredentials   `json:"aws_creds,omitempty"`
	Role     string            `json:"role,omitempty"`
	Vars     map[string]string `json:"vars,omitempty"`
	SSHKeys  map[string]string `json:"ssh_keys,omitempty"`
}

func (e *Session) AssumeSessionRole() (*Session, error) {
	if e.Role == "" {
		return e, nil
	}

	role := e.Role
	e.Role = ""

	return e.AssumeRole(role)
}

func (e *Session) AssumeRole(roleArn string) (*Session, error) {
	expiration := e.Expiration
	maxExpiration := time.Now().Add(time.Hour).Truncate(time.Second)
	if expiration.After(maxExpiration) {
		expiration = maxExpiration
	}

	duration := expiration.Sub(time.Now())

	var creds *AWSCredentials

	selectedRoleArn := roleArn
	parsedArn, err := arn.Parse(roleArn)
	if err == nil {
		creds, err = e.AWSCreds.AssumeRole(parsedArn.String(), duration)
		if err != nil {
			return nil, err
		}
	} else {
		// Unparseable ARN
		identityArn, err := e.AWSCreds.GetCallerIdentity()
		if err != nil {
			return nil, err
		}

		fullRoleArn := arn.ARN{
			Partition: identityArn.Partition,
			Service:   "iam",
			Region:    "",
			AccountID: identityArn.AccountID,
			Resource:  "role/" + roleArn,
		}

		selectedRoleArn = fullRoleArn.String()
		creds, err = e.AWSCreds.AssumeRole(selectedRoleArn, duration)
		if err != nil {
			return nil, fmt.Errorf("Error assuming role '%s' which was interpreted as '%s'\nError: %v", roleArn, selectedRoleArn, err)
		}
	}

	session := &Session{
		SessionVersion: SessionVersion,

		Name:       e.Name,
		Expiration: *creds.Expiration,

		ActiveRole: selectedRoleArn,

		AWSCreds: creds,
		Vars:     make(map[string]string),
		SSHKeys:  make(map[string]string),
	}
	for key, value := range e.Vars {
		session.Vars[key] = value
	}
	for key, value := range e.SSHKeys {
		session.SSHKeys[key] = value
	}

	return session, nil
}

func (e *Session) Spawn(cmd []string) (*int, error) {
	if len(cmd) == 0 {
		return nil, ErrInvalidCommand
	}

	// lookup the path of the executable
	cmdpath, err := exec.LookPath(cmd[0])
	if err != nil {
		return nil, fmt.Errorf("Cannot find executable %s: %v", cmd[0], err)
	}

	// start the agent
	sock, err := e.startProxyKeyring()
	if err != nil {
		return nil, err
	}

	vars := make(map[string]string)
	vars["SSH_AUTH_SOCK"] = sock

	// trap signals
	sigs := make(chan os.Signal)
	signal.Notify(
		sigs,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGQUIT,
		syscall.SIGTERM,
		syscall.SIGUSR1,
		syscall.SIGUSR2,
		syscall.SIGWINCH,
	)

	// start the process
	var attr os.ProcAttr
	attr.Env = e.buildEnviron(vars)
	attr.Files = []*os.File{os.Stdin, os.Stdout, os.Stderr}

	proc, err := os.StartProcess(cmdpath, cmd, &attr)
	if err != nil {
		return nil, fmt.Errorf("Failed to execute command: %v", err)
	}

	// relay trapped signals to the spawned process
	go func() {
		for s := range sigs {
			proc.Signal(s)
		}
	}()

	defer func() {
		signal.Stop(sigs)
		close(sigs)
	}()

	// wait for the process to exit
	state, _ := proc.Wait()

	var exitStatus int
	if !state.Success() {
		if status, ok := state.Sys().(syscall.WaitStatus); ok {
			exitStatus = status.ExitStatus()
		} else {
			exitStatus = 255
		}
	}

	// we only return an error if spawning the process failed, not if
	// the spawned command returned a failure status code.
	return &exitStatus, nil
}

func (e *Session) startProxyKeyring() (string, error) {
	keyring, err := NewProxyKeyring(os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		return "", err
	}

	// load ssh keys
	for comment, key := range e.SSHKeys {
		timeRemaining := e.Expiration.Sub(time.Now())
		addedKey := agent.AddedKey{
			Comment:      comment,
			LifetimeSecs: uint32(timeRemaining.Seconds()),
		}

		addedKey.PrivateKey, err = ssh.ParseRawPrivateKey([]byte(key))
		if err != nil {
			return "", err
		}

		err := keyring.Add(addedKey)
		if err != nil {
			return "", err
		}
	}

	sock, err := keyring.Listen()
	if err != nil {
		return "", err
	}

	go keyring.Serve()

	return sock, err
}

func (e *Session) Variables() *Variables {
	vars := Variables{
		Set: make(map[string]string),
	}

	for key, value := range e.Vars {
		vars.Set[key] = value
	}

	vars.Set["VAULTED_ENV"] = e.Name
	vars.Set["VAULTED_ENV_EXPIRATION"] = e.Expiration.UTC().Format(time.RFC3339)

	if e.ActiveRole != "" {
		vars.Set["VAULTED_ENV_ROLE_ARN"] = e.ActiveRole

		roleArn, err := arn.Parse(e.ActiveRole)
		if err == nil {
			resource := strings.TrimPrefix(roleArn.Resource, "role/")
			vars.Set["VAULTED_ENV_ROLE_PARTITION"] = roleArn.Partition
			vars.Set["VAULTED_ENV_ROLE_ACCOUNT_ID"] = roleArn.AccountID
			vars.Set["VAULTED_ENV_ROLE_NAME"] = path.Base(resource)
			vars.Set["VAULTED_ENV_ROLE_PATH"] = fmt.Sprintf("/%s/", path.Dir(resource))
		}
	}

	if e.AWSCreds != nil {
		vars.Set["AWS_ACCESS_KEY_ID"] = e.AWSCreds.ID
		vars.Set["AWS_SECRET_ACCESS_KEY"] = e.AWSCreds.Secret

		if e.AWSCreds.Token != "" {
			vars.Set["AWS_SESSION_TOKEN"] = e.AWSCreds.Token
			vars.Set["AWS_SECURITY_TOKEN"] = e.AWSCreds.Token
		} else {
			vars.Unset = append(
				vars.Unset,
				"AWS_SESSION_TOKEN",
				"AWS_SECURITY_TOKEN",
			)
		}
	}

	return &vars
}

func (e *Session) buildEnviron(extraVars map[string]string) []string {
	vars := make(map[string]string)
	for _, v := range os.Environ() {
		parts := strings.SplitN(v, "=", 2)
		vars[parts[0]] = parts[1]
	}

	v := e.Variables()
	for _, key := range v.Unset {
		delete(vars, key)
	}
	for key, value := range v.Set {
		vars[key] = value
	}

	for key, value := range extraVars {
		vars[key] = value
	}

	environ := make([]string, 0, len(vars))
	for key, value := range vars {
		environ = append(environ, fmt.Sprintf("%s=%s", key, value))
	}
	return environ
}

type Variables struct {
	Set   map[string]string
	Unset []string
}
<commit_msg>Rename session receiver variable<commit_after>package vaulted

import (
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"strings"
	"syscall"
	"time"

	"github.com/aws/aws-sdk-go/aws/arn"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

var (
	SessionVersion = "2"
)

type Session struct {
	SessionVersion string `json:"version"`

	Name       string    `json:"name"`
	Expiration time.Time `json:"expiration"`

	ActiveRole string `json:"active_role,omitempty"`

	AWSCreds *AWSCredentials   `json:"aws_creds,omitempty"`
	Role     string            `json:"role,omitempty"`
	Vars     map[string]string `json:"vars,omitempty"`
	SSHKeys  map[string]string `json:"ssh_keys,omitempty"`
}

func (s *Session) AssumeSessionRole() (*Session, error) {
	if s.Role == "" {
		return s, nil
	}

	role := s.Role
	s.Role = ""

	return s.AssumeRole(role)
}

func (s *Session) AssumeRole(roleArn string) (*Session, error) {
	expiration := s.Expiration
	maxExpiration := time.Now().Add(time.Hour).Truncate(time.Second)
	if expiration.After(maxExpiration) {
		expiration = maxExpiration
	}

	duration := expiration.Sub(time.Now())

	var creds *AWSCredentials

	selectedRoleArn := roleArn
	parsedArn, err := arn.Parse(roleArn)
	if err == nil {
		creds, err = s.AWSCreds.AssumeRole(parsedArn.String(), duration)
		if err != nil {
			return nil, err
		}
	} else {
		// Unparseable ARN
		identityArn, err := s.AWSCreds.GetCallerIdentity()
		if err != nil {
			return nil, err
		}

		fullRoleArn := arn.ARN{
			Partition: identityArn.Partition,
			Service:   "iam",
			Region:    "",
			AccountID: identityArn.AccountID,
			Resource:  "role/" + roleArn,
		}

		selectedRoleArn = fullRoleArn.String()
		creds, err = s.AWSCreds.AssumeRole(selectedRoleArn, duration)
		if err != nil {
			return nil, fmt.Errorf("Error assuming role '%s' which was interpreted as '%s'\nError: %v", roleArn, selectedRoleArn, err)
		}
	}

	session := &Session{
		SessionVersion: SessionVersion,

		Name:       s.Name,
		Expiration: *creds.Expiration,

		ActiveRole: selectedRoleArn,

		AWSCreds: creds,
		Vars:     make(map[string]string),
		SSHKeys:  make(map[string]string),
	}
	for key, value := range s.Vars {
		session.Vars[key] = value
	}
	for key, value := range s.SSHKeys {
		session.SSHKeys[key] = value
	}

	return session, nil
}

func (s *Session) Spawn(cmd []string) (*int, error) {
	if len(cmd) == 0 {
		return nil, ErrInvalidCommand
	}

	// lookup the path of the executable
	cmdpath, err := exec.LookPath(cmd[0])
	if err != nil {
		return nil, fmt.Errorf("Cannot find executable %s: %v", cmd[0], err)
	}

	// start the agent
	sock, err := s.startProxyKeyring()
	if err != nil {
		return nil, err
	}

	vars := make(map[string]string)
	vars["SSH_AUTH_SOCK"] = sock

	// trap signals
	sigs := make(chan os.Signal)
	signal.Notify(
		sigs,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGQUIT,
		syscall.SIGTERM,
		syscall.SIGUSR1,
		syscall.SIGUSR2,
		syscall.SIGWINCH,
	)

	// start the process
	var attr os.ProcAttr
	attr.Env = s.buildEnviron(vars)
	attr.Files = []*os.File{os.Stdin, os.Stdout, os.Stderr}

	proc, err := os.StartProcess(cmdpath, cmd, &attr)
	if err != nil {
		return nil, fmt.Errorf("Failed to execute command: %v", err)
	}

	// relay trapped signals to the spawned process
	go func() {
		for sig := range sigs {
			proc.Signal(sig)
		}
	}()

	defer func() {
		signal.Stop(sigs)
		close(sigs)
	}()

	// wait for the process to exit
	state, _ := proc.Wait()

	var exitStatus int
	if !state.Success() {
		if status, ok := state.Sys().(syscall.WaitStatus); ok {
			exitStatus = status.ExitStatus()
		} else {
			exitStatus = 255
		}
	}

	// we only return an error if spawning the process failed, not if
	// the spawned command returned a failure status code.
	return &exitStatus, nil
}

func (s *Session) startProxyKeyring() (string, error) {
	keyring, err := NewProxyKeyring(os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		return "", err
	}

	// load ssh keys
	for comment, key := range s.SSHKeys {
		timeRemaining := s.Expiration.Sub(time.Now())
		addedKey := agent.AddedKey{
			Comment:      comment,
			LifetimeSecs: uint32(timeRemaining.Seconds()),
		}

		addedKey.PrivateKey, err = ssh.ParseRawPrivateKey([]byte(key))
		if err != nil {
			return "", err
		}

		err := keyring.Add(addedKey)
		if err != nil {
			return "", err
		}
	}

	sock, err := keyring.Listen()
	if err != nil {
		return "", err
	}

	go keyring.Serve()

	return sock, err
}

func (s *Session) Variables() *Variables {
	vars := Variables{
		Set: make(map[string]string),
	}

	for key, value := range s.Vars {
		vars.Set[key] = value
	}

	vars.Set["VAULTED_ENV"] = s.Name
	vars.Set["VAULTED_ENV_EXPIRATION"] = s.Expiration.UTC().Format(time.RFC3339)

	if s.ActiveRole != "" {
		vars.Set["VAULTED_ENV_ROLE_ARN"] = s.ActiveRole

		roleArn, err := arn.Parse(s.ActiveRole)
		if err == nil {
			resource := strings.TrimPrefix(roleArn.Resource, "role/")
			vars.Set["VAULTED_ENV_ROLE_PARTITION"] = roleArn.Partition
			vars.Set["VAULTED_ENV_ROLE_ACCOUNT_ID"] = roleArn.AccountID
			vars.Set["VAULTED_ENV_ROLE_NAME"] = path.Base(resource)
			vars.Set["VAULTED_ENV_ROLE_PATH"] = fmt.Sprintf("/%s/", path.Dir(resource))
		}
	}

	if s.AWSCreds != nil {
		vars.Set["AWS_ACCESS_KEY_ID"] = s.AWSCreds.ID
		vars.Set["AWS_SECRET_ACCESS_KEY"] = s.AWSCreds.Secret

		if s.AWSCreds.Token != "" {
			vars.Set["AWS_SESSION_TOKEN"] = s.AWSCreds.Token
			vars.Set["AWS_SECURITY_TOKEN"] = s.AWSCreds.Token
		} else {
			vars.Unset = append(
				vars.Unset,
				"AWS_SESSION_TOKEN",
				"AWS_SECURITY_TOKEN",
			)
		}
	}

	return &vars
}

func (s *Session) buildEnviron(extraVars map[string]string) []string {
	vars := make(map[string]string)
	for _, v := range os.Environ() {
		parts := strings.SplitN(v, "=", 2)
		vars[parts[0]] = parts[1]
	}

	v := s.Variables()
	for _, key := range v.Unset {
		delete(vars, key)
	}
	for key, value := range v.Set {
		vars[key] = value
	}

	for key, value := range extraVars {
		vars[key] = value
	}

	environ := make([]string, 0, len(vars))
	for key, value := range vars {
		environ = append(environ, fmt.Sprintf("%s=%s", key, value))
	}
	return environ
}

type Variables struct {
	Set   map[string]string
	Unset []string
}
<|endoftext|>


<commit_before>// Package encoding converts arbitrary objects into byte slices, and vis
// versa. It also contains helper functions for reading and writing length-
// prefixed data. See doc/Encoding.md for the full encoding specification.
package encoding

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"reflect"
)

// A SiaMarshaler can encode itself as a byte slice.
type SiaMarshaler interface {
	MarshalSia() []byte
}

// A SiaUnmarshaler can decode itself from a byte slice. If a decoding error
// occurs, UnmarshalSia should panic.
type SiaUnmarshaler interface {
	UnmarshalSia([]byte)
}

// An Encoder writes objects to an output stream.
type Encoder struct {
	w io.Writer
}

const (
	maxDecodeLen = 10 * 1024 * 1024 // 10 MB
	maxSliceLen  = 4 * 1024 * 1024  // 4 MB
)

var (
	errBadPointer = errors.New("cannot decode into invalid pointer")
)

// Encode writes the encoding of v to the stream. For encoding details, see
// the package docstring.
func (e *Encoder) Encode(v interface{}) error {
	return e.encode(reflect.ValueOf(v))
}

// write catches instances where short writes do not return an error.
func (e *Encoder) write(p []byte) error {
	n, err := e.w.Write(p)
	if n != len(p) {
		return io.ErrShortWrite
	}
	return err
}

// Encode writes the encoding of val to the stream. For encoding details, see
// the package docstring.
func (e *Encoder) encode(val reflect.Value) error {
	// check for MarshalSia interface first
	if m, ok := val.Interface().(SiaMarshaler); ok {
		return WritePrefix(e.w, m.MarshalSia())
	}

	switch val.Kind() {
	case reflect.Ptr:
		// write either a 1 or 0
		if err := e.Encode(!val.IsNil()); err != nil {
			return err
		}
		if !val.IsNil() {
			return e.encode(val.Elem())
		}
	case reflect.Bool:
		if val.Bool() {
			return e.write([]byte{1})
		} else {
			return e.write([]byte{0})
		}
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return e.write(EncInt64(val.Int()))
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return e.write(EncUint64(val.Uint()))
	case reflect.String:
		return WritePrefix(e.w, []byte(val.String()))
	case reflect.Slice:
		// slices are variable length, so prepend the length and then fallthrough to array logic
		if err := e.write(EncUint64(uint64(val.Len()))); err != nil {
			return err
		}
		fallthrough
	case reflect.Array:
		// special case for byte arrays
		if val.Type().Elem().Kind() == reflect.Uint8 {
			// if the array is addressable, we can optimize a bit here
			if val.CanAddr() {
				return e.write(val.Slice(0, val.Len()).Bytes())
			}
			// otherwise we have to copy into a newly allocated slice
			slice := reflect.MakeSlice(reflect.SliceOf(val.Type().Elem()), val.Len(), val.Len())
			reflect.Copy(slice, val)
			return e.write(slice.Bytes())
		}
		// normal slices/arrays are encoded by sequentially encoding their elements
		for i := 0; i < val.Len(); i++ {
			if err := e.encode(val.Index(i)); err != nil {
				return err
			}
		}
	case reflect.Struct:
		for i := 0; i < val.NumField(); i++ {
			if err := e.encode(val.Field(i)); err != nil {
				return err
			}
		}
	}

	// Marshalling should never fail. If it panics, you're doing something wrong,
	// like trying to encode a map or an unexported struct field.
	panic("could not marshal type " + val.Type().String())
}

// NewEncoder returns a new encoder that writes to w.
func NewEncoder(w io.Writer) *Encoder {
	return &Encoder{w}
}

// Marshal returns the encoding of v. For encoding details, see the package
// docstring.
func Marshal(v interface{}) []byte {
	b := new(bytes.Buffer)
	NewEncoder(b).Encode(v) // no error possible when using a bytes.Buffer
	return b.Bytes()
}

// MarshalAll encodes all of its inputs and returns their concatenation.
func MarshalAll(v ...interface{}) []byte {
	b := new(bytes.Buffer)
	enc := NewEncoder(b)
	for i := range v {
		enc.Encode(v[i]) // no error possible when using a bytes.Buffer
	}
	return b.Bytes()
}

// WriteFile writes v to a file. The file will be created if it does not exist.
func WriteFile(filename string, v interface{}) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	return NewEncoder(file).Encode(v)
}

// A Decoder reads and decodes values from an input stream.
type Decoder struct {
	r io.Reader
	n int
}

// Read implements the io.Reader interface. It also keeps track of the total
// number of bytes decoded, and panics if that number exceeds a global
// maximum.
func (d *Decoder) Read(p []byte) (int, error) {
	n, err := d.r.Read(p)
	// enforce an absolute maximum size limit
	if d.n += n; d.n > maxDecodeLen {
		panic("encoded type exceeds size limit")
	}
	return n, err
}

// Decode reads the next encoded value from its input stream and stores it in
// v, which must be a pointer. The decoding rules are the inverse of those
// specified in the package docstring.
func (d *Decoder) Decode(v interface{}) (err error) {
	// v must be a pointer
	pval := reflect.ValueOf(v)
	if pval.Kind() != reflect.Ptr || pval.IsNil() {
		return errBadPointer
	}

	// catch decoding panics and convert them to errors
	// note that this allows us to skip boundary checks during decoding
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("could not decode type %s: %v", pval.Elem().Type().String(), r)
		}
	}()

	// reset the read count
	d.n = 0

	d.decode(pval.Elem())
	return
}

// readN reads n bytes and panics if the read fails.
func (d *Decoder) readN(n int) []byte {
	b := make([]byte, n)
	_, err := io.ReadFull(d, b)
	if err != nil {
		panic(err)
	}
	return b
}

// readPrefix reads a length-prefixed byte slice and panics if the read fails.
func (d *Decoder) readPrefix() []byte {
	b, err := ReadPrefix(d, maxSliceLen)
	if err != nil {
		panic(err)
	}
	return b
}

// decode reads the next encoded value from its input stream and stores it in
// val. The decoding rules are the inverse of those specified in the package
// docstring.
func (d *Decoder) decode(val reflect.Value) {
	// check for UnmarshalSia interface first
	if val.CanAddr() {
		if u, ok := val.Addr().Interface().(SiaUnmarshaler); ok {
			u.UnmarshalSia(d.readPrefix())
			return
		}
	}

	switch val.Kind() {
	case reflect.Ptr:
		var valid bool
		d.decode(reflect.ValueOf(&valid).Elem())
		// nil pointer, nothing to decode
		if !valid {
			return
		}
		// make sure we aren't decoding into nil
		if val.IsNil() {
			val.Set(reflect.New(val.Type().Elem()))
		}
		d.decode(val.Elem())
	case reflect.Bool:
		b := d.readN(1)
		if b[0] > 1 {
			panic("boolean value was not 0 or 1")
		}
		val.SetBool(b[0] == 1)
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		val.SetInt(DecInt64(d.readN(8)))
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		val.SetUint(DecUint64(d.readN(8)))
	case reflect.String:
		val.SetString(string(d.readPrefix()))
	case reflect.Slice:
		// slices are variable length, but otherwise the same as arrays.
		// just have to allocate them first, then we can fallthrough to the array logic.
		sliceLen := DecUint64(d.readN(8))
		// sanity-check the sliceLen, otherwise you can crash a peer by making
		// them allocate a massive slice
		if sliceLen > 1<<31-1 || sliceLen*uint64(val.Type().Elem().Size()) > maxSliceLen {
			panic("slice is too large")
		}
		val.Set(reflect.MakeSlice(val.Type(), int(sliceLen), int(sliceLen)))
		fallthrough
	case reflect.Array:
		// special case for byte arrays (e.g. hashes)
		if val.Type().Elem().Kind() == reflect.Uint8 {
			// convert val to a slice and read into it directly
			b := val.Slice(0, val.Len())
			_, err := io.ReadFull(d, b.Bytes())
			if err != nil {
				panic(err)
			}
			return
		}
		// arrays are unmarshalled by sequentially unmarshalling their elements
		for i := 0; i < val.Len(); i++ {
			d.decode(val.Index(i))
		}
		return
	case reflect.Struct:
		for i := 0; i < val.NumField(); i++ {
			d.decode(val.Field(i))
		}
		return
	default:
		panic("unknown type")
	}
}

// NewDecoder returns a new decoder that reads from r.
func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{r, 0}
}

// Unmarshal decodes the encoded value b and stores it in v, which must be a
// pointer. The decoding rules are the inverse of those specified in the
// package docstring.
func Unmarshal(b []byte, v interface{}) error {
	r := bytes.NewReader(b)
	return NewDecoder(r).Decode(v)
}

// ReadFile reads the contents of a file and decodes them into v.
func ReadFile(filename string, v interface{}) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	return NewDecoder(file).Decode(v)
}
<commit_msg>fix marshal panics<commit_after>// Package encoding converts arbitrary objects into byte slices, and vis
// versa. It also contains helper functions for reading and writing length-
// prefixed data. See doc/Encoding.md for the full encoding specification.
package encoding

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"reflect"
)

// A SiaMarshaler can encode itself as a byte slice.
type SiaMarshaler interface {
	MarshalSia() []byte
}

// A SiaUnmarshaler can decode itself from a byte slice. If a decoding error
// occurs, UnmarshalSia should panic.
type SiaUnmarshaler interface {
	UnmarshalSia([]byte)
}

// An Encoder writes objects to an output stream.
type Encoder struct {
	w io.Writer
}

const (
	maxDecodeLen = 10 * 1024 * 1024 // 10 MB
	maxSliceLen  = 4 * 1024 * 1024  // 4 MB
)

var (
	errBadPointer = errors.New("cannot decode into invalid pointer")
)

// Encode writes the encoding of v to the stream. For encoding details, see
// the package docstring.
func (e *Encoder) Encode(v interface{}) error {
	return e.encode(reflect.ValueOf(v))
}

// write catches instances where short writes do not return an error.
func (e *Encoder) write(p []byte) error {
	n, err := e.w.Write(p)
	if n != len(p) {
		return io.ErrShortWrite
	}
	return err
}

// Encode writes the encoding of val to the stream. For encoding details, see
// the package docstring.
func (e *Encoder) encode(val reflect.Value) error {
	// check for MarshalSia interface first
	if m, ok := val.Interface().(SiaMarshaler); ok {
		return WritePrefix(e.w, m.MarshalSia())
	}

	switch val.Kind() {
	case reflect.Ptr:
		// write either a 1 or 0
		if err := e.Encode(!val.IsNil()); err != nil {
			return err
		}
		if !val.IsNil() {
			return e.encode(val.Elem())
		}
	case reflect.Bool:
		if val.Bool() {
			return e.write([]byte{1})
		} else {
			return e.write([]byte{0})
		}
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return e.write(EncInt64(val.Int()))
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return e.write(EncUint64(val.Uint()))
	case reflect.String:
		return WritePrefix(e.w, []byte(val.String()))
	case reflect.Slice:
		// slices are variable length, so prepend the length and then fallthrough to array logic
		if err := e.write(EncUint64(uint64(val.Len()))); err != nil {
			return err
		}
		fallthrough
	case reflect.Array:
		// special case for byte arrays
		if val.Type().Elem().Kind() == reflect.Uint8 {
			// if the array is addressable, we can optimize a bit here
			if val.CanAddr() {
				return e.write(val.Slice(0, val.Len()).Bytes())
			}
			// otherwise we have to copy into a newly allocated slice
			slice := reflect.MakeSlice(reflect.SliceOf(val.Type().Elem()), val.Len(), val.Len())
			reflect.Copy(slice, val)
			return e.write(slice.Bytes())
		}
		// normal slices/arrays are encoded by sequentially encoding their elements
		for i := 0; i < val.Len(); i++ {
			if err := e.encode(val.Index(i)); err != nil {
				return err
			}
		}
		return nil
	case reflect.Struct:
		for i := 0; i < val.NumField(); i++ {
			if err := e.encode(val.Field(i)); err != nil {
				return err
			}
		}
		return nil
	}

	// Marshalling should never fail. If it panics, you're doing something wrong,
	// like trying to encode a map or an unexported struct field.
	panic("could not marshal type " + val.Type().String())
}

// NewEncoder returns a new encoder that writes to w.
func NewEncoder(w io.Writer) *Encoder {
	return &Encoder{w}
}

// Marshal returns the encoding of v. For encoding details, see the package
// docstring.
func Marshal(v interface{}) []byte {
	b := new(bytes.Buffer)
	NewEncoder(b).Encode(v) // no error possible when using a bytes.Buffer
	return b.Bytes()
}

// MarshalAll encodes all of its inputs and returns their concatenation.
func MarshalAll(v ...interface{}) []byte {
	b := new(bytes.Buffer)
	enc := NewEncoder(b)
	for i := range v {
		enc.Encode(v[i]) // no error possible when using a bytes.Buffer
	}
	return b.Bytes()
}

// WriteFile writes v to a file. The file will be created if it does not exist.
func WriteFile(filename string, v interface{}) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	return NewEncoder(file).Encode(v)
}

// A Decoder reads and decodes values from an input stream.
type Decoder struct {
	r io.Reader
	n int
}

// Read implements the io.Reader interface. It also keeps track of the total
// number of bytes decoded, and panics if that number exceeds a global
// maximum.
func (d *Decoder) Read(p []byte) (int, error) {
	n, err := d.r.Read(p)
	// enforce an absolute maximum size limit
	if d.n += n; d.n > maxDecodeLen {
		panic("encoded type exceeds size limit")
	}
	return n, err
}

// Decode reads the next encoded value from its input stream and stores it in
// v, which must be a pointer. The decoding rules are the inverse of those
// specified in the package docstring.
func (d *Decoder) Decode(v interface{}) (err error) {
	// v must be a pointer
	pval := reflect.ValueOf(v)
	if pval.Kind() != reflect.Ptr || pval.IsNil() {
		return errBadPointer
	}

	// catch decoding panics and convert them to errors
	// note that this allows us to skip boundary checks during decoding
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("could not decode type %s: %v", pval.Elem().Type().String(), r)
		}
	}()

	// reset the read count
	d.n = 0

	d.decode(pval.Elem())
	return
}

// readN reads n bytes and panics if the read fails.
func (d *Decoder) readN(n int) []byte {
	b := make([]byte, n)
	_, err := io.ReadFull(d, b)
	if err != nil {
		panic(err)
	}
	return b
}

// readPrefix reads a length-prefixed byte slice and panics if the read fails.
func (d *Decoder) readPrefix() []byte {
	b, err := ReadPrefix(d, maxSliceLen)
	if err != nil {
		panic(err)
	}
	return b
}

// decode reads the next encoded value from its input stream and stores it in
// val. The decoding rules are the inverse of those specified in the package
// docstring.
func (d *Decoder) decode(val reflect.Value) {
	// check for UnmarshalSia interface first
	if val.CanAddr() {
		if u, ok := val.Addr().Interface().(SiaUnmarshaler); ok {
			u.UnmarshalSia(d.readPrefix())
			return
		}
	}

	switch val.Kind() {
	case reflect.Ptr:
		var valid bool
		d.decode(reflect.ValueOf(&valid).Elem())
		// nil pointer, nothing to decode
		if !valid {
			return
		}
		// make sure we aren't decoding into nil
		if val.IsNil() {
			val.Set(reflect.New(val.Type().Elem()))
		}
		d.decode(val.Elem())
	case reflect.Bool:
		b := d.readN(1)
		if b[0] > 1 {
			panic("boolean value was not 0 or 1")
		}
		val.SetBool(b[0] == 1)
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		val.SetInt(DecInt64(d.readN(8)))
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		val.SetUint(DecUint64(d.readN(8)))
	case reflect.String:
		val.SetString(string(d.readPrefix()))
	case reflect.Slice:
		// slices are variable length, but otherwise the same as arrays.
		// just have to allocate them first, then we can fallthrough to the array logic.
		sliceLen := DecUint64(d.readN(8))
		// sanity-check the sliceLen, otherwise you can crash a peer by making
		// them allocate a massive slice
		if sliceLen > 1<<31-1 || sliceLen*uint64(val.Type().Elem().Size()) > maxSliceLen {
			panic("slice is too large")
		}
		val.Set(reflect.MakeSlice(val.Type(), int(sliceLen), int(sliceLen)))
		fallthrough
	case reflect.Array:
		// special case for byte arrays (e.g. hashes)
		if val.Type().Elem().Kind() == reflect.Uint8 {
			// convert val to a slice and read into it directly
			b := val.Slice(0, val.Len())
			_, err := io.ReadFull(d, b.Bytes())
			if err != nil {
				panic(err)
			}
			return
		}
		// arrays are unmarshalled by sequentially unmarshalling their elements
		for i := 0; i < val.Len(); i++ {
			d.decode(val.Index(i))
		}
		return
	case reflect.Struct:
		for i := 0; i < val.NumField(); i++ {
			d.decode(val.Field(i))
		}
		return
	default:
		panic("unknown type")
	}
}

// NewDecoder returns a new decoder that reads from r.
func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{r, 0}
}

// Unmarshal decodes the encoded value b and stores it in v, which must be a
// pointer. The decoding rules are the inverse of those specified in the
// package docstring.
func Unmarshal(b []byte, v interface{}) error {
	r := bytes.NewReader(b)
	return NewDecoder(r).Decode(v)
}

// ReadFile reads the contents of a file and decodes them into v.
func ReadFile(filename string, v interface{}) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	return NewDecoder(file).Decode(v)
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 Sevki <s@sevki.org>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package builder parses build graphs and coordinates builds
package postprocessor

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

	"bldy.build/build"
	"bldy.build/build/util"
)

type PostProcessor struct {
	projectPath, packagePath string
}

// New returns a new PostProcessor
func New(p string) PostProcessor {
	return PostProcessor{
		packagePath: p,
		projectPath: util.GetProjectPath(),
	}
}

// ProcessDependencies takes relative dependency paths and turns then in to
// absolute paths.
func (pp *PostProcessor) ProcessDependencies(t build.Target) error {

	v := reflect.ValueOf(t)

	deps := v.Elem().FieldByName("Dependencies").Interface().([]string)

	seen := make(map[string]bool)

	for i, d := range deps {
		if _, ok := seen[d]; ok {
			return fmt.Errorf("post process dependencies: %s is duplicated", d)
		} else {
			seen[d] = true
		}
		switch {
		case d[:2] == "//":
			continue
		case d[0] == ':':
			deps[i] = fmt.Sprintf("//%s%s", pp.packagePath, d)
			break
		default:
			errorf := `dependency '%s' in %s is not a valid URL for a target.
		a target url can only start with a '//' or a ':' for relative targets.`

			return fmt.Errorf(errorf, d, t.GetName())
		}

		seen[d] = true
	}

	v.Elem().FieldByName("Dependencies").Set(reflect.ValueOf(deps))
	return nil
}

// ProcesPaths takes paths relative to the target and absolutes them,
// unless they are going to be exported in to the target folder from a dependency.
func (pp *PostProcessor) ProcessPaths(t build.Target, deps []build.Target) error {

	v := reflect.ValueOf(t)

	r := reflect.TypeOf(t).Elem()

	for i := 0; i < r.NumField(); i++ {
		f := r.Field(i)

		tag := f.Tag.Get("build")
		if !(tag == "path" || tag == "expand") {
			continue
		}

		n := v.Elem().FieldByName(f.Name)

		isExported := func(s string) bool {
			for _, d := range deps {
				if _, ok := d.Installs()[s]; ok {
					return true
				}
			}
			return false
		}
		exp := func(s string) string {
			if tag == "path" {
				return pp.absPath(s)
			}
			if tag == "expand" {
				return os.Expand(s, util.Getenv)
			}
			return s
		}
		switch n.Kind() {
		case reflect.String:
			s := n.Interface().(string)
			if isExported(s) {
				continue
			}
			if s == "" {
				continue
			}
			n.SetString(exp(s))
		case reflect.Slice:
			switch n.Type().Elem().Kind() {
			case reflect.String:
				strs := n.Convert(reflect.TypeOf([]string{})).Interface().([]string)
				for i, s := range strs {
					if isExported(s) {
						continue
					}
					strs[i] = exp(s)
				}
				n.Set(reflect.ValueOf(strs))
			}
		case reflect.Map:
			switch n.Type().Elem().Kind() {
			case reflect.String:
				strs := n.Convert(reflect.TypeOf(map[string]string{})).Interface().(map[string]string)
				for k, v := range strs {
					if isExported(k) {
						continue
					}
					delete(strs, k)

					k, v = exp(k), exp(v)

					strs[k] = v
				}
				n.Set(reflect.ValueOf(strs))
			}
		}

	}

	return nil
}

func (pp *PostProcessor) absPath(s string) string {

	if len(s) < 2 {
		log.Fatalf("%s is invalid", s)
	}
	var r string
	switch {
	case s[:2] == "//":
		r = filepath.Join(pp.projectPath, strings.Trim(s, "//"))
	default:
                r = os.Expand(s, util.Getenv)
                if filepath.IsAbs(r) {
                        return r
                }
                r = filepath.Join(pp.projectPath, pp.packagePath, s)
	}
	r = os.Expand(r, util.Getenv)
	return r
}
<commit_msg>Applied gofmt<commit_after>// Copyright 2015-2016 Sevki <s@sevki.org>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package builder parses build graphs and coordinates builds
package postprocessor

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

	"bldy.build/build"
	"bldy.build/build/util"
)

type PostProcessor struct {
	projectPath, packagePath string
}

// New returns a new PostProcessor
func New(p string) PostProcessor {
	return PostProcessor{
		packagePath: p,
		projectPath: util.GetProjectPath(),
	}
}

// ProcessDependencies takes relative dependency paths and turns then in to
// absolute paths.
func (pp *PostProcessor) ProcessDependencies(t build.Target) error {

	v := reflect.ValueOf(t)

	deps := v.Elem().FieldByName("Dependencies").Interface().([]string)

	seen := make(map[string]bool)

	for i, d := range deps {
		if _, ok := seen[d]; ok {
			return fmt.Errorf("post process dependencies: %s is duplicated", d)
		} else {
			seen[d] = true
		}
		switch {
		case d[:2] == "//":
			continue
		case d[0] == ':':
			deps[i] = fmt.Sprintf("//%s%s", pp.packagePath, d)
			break
		default:
			errorf := `dependency '%s' in %s is not a valid URL for a target.
		a target url can only start with a '//' or a ':' for relative targets.`

			return fmt.Errorf(errorf, d, t.GetName())
		}

		seen[d] = true
	}

	v.Elem().FieldByName("Dependencies").Set(reflect.ValueOf(deps))
	return nil
}

// ProcesPaths takes paths relative to the target and absolutes them,
// unless they are going to be exported in to the target folder from a dependency.
func (pp *PostProcessor) ProcessPaths(t build.Target, deps []build.Target) error {

	v := reflect.ValueOf(t)

	r := reflect.TypeOf(t).Elem()

	for i := 0; i < r.NumField(); i++ {
		f := r.Field(i)

		tag := f.Tag.Get("build")
		if !(tag == "path" || tag == "expand") {
			continue
		}

		n := v.Elem().FieldByName(f.Name)

		isExported := func(s string) bool {
			for _, d := range deps {
				if _, ok := d.Installs()[s]; ok {
					return true
				}
			}
			return false
		}
		exp := func(s string) string {
			if tag == "path" {
				return pp.absPath(s)
			}
			if tag == "expand" {
				return os.Expand(s, util.Getenv)
			}
			return s
		}
		switch n.Kind() {
		case reflect.String:
			s := n.Interface().(string)
			if isExported(s) {
				continue
			}
			if s == "" {
				continue
			}
			n.SetString(exp(s))
		case reflect.Slice:
			switch n.Type().Elem().Kind() {
			case reflect.String:
				strs := n.Convert(reflect.TypeOf([]string{})).Interface().([]string)
				for i, s := range strs {
					if isExported(s) {
						continue
					}
					strs[i] = exp(s)
				}
				n.Set(reflect.ValueOf(strs))
			}
		case reflect.Map:
			switch n.Type().Elem().Kind() {
			case reflect.String:
				strs := n.Convert(reflect.TypeOf(map[string]string{})).Interface().(map[string]string)
				for k, v := range strs {
					if isExported(k) {
						continue
					}
					delete(strs, k)

					k, v = exp(k), exp(v)

					strs[k] = v
				}
				n.Set(reflect.ValueOf(strs))
			}
		}

	}

	return nil
}

func (pp *PostProcessor) absPath(s string) string {

	if len(s) < 2 {
		log.Fatalf("%s is invalid", s)
	}
	var r string
	switch {
	case s[:2] == "//":
		r = filepath.Join(pp.projectPath, strings.Trim(s, "//"))
	default:
		r = os.Expand(s, util.Getenv)
		if filepath.IsAbs(r) {
			return r
		}
		r = filepath.Join(pp.projectPath, pp.packagePath, s)
	}
	r = os.Expand(r, util.Getenv)
	return r
}
<|endoftext|>


<commit_before>package myaws

import (
	"fmt"
	"os"

	"github.com/spf13/cobra"
)

// UsageError prints a error message and show usage, then exit with return -1.
func UsageError(cmd *cobra.Command, msg interface{}) {
	fmt.Println("Error:", msg)
	cmd.Usage()
	os.Exit(-1)
}
<commit_msg>Remove myaws.UsageError<commit_after><|endoftext|>


<commit_before>package main

import (
	"bytes"
	"database/sql"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"strings"
	"syscall"
	"time"
	"unsafe"

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

const insertBufferSize = 1048576

// Type definitions
type (
	// dbInfo contains information about a database
	dbInfo struct {
		db      *sql.DB
		user    string
		pass    string
		host    string
		port    string
		sock    string
		schema  string
		table   string
		where   string
		columns []string
	}
)

func main() {
	start := time.Now()

	// Trap for SIGINT, may need to trap other signals in the future as well
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)

	// Syscall in case signal is sent during terminal echo off
	var oldState syscall.Termios
	syscall.Syscall6(syscall.SYS_IOCTL, uintptr(0), syscall.TCGETS, uintptr(unsafe.Pointer(&oldState)), 0, 0, 0)

	var timer time.Time
	go func() {
		for sig := range sigChan {
			if time.Now().Sub(timer) < time.Second*5 {
				syscall.Syscall6(syscall.SYS_IOCTL, uintptr(0), syscall.TCSETS, uintptr(unsafe.Pointer(&oldState)), 0, 0, 0)
				os.Exit(0)
			}

			fmt.Println()
			fmt.Println(sig, "signal caught!")
			fmt.Println("Send signal again within 3 seconds to exit")

			timer = time.Now()
		}
	}()

	// Profiling flags
	var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	var memprofile = flag.String("memprofile", "", "write memory profile to this file")

	// Source flags
	fSrcUser := flag.String("srcuser", "", "Source Username")
	fSrcPass := flag.String("srcpassword", "", "Source Password")
	fSrcHost := flag.String("srchost", "", "Source Database")
	fSrcPort := flag.String("srcport", "3306", "Source MySQL Port")
	fSrcSock := flag.String("srcsocket", "", "Source MySQL Socket")
	fSrcTable := flag.String("srctable", "", "Source Table (must be fully qualified ex. schema.tablename)")
	fSrcWhere := flag.String("where", "", "Where clause to apply to source table select")

	// Target flags
	fTgtUser := flag.String("tgtuser", "", "Target Username")
	fTgtPass := flag.String("tgtpassword", "", "Target Password")
	fTgtHost := flag.String("tgthost", "", "Target Database")
	fTgtPort := flag.String("tgtport", "3306", "Target MySQL Port")
	fTgtTable := flag.String("tgttable", "", "Target Table (must be fully qualified ex. schema.tablename)")

	flag.Parse()

	// CPU Profiling
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		checkErr(err)
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// Default to localhost if no host or socket provided
	if *fSrcSock == "" && *fSrcHost == "" {
		*fSrcSock = "/var/lib/mysql/mysql.sock"
	}

	// Need to provide a target
	if *fTgtHost == "" {
		fmt.Println("You must provide a target database")
		os.Exit(1)
	}

	// A fully qualified table must be provided
	if *fTgtTable == "" && *fSrcTable != "" {
		*fTgtTable = *fSrcTable
	} else if *fSrcTable == "" || !strings.Contains(*fSrcTable, ".") {
		fmt.Println("You must provide a fully qualifed table to move")
		os.Exit(1)
	}

	// If password is blank prompt user - Not perfect as it prints the password typed to the screen
	if *fSrcPass == "" {
		fmt.Println("Enter password: ")
		pwd, err := readPassword(0)
		checkErr(err)
		*fSrcPass = string(pwd)
	}

	if *fTgtUser == "" && *fSrcUser != "" {
		*fTgtUser = *fSrcUser
	}

	if *fTgtPass == "" && *fSrcPass != "" {
		*fTgtPass = *fSrcPass
	}

	// Add where keyword if where clause is supplied
	fmt.Println(*fSrcWhere)
	if *fSrcWhere != "" {
		*fSrcWhere = " where " + *fSrcWhere
	}

	// Split the table into schema and table name
	srcSplit := strings.Split(*fSrcTable, ".")
	tgtSplit := strings.Split(*fTgtTable, ".")

	source := dbInfo{user: *fSrcUser, pass: *fSrcPass, host: *fSrcHost, port: *fSrcPort, sock: *fSrcSock, schema: srcSplit[0], table: srcSplit[1], where: *fSrcWhere}
	target := dbInfo{user: *fTgtUser, pass: *fTgtPass, host: *fTgtHost, port: *fTgtPort, schema: tgtSplit[0], table: tgtSplit[1]}

	// Create a *sql.DB connection to the source database
	sourceDB, err := source.dbConn()
	defer sourceDB.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	source.db = sourceDB

	// Create a *sql.DB connection to the target database
	targetDB, err := target.dbConn()
	defer targetDB.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	target.db = targetDB

	createStmt := source.readTable()

	target.columns = source.getColumns()

	checkSchema(&source, &target)
	writeTable(&source, &target, createStmt)

	dataChan := make(chan []sql.RawBytes)
	quitChan := make(chan bool)
	goChan := make(chan bool)
	go readRows(&source, &target, dataChan, quitChan, goChan)
	target.writeRows(dataChan, goChan)

	<-quitChan
	close(quitChan)
	close(goChan)

	//  fmt.Println(source)
	//  fmt.Println(target)

	// Memory Profiling
	if *memprofile != "" {
		f, err := os.Create(*memprofile)
		checkErr(err)
		pprof.WriteHeapProfile(f)
		defer f.Close()
	}

	fmt.Println()
	fmt.Println("Total runtime =", time.Since(start))
}

// Pass the buck error catching
func checkErr(e error) {
	if e != nil {
		log.Panic(e)
	}
}

// readPassword is borrowed from the crypto/ssh/terminal sub repo to accept a password from stdin without local echo.
// http://godoc.org/code.google.com/p/go.crypto/ssh/terminal#Terminal.ReadPassword
func readPassword(fd int) ([]byte, error) {
	var oldState syscall.Termios
	if _, _, err := syscall.Syscall6(syscall.SYS_IOCTL, uintptr(fd), syscall.TCGETS, uintptr(unsafe.Pointer(&oldState)), 0, 0, 0); err != 0 {
		return nil, err
	}

	newState := oldState
	newState.Lflag &^= syscall.ECHO
	newState.Lflag |= syscall.ICANON | syscall.ISIG
	newState.Iflag |= syscall.ICRNL
	if _, _, err := syscall.Syscall6(syscall.SYS_IOCTL, uintptr(fd), syscall.TCSETS, uintptr(unsafe.Pointer(&newState)), 0, 0, 0); err != 0 {
		return nil, err
	}

	defer func() {
		syscall.Syscall6(syscall.SYS_IOCTL, uintptr(fd), syscall.TCSETS, uintptr(unsafe.Pointer(&oldState)), 0, 0, 0)
	}()

	var buf [16]byte
	var ret []byte
	for {
		n, err := syscall.Read(fd, buf[:])
		if err != nil {
			return nil, err
		}

		if n == 0 {
			if len(ret) == 0 {
				return nil, io.EOF
			}
			break
		}

		if buf[n-1] == '\n' {
			n--
		}

		ret = append(ret, buf[:n]...)
		if n < len(buf) {
			break
		}
	}

	return ret, nil
}

func (d *dbInfo) dbConn() (*sql.DB, error) {
	// Determine tcp or socket connection
	var db *sql.DB
	var err error
	if d.sock != "" {
		db, err = sql.Open("mysql", d.user+":"+d.pass+"@unix("+d.sock+")/?allowCleartextPasswords=1")
		checkErr(err)
	} else if d.host != "" {
		db, err = sql.Open("mysql", d.user+":"+d.pass+"@tcp("+d.host+":"+d.port+")/?allowCleartextPasswords=1")
		checkErr(err)
	} else {
		fmt.Println("should be no else")
	}

	// Ping database to verify credentials
	err = db.Ping()

	return db, err
}

// Adds backtick quotes in cases where identifiers are all numeric or match reserved keywords
func addQuotes(s string) string {
	s = "`" + s + "`"
	return s
}

func (src *dbInfo) readTable() string {
	var err error
	var ignore string
	var stmt string
	err = src.db.QueryRow("show create table "+addQuotes(src.schema)+"."+addQuotes(src.table)).Scan(&ignore, &stmt)
	checkErr(err)

	return stmt
}

func (src *dbInfo) getColumns() []string {
	var cols = []string{}
	rows, err := src.db.Query("select data_type from information_schema.columns where table_schema = '" + src.schema + "' and table_name = '" + src.table + "'")
	defer rows.Close()
	checkErr(err)

	var dataType string
	for rows.Next() {
		err = rows.Scan(&dataType)
		checkErr(err)

		cols = append(cols, dataType)
	}
	checkErr(err)

	return cols
}

// checkSchema creates a schema if it does not already exist
func checkSchema(src, tgt *dbInfo) {
	var exists string
	err := tgt.db.QueryRow("show databases like '" + tgt.schema + "'").Scan(&exists)

	if err != nil {
		var charSet string
		err := src.db.QueryRow("select default_character_set_name from information_schema.schemata where schema_name = '" + src.schema + "'").Scan(&charSet)

		_, err = tgt.db.Exec("create database " + addQuotes(tgt.schema) + " default character set " + charSet)
		checkErr(err)

		fmt.Println("       Created schema", tgt.schema)
	}
}

func writeTable(src, tgt *dbInfo, tableCreate string) {
	// Start db transaction
	tx, err := tgt.db.Begin()
	checkErr(err)

	_, err = tx.Exec("use " + tgt.schema)

	// Drop table if exists
	_, err = tx.Exec("drop table if exists " + addQuotes(tgt.table))
	checkErr(err)

	// Change table name if different
	if src.table != tgt.table {
		tableCreate = strings.Replace(tableCreate, src.table, tgt.table, 1)
	}

	// Create table
	_, err = tx.Exec(tableCreate)
	checkErr(err)

	// Commit transaction
	err = tx.Commit()
	checkErr(err)
}

func readRows(src, tgt *dbInfo, dataChan chan []sql.RawBytes, quitChan chan bool, goChan chan bool) {
	rows, err := src.db.Query("select * from " + addQuotes(src.schema) + "." + addQuotes(src.table) + src.where)
	defer rows.Close()
	checkErr(err)

	cols, err := rows.Columns()
	checkErr(err)

	// Need to scan into empty interface since we don't know how many columns or their types
	scanVals := make([]interface{}, len(cols))
	vals := make([]sql.RawBytes, len(cols))
	stray := make([]sql.RawBytes, len(cols))
	for i := range vals {
		scanVals[i] = &vals[i]
	}

	var rowNum int64
	for rows.Next() {
		err := rows.Scan(scanVals...)
		checkErr(err)

		copy(stray, vals)

		dataChan <- stray
		rowNum++

		// Block until writeRows() signals it is safe to proceed
		// This is necessary because sql.RawBytes is a memory pointer and rows.Next() will loop and change the memory address before writeRows can properly process the values
		<-goChan
	}

	err = rows.Err()
	checkErr(err)

	close(dataChan)
	quitChan <- true

	fmt.Println(rowNum, "rows inserted")
}

func (target *dbInfo) writeRows(dataChan chan []sql.RawBytes, goChan chan bool) {
	var cleaned []byte
	buf := bytes.NewBuffer(make([]byte, 0, insertBufferSize))

	sqlPrefix := "insert into " + addQuotes(target.schema) + "." + addQuotes(target.table) + " values ("
	prefix, _ := buf.WriteString(sqlPrefix)

	append := false
	for data := range dataChan {
		if append {
			buf.WriteString(",(")
		}
		append = true

		for i, col := range data {
			if col == nil {
				buf.WriteString("NULL")
			} else if len(col) == 0 {
				buf.WriteString("''")
			} else {
				switch target.columns[i] {
				case "tinytext":
					fallthrough
				case "text":
					fallthrough
				case "mediumtext":
					fallthrough
				case "longtext":
					fallthrough
				case "varchar":
					fallthrough
				case "char":
					cleaned = col
					cleaned = bytes.Replace(cleaned, []byte(`\`), []byte(`\\`), -1)
					cleaned = bytes.Replace(cleaned, []byte(`'`), []byte(`\'`), -1)
					buf.WriteString("'")
					buf.Write(cleaned)
					buf.WriteString("'")
				default:
					buf.WriteString("'")
					buf.Write(col)
					buf.WriteString("'")
				}
			}

			if i < len(target.columns)-1 {
				buf.WriteString(",")
			}
		}

		buf.WriteString(")")

		if buf.Len() > insertBufferSize {
			//buf.WriteTo(os.Stdout) // DEBUG
			//fmt.Println()          // DEBUG
			_, err := target.db.Exec(buf.String())
			if err != nil {
				fmt.Println(buf.String())
			}
			checkErr(err)

			buf.Reset()
			buf.WriteString(sqlPrefix)
			append = false
		}

		// Allow read function to loop over rows
		goChan <- true
	}

	// Apply left over data
	if buf.Len() > prefix {
		//buf.WriteTo(os.Stdout) // DEBUG
		//fmt.Println()          // DEBUG
		_, err := target.db.Exec(buf.String())
		checkErr(err)
	}
}
<commit_msg>Version info, use ssh/terminal for reading pwds<commit_after>package main

import (
	"bytes"
	"database/sql"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"strings"
	"time"

	"golang.org/x/crypto/ssh/terminal"

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

const (
	// Insert buffer size
	insertBufferSize = 1048576 // 1MB

	// Timeout length where ctrl+c is ignored.
	signalTimeout = 3 // Seconds
)

// Type definitions
type (
	// dbInfo contains information about a database
	dbInfo struct {
		db      *sql.DB
		user    string
		pass    string
		host    string
		port    string
		sock    string
		schema  string
		table   string
		where   string
		columns []string
	}
)

// Version information supplied by build script
var versionInformation string

// ShowUsage prints a help screen
func showUsage() {
	fmt.Printf("\tmytablecopy version %s\n", versionInformation)
	fmt.Println(`
			  `)
}

func main() {
	start := time.Now()

	// Catch signals
	catchNotifications()

	// Profiling flags
	var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	var memprofile = flag.String("memprofile", "", "write memory profile to this file")

	// Source flags
	fSrcUser := flag.String("srcuser", "", "Source Username")
	fSrcPass := flag.String("srcpassword", "", "Source Password")
	fSrcHost := flag.String("srchost", "", "Source Database")
	fSrcPort := flag.String("srcport", "3306", "Source MySQL Port")
	fSrcSock := flag.String("srcsocket", "", "Source MySQL Socket")
	fSrcTable := flag.String("srctable", "", "Source Table (must be fully qualified ex. schema.tablename)")
	fSrcWhere := flag.String("where", "", "Where clause to apply to source table select")

	// Target flags
	fTgtUser := flag.String("tgtuser", "", "Target Username")
	fTgtPass := flag.String("tgtpassword", "", "Target Password")
	fTgtHost := flag.String("tgthost", "", "Target Database")
	fTgtPort := flag.String("tgtport", "3306", "Target MySQL Port")
	fTgtTable := flag.String("tgttable", "", "Target Table (must be fully qualified ex. schema.tablename)")

	// Other flags
	version := flag.Bool("version", false, "Version information")
	help := flag.Bool("help", false, "Show usage")
	h := flag.Bool("h", false, "Show usage")

	flag.Parse()

	// Print usage
	if flag.NFlag() == 0 || *help == true || *h == true {
		showUsage()

		os.Exit(0)
	}

	if *version {
		fmt.Printf("mytablecopy version %s\n", versionInformation)
		os.Exit(0)
	}

	// CPU Profiling
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		checkErr(err)
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// Default to localhost if no host or socket provided
	if *fSrcSock == "" && *fSrcHost == "" {
		*fSrcSock = "/var/lib/mysql/mysql.sock"
	}

	// Need to provide a target
	if *fTgtHost == "" {
		fmt.Println("You must provide a target database")
		os.Exit(1)
	}

	// A fully qualified table must be provided
	if *fTgtTable == "" && *fSrcTable != "" {
		*fTgtTable = *fSrcTable
	} else if *fSrcTable == "" || !strings.Contains(*fSrcTable, ".") {
		fmt.Println("You must provide a fully qualifed table to move")
		os.Exit(1)
	}

	// If password is blank prompt user - Not perfect as it prints the password typed to the screen
	if *fSrcPass == "" {
		fmt.Println("Enter password: ")
		//		pwd, err := readPassword(0)
		pwd, err := terminal.ReadPassword(int(os.Stdin.Fd()))
		if err != nil {
			if err != io.EOF {
				checkErr(err)
			}
		}

		*fSrcPass = string(pwd)
	}

	if *fTgtUser == "" && *fSrcUser != "" {
		*fTgtUser = *fSrcUser
	}

	if *fTgtPass == "" && *fSrcPass != "" {
		*fTgtPass = *fSrcPass
	}

	// Add where keyword if where clause is supplied
	if *fSrcWhere != "" {
		*fSrcWhere = " where " + *fSrcWhere
	}

	// Split the table into schema and table name
	srcSplit := strings.Split(*fSrcTable, ".")
	tgtSplit := strings.Split(*fTgtTable, ".")

	source := dbInfo{user: *fSrcUser, pass: *fSrcPass, host: *fSrcHost, port: *fSrcPort, sock: *fSrcSock, schema: srcSplit[0], table: srcSplit[1], where: *fSrcWhere}
	target := dbInfo{user: *fTgtUser, pass: *fTgtPass, host: *fTgtHost, port: *fTgtPort, schema: tgtSplit[0], table: tgtSplit[1]}

	// Create a *sql.DB connection to the source database
	sourceDB, err := source.dbConn()
	defer sourceDB.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	source.db = sourceDB

	// Create a *sql.DB connection to the target database
	targetDB, err := target.dbConn()
	defer targetDB.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	target.db = targetDB

	createStmt := source.readTable()

	target.columns = source.getColumns()

	checkSchema(&source, &target)
	writeTable(&source, &target, createStmt)

	dataChan := make(chan []sql.RawBytes)
	quitChan := make(chan bool)
	goChan := make(chan bool)
	go readRows(&source, &target, dataChan, quitChan, goChan)
	target.writeRows(dataChan, goChan)

	<-quitChan
	close(quitChan)
	close(goChan)

	//  fmt.Println(source)
	//  fmt.Println(target)

	// Memory Profiling
	if *memprofile != "" {
		f, err := os.Create(*memprofile)
		checkErr(err)
		pprof.WriteHeapProfile(f)
		defer f.Close()
	}

	fmt.Println()
	fmt.Println("Total runtime =", time.Since(start))
}

// Pass the buck error catching
func checkErr(e error) {
	if e != nil {
		log.Panic(e)
	}
}

// Catch signals
func catchNotifications() {
	state, err := terminal.GetState(int(os.Stdin.Fd()))
	checkErr(err)

	// Deal with SIGINT
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)

	var timer time.Time
	go func() {
		for sig := range sigChan {
			// Prevent exiting on accidental signal send
			if time.Now().Sub(timer) < time.Second*signalTimeout {
				terminal.Restore(int(os.Stdin.Fd()), state)
				os.Exit(0)
			}

			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, sig, "signal caught!")
			fmt.Fprintf(os.Stderr, "Send signal again within %v seconds to exit\n", signalTimeout)
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")

			timer = time.Now()
		}
	}()
}

func (d *dbInfo) dbConn() (*sql.DB, error) {
	// Determine tcp or socket connection
	var db *sql.DB
	var err error
	if d.sock != "" {
		db, err = sql.Open("mysql", d.user+":"+d.pass+"@unix("+d.sock+")/?allowCleartextPasswords=1")
		checkErr(err)
	} else if d.host != "" {
		db, err = sql.Open("mysql", d.user+":"+d.pass+"@tcp("+d.host+":"+d.port+")/?allowCleartextPasswords=1")
		checkErr(err)
	} else {
		fmt.Println("should be no else")
	}

	// Ping database to verify credentials
	err = db.Ping()

	return db, err
}

// Adds backtick quotes in cases where identifiers are all numeric or match reserved keywords
func addQuotes(s string) string {
	s = "`" + s + "`"
	return s
}

func (src *dbInfo) readTable() string {
	var err error
	var ignore string
	var stmt string
	err = src.db.QueryRow("show create table "+addQuotes(src.schema)+"."+addQuotes(src.table)).Scan(&ignore, &stmt)
	checkErr(err)

	return stmt
}

func (src *dbInfo) getColumns() []string {
	var cols = []string{}
	rows, err := src.db.Query("select data_type from information_schema.columns where table_schema = '" + src.schema + "' and table_name = '" + src.table + "'")
	defer rows.Close()
	checkErr(err)

	var dataType string
	for rows.Next() {
		err = rows.Scan(&dataType)
		checkErr(err)

		cols = append(cols, dataType)
	}
	checkErr(err)

	return cols
}

// checkSchema creates a schema if it does not already exist
func checkSchema(src, tgt *dbInfo) {
	var exists string
	err := tgt.db.QueryRow("show databases like '" + tgt.schema + "'").Scan(&exists)

	if err != nil {
		var charSet string
		err := src.db.QueryRow("select default_character_set_name from information_schema.schemata where schema_name = '" + src.schema + "'").Scan(&charSet)

		_, err = tgt.db.Exec("create database " + addQuotes(tgt.schema) + " default character set " + charSet)
		checkErr(err)

		fmt.Println("       Created schema", tgt.schema)
	}
}

func writeTable(src, tgt *dbInfo, tableCreate string) {
	// Start db transaction
	tx, err := tgt.db.Begin()
	checkErr(err)

	_, err = tx.Exec("use " + tgt.schema)

	// Drop table if exists
	_, err = tx.Exec("drop table if exists " + addQuotes(tgt.table))
	checkErr(err)

	// Change table name if different
	if src.table != tgt.table {
		tableCreate = strings.Replace(tableCreate, src.table, tgt.table, 1)
	}

	// Create table
	_, err = tx.Exec(tableCreate)
	checkErr(err)

	// Commit transaction
	err = tx.Commit()
	checkErr(err)
}

func readRows(src, tgt *dbInfo, dataChan chan []sql.RawBytes, quitChan chan bool, goChan chan bool) {
	rows, err := src.db.Query("select * from " + addQuotes(src.schema) + "." + addQuotes(src.table) + src.where)
	defer rows.Close()
	checkErr(err)

	cols, err := rows.Columns()
	checkErr(err)

	// Need to scan into empty interface since we don't know how many columns or their types
	scanVals := make([]interface{}, len(cols))
	vals := make([]sql.RawBytes, len(cols))
	stray := make([]sql.RawBytes, len(cols))
	for i := range vals {
		scanVals[i] = &vals[i]
	}

	var rowNum int64
	for rows.Next() {
		err := rows.Scan(scanVals...)
		checkErr(err)

		copy(stray, vals)

		dataChan <- stray
		rowNum++

		// Block until writeRows() signals it is safe to proceed
		// This is necessary because sql.RawBytes is a memory pointer and rows.Next() will loop and change the memory address before writeRows can properly process the values
		<-goChan
	}

	err = rows.Err()
	checkErr(err)

	close(dataChan)
	quitChan <- true

	fmt.Println(rowNum, "rows inserted")
}

func (target *dbInfo) writeRows(dataChan chan []sql.RawBytes, goChan chan bool) {
	var cleaned []byte
	buf := bytes.NewBuffer(make([]byte, 0, insertBufferSize))

	sqlPrefix := "insert into " + addQuotes(target.schema) + "." + addQuotes(target.table) + " values ("
	prefix, _ := buf.WriteString(sqlPrefix)

	append := false
	for data := range dataChan {
		if append {
			buf.WriteString(",(")
		}
		append = true

		for i, col := range data {
			if col == nil {
				buf.WriteString("NULL")
			} else if len(col) == 0 {
				buf.WriteString("''")
			} else {
				switch target.columns[i] {
				case "tinytext":
					fallthrough
				case "text":
					fallthrough
				case "mediumtext":
					fallthrough
				case "longtext":
					fallthrough
				case "varchar":
					fallthrough
				case "char":
					cleaned = col
					cleaned = bytes.Replace(cleaned, []byte(`\`), []byte(`\\`), -1)
					cleaned = bytes.Replace(cleaned, []byte(`'`), []byte(`\'`), -1)
					buf.WriteString("'")
					buf.Write(cleaned)
					buf.WriteString("'")
				default:
					buf.WriteString("'")
					buf.Write(col)
					buf.WriteString("'")
				}
			}

			if i < len(target.columns)-1 {
				buf.WriteString(",")
			}
		}

		buf.WriteString(")")

		if buf.Len() > insertBufferSize {
			//buf.WriteTo(os.Stdout) // DEBUG
			//fmt.Println()          // DEBUG
			_, err := target.db.Exec(buf.String())
			if err != nil {
				fmt.Println(buf.String())
			}
			checkErr(err)

			buf.Reset()
			buf.WriteString(sqlPrefix)
			append = false
		}

		// Allow read function to loop over rows
		goChan <- true
	}

	// Apply left over data
	if buf.Len() > prefix {
		//buf.WriteTo(os.Stdout) // DEBUG
		//fmt.Println()          // DEBUG
		_, err := target.db.Exec(buf.String())
		checkErr(err)
	}
}
<|endoftext|>


<commit_before>// Package natty provides a Go language wrapper to the natty NAT traversal
// utility.  See https://github.com/getlantern/t.
//
// See natty_test for an example of Natty in use, including debug logging
// showing the messages that are sent across the signaling channel.
package natty

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/oxtoacart/byteexec"
)

const (
	UDP = Protocol("udp")
	TCP = Protocol("tcp")
)

var (
	reallyHighTimeout = 100000 * time.Hour
)

type Protocol string

// A FiveTuple is the result of a successful NAT traversal.
type FiveTuple struct {
	Proto  Protocol
	Local  string
	Remote string
}

// UDPAddrs returns a pair of UDPAddrs representing the Local and Remote
// addresses of this FiveTuple. If the FiveTuple's Proto is not UDP, this method
// returns an error.
func (ft *FiveTuple) UDPAddrs() (local *net.UDPAddr, remote *net.UDPAddr, err error) {
	if ft.Proto != UDP {
		err = fmt.Errorf("FiveTuple.Proto was not UDP!: %s", ft.Proto)
		return
	}
	local, err = net.ResolveUDPAddr("udp", ft.Local)
	if err != nil {
		err = fmt.Errorf("Unable to resolve local UDP address %s: %s", ft.Local)
		return
	}
	remote, err = net.ResolveUDPAddr("udp", ft.Remote)
	if err != nil {
		err = fmt.Errorf("Unable to resolve remote UDP address %s: %s", ft.Remote)
	}
	return
}

// Traversal represents a single NAT traversal using natty, whose result is
// available via the methods FiveTuple() and FiveTupleTimeout().
//
// Consumers should make sure to call Close() after finishing with this Natty
// in order to make sure the underlying natty process and associated resources
// are closed.
type Traversal struct {
	debugOut           io.Writer       // target for output from natty's stderr
	cmd                *exec.Cmd       // the natty command
	stdin              io.WriteCloser  // pipe to natty's stdin
	stdout             io.ReadCloser   // pipe from natty's stdout
	stdoutbuf          *bufio.Reader   // buffered stdout
	stderr             io.ReadCloser   // pipe from natty's stderr
	msgInCh            chan string     // channel for messages inbound to this Natty
	msgOutCh           chan string     // channel for messages outbound from this Natty
	peerGotFiveTupleCh chan bool       // channel to signal once we know that our peer received their own FiveTuple
	fiveTupleCh        chan *FiveTuple // intermediary channel for the FiveTuple emitted by the natty command
	errCh              chan error      // intermediary channel for any error encountered while running natty
	fiveTupleOutCh     chan *FiveTuple // channel for FiveTuple output
	errOutCh           chan error      // channel for error output
	fiveTupleOut       *FiveTuple      // the output FiveTuple
	errOut             error           // the output error
	outMutex           sync.Mutex      // mutex for synchronizing access to output variables
}

// Offer starts a Traversal as an Offerer, meaning that it will make an offer to
// initiate an ICE session. Call FiveTuple() or FiveTupleTimeout() to get the
// FiveTuple resulting from Traversal.
//
// debugOut (optional) is an optional io.Writer to which debug output from natty
// will be written.
//
func Offer(debugOut io.Writer) *Traversal {
	t := &Traversal{
		debugOut: debugOut,
	}
	t.run([]string{"-offer"})
	return t
}

// Answer starts a Traversal as an Answerer, meaning that it will accept offers
// to initiate an ICE session. Call FiveTuple() or FiveTupleTimeout() to get the
// FiveTuple resulting from Traversal.
//
// debugOut (optional) is an optional io.Writer to which debug output from Natty
// will be written.
//
func Answer(debugOut io.Writer) *Traversal {
	t := &Traversal{
		debugOut: debugOut,
	}
	t.run([]string{})
	return t
}

// MsgIn is used to pass this Traversal a message from the peer t. This method
// is buffered and will typically not block.
func (t *Traversal) MsgIn(msg string) {
	t.msgInCh <- msg
}

// NextMsgOut gets the next message to pass to the peer.  If done is true, there
// are no more messages to be read, and the currently returned message should be
// ignored.
func (t *Traversal) NextMsgOut() (msg string, done bool) {
	m, ok := <-t.msgOutCh
	return m, !ok
}

// FiveTuple gets the FiveTuple from the Traversal, blocking until such is
// available.
func (t *Traversal) FiveTuple() (*FiveTuple, error) {
	return t.FiveTupleTimeout(reallyHighTimeout)
}

// FiveTupleTimeout gets the FiveTuple from the Traversal, blocking until such
// is available or the timeout is hit.
func (t *Traversal) FiveTupleTimeout(timeout time.Duration) (*FiveTuple, error) {
	t.outMutex.Lock()
	defer t.outMutex.Unlock()

	if t.fiveTupleOut == nil && t.errOut == nil {
		// We don't have a result yet, wait for one
		select {
		case ft := <-t.fiveTupleOutCh:
			t.fiveTupleOut = ft
		case err := <-t.errOutCh:
			t.errOut = err
		case <-time.After(timeout):
			// Return an error, but don't store it (lets caller try again)
			return nil, fmt.Errorf("Timed out waiting for five-tuple")
		}
	}

	return t.fiveTupleOut, t.errOut
}

// Close closes this Traversal, terminating any outstanding natty process by
// sending SIGKILL. Close blocks until the natty process has terminated, at
// which point any ports that it bound should be available for use.
func (t *Traversal) Close() error {
	t.closePipes()

	if t.cmd != nil && t.cmd.Process != nil {
		t.cmd.Process.Kill()
		t.cmd.Process.Wait()
	}
	return nil
}

// run runs the natty command to obtain a FiveTuple. The actual running of
// natty happens on a goroutine so that run itself doesn't block.
func (t *Traversal) run(params []string) {
	t.msgInCh = make(chan string, 100)
	t.msgOutCh = make(chan string, 100)

	// Note - these channels are a buffered in order to prevent deadlocks
	t.peerGotFiveTupleCh = make(chan bool, 10)
	t.fiveTupleCh = make(chan *FiveTuple, 10)
	t.errCh = make(chan error, 10)
	t.fiveTupleOutCh = make(chan *FiveTuple, 10)
	t.errOutCh = make(chan error, 10)

	err := t.initCommand(params)

	go func() {
		if err != nil {
			t.errOutCh <- err
			return
		}

		ft, err := t.doRun(params)
		// Once doRun is finished, inform client of the FiveTuple
		if err != nil {
			t.errOutCh <- err
		} else {
			t.fiveTupleOutCh <- ft
		}
	}()
}

// doRun does the running, including resource cleanup.  doRun blocks until
// Close() has finished, meaning that natty is no longer running and whatever
// port it returned in the FiveTuple can now be used for other things.
func (t *Traversal) doRun(params []string) (*FiveTuple, error) {
	defer t.Close()

	go t.processStdout()
	go t.processStderr()

	// Start the natty command
	t.errCh <- t.cmd.Start()

	// Process incoming messages
	go func() {
		for {
			msg := <-t.msgInCh

			if IsFiveTuple(msg) {
				t.peerGotFiveTupleCh <- true
				continue
			}

			// Forward message to natty process
			_, err := t.stdin.Write([]byte(msg))
			if err == nil {
				_, err = t.stdin.Write([]byte("\n"))
			}
			if err != nil {
				t.errCh <- err
			}
		}
	}()

	for {
		select {
		case result := <-t.fiveTupleCh:
			// Wait for peer to get FiveTuple before returning.  If we didn't do
			// this, our natty instance might stop running before the peer
			// finishes its work to get its own FiveTuple.
			<-t.peerGotFiveTupleCh
			return result, nil
		case err := <-t.errCh:
			if err != nil && err != io.EOF {
				return nil, err
			}
		}
	}
}

// initCommand sets up the natty command
func (t *Traversal) initCommand(params []string) error {
	if t.debugOut == nil {
		// Discard stderr output by default
		t.debugOut = ioutil.Discard
	} else {
		// Tell natty to log debug output
		params = append(params, "-debug")
	}

	nattyBytes, err := Asset("natty")
	if err != nil {
		return err
	}
	be, err := byteexec.NewByteExec(nattyBytes)
	if err != nil {
		return err
	}

	t.cmd = be.Command(params...)
	t.stdin, err = t.cmd.StdinPipe()
	if err != nil {
		return err
	}
	t.stdout, err = t.cmd.StdoutPipe()
	if err != nil {
		return err
	}
	t.stderr, err = t.cmd.StderrPipe()
	if err != nil {
		return err
	}

	t.stdoutbuf = bufio.NewReader(t.stdout)

	return nil
}

// processStdout reads the output from natty and sends it to the msgOutCh. If
// it finds a FiveTuple, it records that.
func (t *Traversal) processStdout() {
	for {
		// Read next message from natty
		msg, err := t.stdoutbuf.ReadString('\n')
		if err != nil {
			t.errCh <- err
			return
		}

		// Request send of message to peer
		t.msgOutCh <- msg

		if IsFiveTuple(msg) {
			// We got a FiveTuple!
			fiveTuple := &FiveTuple{}
			err = json.Unmarshal([]byte(msg), fiveTuple)
			if err != nil {
				t.errCh <- err
				return
			}
			t.fiveTupleCh <- fiveTuple
		}
	}
}

// processStderr copies the output from natty's stderr to the configured
// DebugOut
func (t *Traversal) processStderr() {
	_, err := io.Copy(t.debugOut, t.stderr)
	t.errCh <- err
}

func (t *Traversal) closePipes() {
	if t.stdin != nil {
		t.stdin.Close()
	}
	if t.stdout != nil {
		t.stdout.Close()
	}
	if t.stderr != nil {
		t.stderr.Close()
	}
}

func IsFiveTuple(msg string) bool {
	return strings.Contains(msg, "\"type\":\"5-tuple\"")
}
<commit_msg>Pointing at byteexec in getlantern<commit_after>// Package natty provides a Go language wrapper to the natty NAT traversal
// utility.  See https://github.com/getlantern/t.
//
// See natty_test for an example of Natty in use, including debug logging
// showing the messages that are sent across the signaling channel.
package natty

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/getlantern/byteexec"
)

const (
	UDP = Protocol("udp")
	TCP = Protocol("tcp")
)

var (
	reallyHighTimeout = 100000 * time.Hour
)

type Protocol string

// A FiveTuple is the result of a successful NAT traversal.
type FiveTuple struct {
	Proto  Protocol
	Local  string
	Remote string
}

// UDPAddrs returns a pair of UDPAddrs representing the Local and Remote
// addresses of this FiveTuple. If the FiveTuple's Proto is not UDP, this method
// returns an error.
func (ft *FiveTuple) UDPAddrs() (local *net.UDPAddr, remote *net.UDPAddr, err error) {
	if ft.Proto != UDP {
		err = fmt.Errorf("FiveTuple.Proto was not UDP!: %s", ft.Proto)
		return
	}
	local, err = net.ResolveUDPAddr("udp", ft.Local)
	if err != nil {
		err = fmt.Errorf("Unable to resolve local UDP address %s: %s", ft.Local)
		return
	}
	remote, err = net.ResolveUDPAddr("udp", ft.Remote)
	if err != nil {
		err = fmt.Errorf("Unable to resolve remote UDP address %s: %s", ft.Remote)
	}
	return
}

// Traversal represents a single NAT traversal using natty, whose result is
// available via the methods FiveTuple() and FiveTupleTimeout().
//
// Consumers should make sure to call Close() after finishing with this Natty
// in order to make sure the underlying natty process and associated resources
// are closed.
type Traversal struct {
	debugOut           io.Writer       // target for output from natty's stderr
	cmd                *exec.Cmd       // the natty command
	stdin              io.WriteCloser  // pipe to natty's stdin
	stdout             io.ReadCloser   // pipe from natty's stdout
	stdoutbuf          *bufio.Reader   // buffered stdout
	stderr             io.ReadCloser   // pipe from natty's stderr
	msgInCh            chan string     // channel for messages inbound to this Natty
	msgOutCh           chan string     // channel for messages outbound from this Natty
	peerGotFiveTupleCh chan bool       // channel to signal once we know that our peer received their own FiveTuple
	fiveTupleCh        chan *FiveTuple // intermediary channel for the FiveTuple emitted by the natty command
	errCh              chan error      // intermediary channel for any error encountered while running natty
	fiveTupleOutCh     chan *FiveTuple // channel for FiveTuple output
	errOutCh           chan error      // channel for error output
	fiveTupleOut       *FiveTuple      // the output FiveTuple
	errOut             error           // the output error
	outMutex           sync.Mutex      // mutex for synchronizing access to output variables
}

// Offer starts a Traversal as an Offerer, meaning that it will make an offer to
// initiate an ICE session. Call FiveTuple() or FiveTupleTimeout() to get the
// FiveTuple resulting from Traversal.
//
// debugOut (optional) is an optional io.Writer to which debug output from natty
// will be written.
//
func Offer(debugOut io.Writer) *Traversal {
	t := &Traversal{
		debugOut: debugOut,
	}
	t.run([]string{"-offer"})
	return t
}

// Answer starts a Traversal as an Answerer, meaning that it will accept offers
// to initiate an ICE session. Call FiveTuple() or FiveTupleTimeout() to get the
// FiveTuple resulting from Traversal.
//
// debugOut (optional) is an optional io.Writer to which debug output from Natty
// will be written.
//
func Answer(debugOut io.Writer) *Traversal {
	t := &Traversal{
		debugOut: debugOut,
	}
	t.run([]string{})
	return t
}

// MsgIn is used to pass this Traversal a message from the peer t. This method
// is buffered and will typically not block.
func (t *Traversal) MsgIn(msg string) {
	t.msgInCh <- msg
}

// NextMsgOut gets the next message to pass to the peer.  If done is true, there
// are no more messages to be read, and the currently returned message should be
// ignored.
func (t *Traversal) NextMsgOut() (msg string, done bool) {
	m, ok := <-t.msgOutCh
	return m, !ok
}

// FiveTuple gets the FiveTuple from the Traversal, blocking until such is
// available.
func (t *Traversal) FiveTuple() (*FiveTuple, error) {
	return t.FiveTupleTimeout(reallyHighTimeout)
}

// FiveTupleTimeout gets the FiveTuple from the Traversal, blocking until such
// is available or the timeout is hit.
func (t *Traversal) FiveTupleTimeout(timeout time.Duration) (*FiveTuple, error) {
	t.outMutex.Lock()
	defer t.outMutex.Unlock()

	if t.fiveTupleOut == nil && t.errOut == nil {
		// We don't have a result yet, wait for one
		select {
		case ft := <-t.fiveTupleOutCh:
			t.fiveTupleOut = ft
		case err := <-t.errOutCh:
			t.errOut = err
		case <-time.After(timeout):
			// Return an error, but don't store it (lets caller try again)
			return nil, fmt.Errorf("Timed out waiting for five-tuple")
		}
	}

	return t.fiveTupleOut, t.errOut
}

// Close closes this Traversal, terminating any outstanding natty process by
// sending SIGKILL. Close blocks until the natty process has terminated, at
// which point any ports that it bound should be available for use.
func (t *Traversal) Close() error {
	t.closePipes()

	if t.cmd != nil && t.cmd.Process != nil {
		t.cmd.Process.Kill()
		t.cmd.Process.Wait()
	}
	return nil
}

// run runs the natty command to obtain a FiveTuple. The actual running of
// natty happens on a goroutine so that run itself doesn't block.
func (t *Traversal) run(params []string) {
	t.msgInCh = make(chan string, 100)
	t.msgOutCh = make(chan string, 100)

	// Note - these channels are a buffered in order to prevent deadlocks
	t.peerGotFiveTupleCh = make(chan bool, 10)
	t.fiveTupleCh = make(chan *FiveTuple, 10)
	t.errCh = make(chan error, 10)
	t.fiveTupleOutCh = make(chan *FiveTuple, 10)
	t.errOutCh = make(chan error, 10)

	err := t.initCommand(params)

	go func() {
		if err != nil {
			t.errOutCh <- err
			return
		}

		ft, err := t.doRun(params)
		// Once doRun is finished, inform client of the FiveTuple
		if err != nil {
			t.errOutCh <- err
		} else {
			t.fiveTupleOutCh <- ft
		}
	}()
}

// doRun does the running, including resource cleanup.  doRun blocks until
// Close() has finished, meaning that natty is no longer running and whatever
// port it returned in the FiveTuple can now be used for other things.
func (t *Traversal) doRun(params []string) (*FiveTuple, error) {
	defer t.Close()

	go t.processStdout()
	go t.processStderr()

	// Start the natty command
	t.errCh <- t.cmd.Start()

	// Process incoming messages
	go func() {
		for {
			msg := <-t.msgInCh

			if IsFiveTuple(msg) {
				t.peerGotFiveTupleCh <- true
				continue
			}

			// Forward message to natty process
			_, err := t.stdin.Write([]byte(msg))
			if err == nil {
				_, err = t.stdin.Write([]byte("\n"))
			}
			if err != nil {
				t.errCh <- err
			}
		}
	}()

	for {
		select {
		case result := <-t.fiveTupleCh:
			// Wait for peer to get FiveTuple before returning.  If we didn't do
			// this, our natty instance might stop running before the peer
			// finishes its work to get its own FiveTuple.
			<-t.peerGotFiveTupleCh
			return result, nil
		case err := <-t.errCh:
			if err != nil && err != io.EOF {
				return nil, err
			}
		}
	}
}

// initCommand sets up the natty command
func (t *Traversal) initCommand(params []string) error {
	if t.debugOut == nil {
		// Discard stderr output by default
		t.debugOut = ioutil.Discard
	} else {
		// Tell natty to log debug output
		params = append(params, "-debug")
	}

	nattyBytes, err := Asset("natty")
	if err != nil {
		return err
	}
	be, err := byteexec.NewByteExec(nattyBytes)
	if err != nil {
		return err
	}

	t.cmd = be.Command(params...)
	t.stdin, err = t.cmd.StdinPipe()
	if err != nil {
		return err
	}
	t.stdout, err = t.cmd.StdoutPipe()
	if err != nil {
		return err
	}
	t.stderr, err = t.cmd.StderrPipe()
	if err != nil {
		return err
	}

	t.stdoutbuf = bufio.NewReader(t.stdout)

	return nil
}

// processStdout reads the output from natty and sends it to the msgOutCh. If
// it finds a FiveTuple, it records that.
func (t *Traversal) processStdout() {
	for {
		// Read next message from natty
		msg, err := t.stdoutbuf.ReadString('\n')
		if err != nil {
			t.errCh <- err
			return
		}

		// Request send of message to peer
		t.msgOutCh <- msg

		if IsFiveTuple(msg) {
			// We got a FiveTuple!
			fiveTuple := &FiveTuple{}
			err = json.Unmarshal([]byte(msg), fiveTuple)
			if err != nil {
				t.errCh <- err
				return
			}
			t.fiveTupleCh <- fiveTuple
		}
	}
}

// processStderr copies the output from natty's stderr to the configured
// DebugOut
func (t *Traversal) processStderr() {
	_, err := io.Copy(t.debugOut, t.stderr)
	t.errCh <- err
}

func (t *Traversal) closePipes() {
	if t.stdin != nil {
		t.stdin.Close()
	}
	if t.stdout != nil {
		t.stdout.Close()
	}
	if t.stderr != nil {
		t.stderr.Close()
	}
}

func IsFiveTuple(msg string) bool {
	return strings.Contains(msg, "\"type\":\"5-tuple\"")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

var Host = "127.0.0.1"
var Port = ":9991"
var Input = "Input from other side, пока, £, 语汉"

// ReadCloser is used to wrap strings.Reader for implementing io.ReadCloser interface
type ReadCloser struct {
	r *strings.Reader
}

func (rc ReadCloser) Read(p []byte) (n int, err error) {
	return rc.r.Read(p)
}

func (rc ReadCloser) Close() error {
	return nil
}

func TestTCP(t *testing.T) {
	// Capture STDIN and STDOUT streams
	oldStdin := os.Stdin
	oldStdout := os.Stdout

	// Bytes written to w1 are read from os.Stdin
	r1, w1, e := os.Pipe()
	assert.Nil(t, e)
	os.Stdin = r1

	// Bytes written os.Stdout are read from r2
	// r2, w2, e := os.Pipe()
	// assert.Nil(t, e)
	// os.Stdout = w2

	// Send test data to listener from goroutine and wait for potentials errors at the end of the test
	go func() {
		// Wait for main thread starts listener
		time.Sleep(200 * time.Millisecond)
		con, err := net.Dial("tcp", Host+Port)
		assert.Nil(t, err)

		// Transfer data
		_, err = w1.Write([]byte(Input))
		assert.Nil(t, err)
		TransferStreams(con)

		// Wait for data will be transferred
		time.Sleep(200 * time.Millisecond)
	}()

	ln, err := net.Listen("tcp", Port)
	assert.Nil(t, err)

	con, err := ln.Accept()
	assert.Nil(t, err)

	buf := make([]byte, 1024)
	n, err := con.Read(buf)
	assert.Nil(t, err)

	// Restore STDIN and STDOUT
	os.Stdin = oldStdin
	os.Stdout = oldStdout

	assert.Equal(t, Input, string(buf[0:n]))
}

func TestUDP(t *testing.T) {
	// Send test data to listener from goroutine and wait for potentials errors at the end of the test
	go func() {
		// Wait for main thread starts listener
		time.Sleep(200 * time.Millisecond)
		con, err := net.Dial("udp", Host+Port)
		assert.Nil(t, err)

		// Transfer data
		addr, err := net.ResolveUDPAddr("udp", Host+Port)
		assert.Nil(t, err)
		fmt.Println(con.RemoteAddr())
		c1 := copyPackets(ReadCloser{r: strings.NewReader(Input)}, con, addr)

		// Wait for data will be transferred
		time.Sleep(200 * time.Millisecond)
		select {
		case progress := <-c1:
			t.Logf("Remote connection is closed: %+v\n", progress)
		default:
			t.Fatal("handle() must write to result channel")
		}
	}()

	con, err := net.ListenPacket("udp", Port)
	assert.Nil(t, err)

	buf := make([]byte, 1024)
	n, _, err := con.ReadFrom(buf)
	assert.Nil(t, err)

	assert.Equal(t, Input, string(buf[0:n]))
}
<commit_msg>Call TransferPackets from test<commit_after>package main

import (
	"net"
	"os"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

var Host = "127.0.0.1"
var Port = ":9991"
var Input = "Input from other side, пока, £, 语汉"

func TestTransferStreams(t *testing.T) {
	oldStdin := os.Stdin

	// Bytes written to w1 are read from os.Stdin
	r, w, e := os.Pipe()
	assert.Nil(t, e)
	os.Stdin = r

	// Send data to server from goroutine and wait for potentials errors at the end of the test
	go func() {
		// Wait for main thread starts server
		time.Sleep(200 * time.Millisecond)
		con, err := net.Dial("tcp", Host+Port)
		assert.Nil(t, err)

		// Client sends data
		_, err = w.Write([]byte(Input))
		assert.Nil(t, err)
		TransferStreams(con)
	}()

	// Server receives data
	ln, err := net.Listen("tcp", Port)
	assert.Nil(t, err)

	con, err := ln.Accept()
	assert.Nil(t, err)

	buf := make([]byte, 1024)
	n, err := con.Read(buf)
	assert.Nil(t, err)

	assert.Equal(t, Input, string(buf[0:n]))

	os.Stdin = oldStdin
}

func TestTransferPackets(t *testing.T) {
	oldStdin := os.Stdin

	// Bytes written to w1 are read from os.Stdin
	r, w, e := os.Pipe()
	assert.Nil(t, e)
	os.Stdin = r

	// Send test data to server from goroutine and wait for potentials errors at the end of the test
	go func() {
		// Wait for main thread starts server
		time.Sleep(200 * time.Millisecond)
		con, err := net.Dial("udp", Host+Port)
		assert.Nil(t, err)

		// Client sends data
		_, err = w.Write([]byte(Input))
		assert.Nil(t, err)
		TransferStreams(con)
	}()

	con, err := net.ListenPacket("udp", Port)
	assert.Nil(t, err)

	buf := make([]byte, 1024)
	n, _, err := con.ReadFrom(buf)
	assert.Nil(t, err)

	assert.Equal(t, Input, string(buf[0:n]))

	os.Stdin = oldStdin
}
<|endoftext|>