<commit_before>package sqlsess

import (
	"database/sql"
	"net/http"
	"time"

	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"

	"github.com/BurntSushi/locker"
)

var (
	SessionLastUpdated = "__sess_last_updated"
	CookieIdName       = "sess_sessionid"
	SqlTableName       = "sess_session"
	SqlCreateSession   = `
	CREATE TABLE IF NOT EXISTS ` + SqlTableName + ` (
		id BYTEA NOT NULL,
		name VARCHAR (255) NOT NULL,
		key TEXT NOT NULL,
		value TEXT NOT NULL,
		PRIMARY KEY (id, name, key)
	)
	`
)

type Store struct {
	*sql.DB
	hashKey, blockKey []byte
}

func Open(db *sql.DB) (*Store, error) {
	if _, err := db.Exec(SqlCreateSession); err != nil {
		return nil, err
	}

	s := &Store{
		DB:       db,
		hashKey:  securecookie.GenerateRandomKey(64),
		blockKey: securecookie.GenerateRandomKey(32),
	}
	return s, nil
}

func (s *Store) Clean(inactive time.Duration) error {
	locker.Lock("clean")
	defer locker.Unlock("clean")

	rows, err := s.Query(`
		SELECT id, value
		FROM ` + SqlTableName + `
		WHERE key = $1
	`, SessionLastUpdated)
	if err != nil {
		return err
	}
	defer rows.Close()

	cutoff := time.Now().UTC().Add(-inactive)
	for rows.Next() {
		var id []byte
		var last string
		if err := rows.Scan(&id, &last); err != nil {
			return err
		}
		lastUp, err := time.Parse(time.RFC3339Nano, last)
		if err != nil {
			return err
		}
		if lastUp.Before(cutoff) {
			_, err = s.Exec(
				"DELETE FROM " + SqlTableName + " WHERE id = $1", id)
			if err != nil {
				return err
			}
		}
	}
	return rows.Err()
}

func (s *Store) Delete(sess *sessions.Session) error {
	id := []byte(sess.ID)
	_, err := s.Exec("DELETE FROM " + SqlTableName + " WHERE id = $1", id)
	return err
}

func (s *Store) Get(r *http.Request, name string) (*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(s, name)
}

func (s *Store) New(r *http.Request, name string) (*sessions.Session, error) {
	sess := sessions.NewSession(s, name)
	sess.ID = s.id(r)

	RLock(sess)
	defer RUnlock(sess)

	rows, err := s.Query(`
		SELECT key, value
		FROM ` + SqlTableName + `
		WHERE id = $1 AND name = $2
	`, []byte(sess.ID), name)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var k, v string
		if err := rows.Scan(&k, &v); err != nil {
			return nil, err
		}
		sess.Values[k] = v
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return sess, nil
}

func (s *Store) Save(
	r *http.Request,
	w http.ResponseWriter,
	sess *sessions.Session,
) error {
	Lock(sess)
	defer Unlock(sess)

	s.writeCookie(r, w, CookieIdName, sess.ID)

	id := []byte(sess.ID)
	sess.Values[SessionLastUpdated] = time.Now().UTC()
	tx, err := s.Begin()
	if err != nil {
		return err
	}

	_, err = tx.Exec("DELETE FROM " + SqlTableName + " WHERE id = $1", id)
	if err != nil {
		tx.Rollback()
		return err
	}

	prep, err := tx.Prepare(`
		INSERT INTO ` + SqlTableName + `
			(id, name, key, value) VALUES ($1, $2, $3, $4)
	`)
	if err != nil {
		tx.Rollback()
		return err
	}
	defer prep.Close()

	for k, v := range sess.Values {
		if _, err := prep.Exec(id, sess.Name(), k, v); err != nil {
			tx.Rollback()
			return err
		}
	}
	return tx.Commit()
}

// id returns either the session id from a user's cookie or generates
// a fresh one if the cookie is inaccessible or missing.
func (s *Store) id(r *http.Request) string {
	id := s.readCookie(r, CookieIdName)
	if len(id) == 0 {
		id = string(securecookie.GenerateRandomKey(64))
	}
	return id
}

// SetKeys sets the hash and block keys used to read and write the session
// cookie. A hash key is required and is used to to authenticate a cookie
// value using HMAC. It's recommend to be 32 or 64 bytes.
//
// A block key is optional and is used to encrypt the cookie value. If it's
// set to nil, then encryption will not be used. This package uses AES, so
// the block key must have length 16, 24 or 32 bytes corresponding to
// AES-128, AES-192 or AES-256. If the block key violates these constraints,
// SetKeys will panic.
//
// This method is exposed so that multiple instantiations of session stores
// can share the same cookie. This particularly useful if you want to be able
// to restart your web server without invalidating existing user sessions.
//
// If this method is not called, then a fresh set of keys is created
// automatically, but will invalidate all existing user sessions.
func (s *Store) SetKeys(hash, block []byte) {
	validLen := len(block) == 16 || len(block) == 24 || len(block) == 32
	if block != nil && !validLen {
		panic("invalid block key")
	}
	s.hashKey, s.blockKey = hash, block
}

func (s *Store) cookrw() *securecookie.SecureCookie {
	return securecookie.New(s.hashKey, s.blockKey)
}

// Returns an empty string if the cookie doesn't exist or if there was
// a problem decoding it.
func (s *Store) readCookie(r *http.Request, cname string) string {
	if cook, err := r.Cookie(cname); err == nil {
		var v string
		if err = s.cookrw().Decode(cname, cook.Value, &v); err == nil {
			return v
		}
	}
	return ""
}

// Writes the value to the named cookie with encryption.
func (s *Store) writeCookie(
	r *http.Request,
	w http.ResponseWriter,
	cname, cvalue string,
) {
	if encoded, err := s.cookrw().Encode(cname, cvalue); err == nil {
		cook := &http.Cookie{
			Name:     cname,
			Value:    encoded,
			Path:     "/",
			HttpOnly: true,
		}
		http.SetCookie(w, cook)
	}
}

func Lock(sess *sessions.Session)    { locker.Lock(sess.ID) }
func Unlock(sess *sessions.Session)  { locker.Unlock(sess.ID) }
func RLock(sess *sessions.Session)   { locker.RLock(sess.ID) }
func RUnlock(sess *sessions.Session) { locker.RUnlock(sess.ID) }
<commit_msg>gofmt<commit_after>package sqlsess

import (
	"database/sql"
	"net/http"
	"time"

	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"

	"github.com/BurntSushi/locker"
)

var (
	SessionLastUpdated = "__sess_last_updated"
	CookieIdName       = "sess_sessionid"
	SqlTableName       = "sess_session"
	SqlCreateSession   = `
	CREATE TABLE IF NOT EXISTS ` + SqlTableName + ` (
		id BYTEA NOT NULL,
		name VARCHAR (255) NOT NULL,
		key TEXT NOT NULL,
		value TEXT NOT NULL,
		PRIMARY KEY (id, name, key)
	)
	`
)

type Store struct {
	*sql.DB
	hashKey, blockKey []byte
}

func Open(db *sql.DB) (*Store, error) {
	if _, err := db.Exec(SqlCreateSession); err != nil {
		return nil, err
	}

	s := &Store{
		DB:       db,
		hashKey:  securecookie.GenerateRandomKey(64),
		blockKey: securecookie.GenerateRandomKey(32),
	}
	return s, nil
}

func (s *Store) Clean(inactive time.Duration) error {
	locker.Lock("clean")
	defer locker.Unlock("clean")

	rows, err := s.Query(`
		SELECT id, value
		FROM `+SqlTableName+`
		WHERE key = $1
	`, SessionLastUpdated)
	if err != nil {
		return err
	}
	defer rows.Close()

	cutoff := time.Now().UTC().Add(-inactive)
	for rows.Next() {
		var id []byte
		var last string
		if err := rows.Scan(&id, &last); err != nil {
			return err
		}
		lastUp, err := time.Parse(time.RFC3339Nano, last)
		if err != nil {
			return err
		}
		if lastUp.Before(cutoff) {
			_, err = s.Exec(
				"DELETE FROM "+SqlTableName+" WHERE id = $1", id)
			if err != nil {
				return err
			}
		}
	}
	return rows.Err()
}

func (s *Store) Delete(sess *sessions.Session) error {
	id := []byte(sess.ID)
	_, err := s.Exec("DELETE FROM "+SqlTableName+" WHERE id = $1", id)
	return err
}

func (s *Store) Get(r *http.Request, name string) (*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(s, name)
}

func (s *Store) New(r *http.Request, name string) (*sessions.Session, error) {
	sess := sessions.NewSession(s, name)
	sess.ID = s.id(r)

	RLock(sess)
	defer RUnlock(sess)

	rows, err := s.Query(`
		SELECT key, value
		FROM `+SqlTableName+`
		WHERE id = $1 AND name = $2
	`, []byte(sess.ID), name)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var k, v string
		if err := rows.Scan(&k, &v); err != nil {
			return nil, err
		}
		sess.Values[k] = v
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return sess, nil
}

func (s *Store) Save(
	r *http.Request,
	w http.ResponseWriter,
	sess *sessions.Session,
) error {
	Lock(sess)
	defer Unlock(sess)

	s.writeCookie(r, w, CookieIdName, sess.ID)

	id := []byte(sess.ID)
	sess.Values[SessionLastUpdated] = time.Now().UTC()
	tx, err := s.Begin()
	if err != nil {
		return err
	}

	_, err = tx.Exec("DELETE FROM "+SqlTableName+" WHERE id = $1", id)
	if err != nil {
		tx.Rollback()
		return err
	}

	prep, err := tx.Prepare(`
		INSERT INTO ` + SqlTableName + `
			(id, name, key, value) VALUES ($1, $2, $3, $4)
	`)
	if err != nil {
		tx.Rollback()
		return err
	}
	defer prep.Close()

	for k, v := range sess.Values {
		if _, err := prep.Exec(id, sess.Name(), k, v); err != nil {
			tx.Rollback()
			return err
		}
	}
	return tx.Commit()
}

// id returns either the session id from a user's cookie or generates
// a fresh one if the cookie is inaccessible or missing.
func (s *Store) id(r *http.Request) string {
	id := s.readCookie(r, CookieIdName)
	if len(id) == 0 {
		id = string(securecookie.GenerateRandomKey(64))
	}
	return id
}

// SetKeys sets the hash and block keys used to read and write the session
// cookie. A hash key is required and is used to to authenticate a cookie
// value using HMAC. It's recommend to be 32 or 64 bytes.
//
// A block key is optional and is used to encrypt the cookie value. If it's
// set to nil, then encryption will not be used. This package uses AES, so
// the block key must have length 16, 24 or 32 bytes corresponding to
// AES-128, AES-192 or AES-256. If the block key violates these constraints,
// SetKeys will panic.
//
// This method is exposed so that multiple instantiations of session stores
// can share the same cookie. This particularly useful if you want to be able
// to restart your web server without invalidating existing user sessions.
//
// If this method is not called, then a fresh set of keys is created
// automatically, but will invalidate all existing user sessions.
func (s *Store) SetKeys(hash, block []byte) {
	validLen := len(block) == 16 || len(block) == 24 || len(block) == 32
	if block != nil && !validLen {
		panic("invalid block key")
	}
	s.hashKey, s.blockKey = hash, block
}

func (s *Store) cookrw() *securecookie.SecureCookie {
	return securecookie.New(s.hashKey, s.blockKey)
}

// Returns an empty string if the cookie doesn't exist or if there was
// a problem decoding it.
func (s *Store) readCookie(r *http.Request, cname string) string {
	if cook, err := r.Cookie(cname); err == nil {
		var v string
		if err = s.cookrw().Decode(cname, cook.Value, &v); err == nil {
			return v
		}
	}
	return ""
}

// Writes the value to the named cookie with encryption.
func (s *Store) writeCookie(
	r *http.Request,
	w http.ResponseWriter,
	cname, cvalue string,
) {
	if encoded, err := s.cookrw().Encode(cname, cvalue); err == nil {
		cook := &http.Cookie{
			Name:     cname,
			Value:    encoded,
			Path:     "/",
			HttpOnly: true,
		}
		http.SetCookie(w, cook)
	}
}

func Lock(sess *sessions.Session)    { locker.Lock(sess.ID) }
func Unlock(sess *sessions.Session)  { locker.Unlock(sess.ID) }
func RLock(sess *sessions.Session)   { locker.RLock(sess.ID) }
func RUnlock(sess *sessions.Session) { locker.RUnlock(sess.ID) }
<|endoftext|>


<commit_before>package run

import (
	"bufio"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

func performRequest(url string) (bool, time.Duration) {
	t0 := time.Now()
	res, err := http.Get(url)
	t1 := time.Now()
	if err == nil {
		res.Body.Close()
	}

	return err == nil && res.StatusCode == 200, t1.Sub(t0)
}

func sleepUntil(t time.Time, c chan time.Duration) {
	tDone := t.Sub(time.Now())
	c <- -tDone

	for tDone > 0 {
		var sleepDuration time.Duration
		if tDone < 10*time.Millisecond {
			sleepDuration = tDone
		} else {
			sleepDuration = tDone / 2
		}

		time.Sleep(sleepDuration)
		tDone = t.Sub(time.Now())
		c <- -tDone
	}
}

func printStatus(sent int, done int, successful int, lag time.Duration) {
	nanoLag := float64(lag.Nanoseconds()) / float64(1000000)
	fmt.Printf("\rsent: %8d, done: %8d, waiting: %8d, successful: %8d (%3.2f%%), lag: %8.3fms", sent, done, sent-done, successful, 100*float64(successful)/float64(done), nanoLag)
}

func Run(prefix *string, flood *bool, speedup *float64) {
	var wg sync.WaitGroup

	t0 := time.Now()
	var tCorrection float64

	first := true

	chanSent := make(chan int)
	chanDone := make(chan int)
	chanSuccess := make(chan int)
	chanLag := make(chan time.Duration)

	go func() {
		sent := 0
		done := 0
		success := 0
		v := 0
		var lag time.Duration

		for {
			select {
			case v = <-chanSent:
				sent += v
			case v = <-chanDone:
				done += v
			case v = <-chanSuccess:
				success += v
			case lag = <-chanLag:
			}

			printStatus(sent, done, success, lag)
		}
	}()

	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line := scanner.Text()

		lineParts := strings.SplitN(line, " ", 2)
		delay, _ := strconv.ParseFloat(lineParts[0], 64)
		url := *prefix + lineParts[1]

		if first {
			tCorrection = delay
			first = false
		}

		if !*flood {
			runAt := t0.Add(time.Duration((delay - tCorrection) / *speedup * 1000000000))
			sleepUntil(runAt, chanLag)
		}

		wg.Add(1)
		go func() {
			defer wg.Done()

			chanSent <- 1
			var success, _ = performRequest(url)
			chanDone <- 1
			if success {
				chanSuccess <- 1
			}
		}()
	}

	wg.Wait()
	fmt.Println()
}
<commit_msg>split statusline into two<commit_after>package run

import (
	"bufio"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

func performRequest(url string) (bool, time.Duration) {
	t0 := time.Now()
	res, err := http.Get(url)
	t1 := time.Now()
	if err == nil {
		res.Body.Close()
	}

	return err == nil && res.StatusCode == 200, t1.Sub(t0)
}

func sleepUntil(t time.Time, c chan time.Duration) {
	tDone := t.Sub(time.Now())
	c <- -tDone

	for tDone > 0 {
		var sleepDuration time.Duration
		if tDone < 10*time.Millisecond {
			sleepDuration = tDone
		} else {
			sleepDuration = tDone / 2
		}

		time.Sleep(sleepDuration)
		tDone = t.Sub(time.Now())
		c <- -tDone
	}
}

func printStatusHeader() {
	fmt.Println("lag [ms]       sent      done   waiting        successful")
}

func printStatus(sent int, done int, successful int, lag time.Duration) {
	nanoLag := float64(lag.Nanoseconds()) / float64(1000000)
	fmt.Printf("\r%8.3f   %8d  %8d  %8d  %8d %3.2f%%", nanoLag, sent, done, sent-done, successful, 100*float64(successful)/float64(done))
}

func Run(prefix *string, flood *bool, speedup *float64) {
	var wg sync.WaitGroup

	t0 := time.Now()
	var tCorrection float64

	first := true

	chanSent := make(chan int)
	chanDone := make(chan int)
	chanSuccess := make(chan int)
	chanLag := make(chan time.Duration)

	printStatusHeader()

	go func() {
		sent := 0
		done := 0
		success := 0
		v := 0
		var lag time.Duration

		for {
			select {
			case v = <-chanSent:
				sent += v
			case v = <-chanDone:
				done += v
			case v = <-chanSuccess:
				success += v
			case lag = <-chanLag:
			}

			printStatus(sent, done, success, lag)
		}
	}()

	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line := scanner.Text()

		lineParts := strings.SplitN(line, " ", 2)
		delay, _ := strconv.ParseFloat(lineParts[0], 64)
		url := *prefix + lineParts[1]

		if first {
			tCorrection = delay
			first = false
		}

		if !*flood {
			runAt := t0.Add(time.Duration((delay - tCorrection) / *speedup * 1000000000))
			sleepUntil(runAt, chanLag)
		}

		wg.Add(1)
		go func() {
			defer wg.Done()

			chanSent <- 1
			var success, _ = performRequest(url)
			chanDone <- 1
			if success {
				chanSuccess <- 1
			}
		}()
	}

	wg.Wait()
	fmt.Println()
}
<|endoftext|>


<commit_before>// Package run is the entry point of elvish.
package run

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"os/signal"
	"os/user"
	"syscall"
	"time"

	"github.com/elves/elvish/edit"
	"github.com/elves/elvish/eval"
	"github.com/elves/elvish/parse"
	"github.com/elves/elvish/store"
	"github.com/elves/elvish/stub"
	"github.com/elves/elvish/sys"
	"github.com/elves/elvish/util"
)

var Logger = util.GetLogger("[main] ")

var (
	log    = flag.String("log", "", "a file to write debug log to")
	dbname = flag.String("db", "", "path to the database")
	help   = flag.Bool("help", false, "show usage help and quit")
)

func usage() {
	fmt.Println("usage: elvish [flags] [script]")
	fmt.Println("flags:")
	flag.PrintDefaults()
}

// Main is the entry point of elvish.
func Main() {
	defer rescue()

	flag.Usage = usage
	flag.Parse()
	args := flag.Args()

	if len(args) > 1 {
		usage()
		os.Exit(2)
	}

	if *help {
		usage()
		os.Exit(0)
	}

	if *log != "" {
		err := util.SetOutputFile(*log)
		if err != nil {
			fmt.Println(err)
		}
	}

	go handleHupAndQuit()
	go logSignals()

	ev, st := newEvalerAndStore()
	defer func() {
		err := st.Close()
		if err != nil {
			fmt.Println("failed to close database:", err)
		}
	}()

	stub, err := stub.NewStub(os.Stderr)
	if err != nil {
		fmt.Println("failed to spawn stub:", err)
	} else {
		ev.Stub = stub
	}

	if len(args) == 1 {
		script(ev, args[0])
	} else if !sys.IsATTY(0) {
		script(ev, "/dev/stdin")
	} else {
		interact(ev, st)
	}
}

func rescue() {
	r := recover()
	if r != nil {
		println()
		fmt.Println(r)
		print(sys.DumpStack())
		println("\nexecing recovery shell /bin/sh")
		syscall.Exec("/bin/sh", []string{"/bin/sh"}, os.Environ())
	}
}

func script(ev *eval.Evaler, fname string) {
	err := ev.Source(fname)
	if err != nil {
		printError(err)
		os.Exit(1)
	}
}

func interact(ev *eval.Evaler, st *store.Store) {
	// Build Editor.
	sigch := make(chan os.Signal)
	signal.Notify(sigch)
	ed := edit.NewEditor(os.Stdin, sigch, ev, st)

	// Source rc.elv.
	datadir, err := store.EnsureDataDir()
	printError(err)
	if err == nil {
		// XXX
		err := ev.Source(datadir + "/rc.elv")
		if err != nil && !os.IsNotExist(err) {
			printError(err)
		}
	}

	// Build prompt and rprompt.
	username := "???"
	user, err := user.Current()
	if err == nil {
		username = user.Username
	}
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "???"
	}
	rpromptStr := username + "@" + hostname
	prompt := func() string {
		return util.Getwd() + "> "
	}
	rprompt := func() string {
		return rpromptStr
	}

	// Build readLine function.
	readLine := func() edit.LineRead {
		return ed.ReadLine(prompt, rprompt)
	}

	usingBasic := false
	cmdNum := 0

	for {
		cmdNum++
		// name := fmt.Sprintf("<tty %d>", cmdNum)

		lr := readLine()

		if lr.EOF {
			break
		} else if lr.Err != nil {
			fmt.Println("Editor error:", lr.Err)
			if !usingBasic {
				fmt.Println("Falling back to basic line editor")
				readLine = basicReadLine
				usingBasic = true
			} else {
				fmt.Println("Don't know what to do, pid is", os.Getpid())
				fmt.Println("Restarting editor in 1s")
				time.Sleep(time.Second)
			}
			continue
		}

		n, err := parse.Parse(lr.Line)
		printError(err)

		if err == nil {
			err := ev.EvalInteractive(lr.Line, n)
			printError(err)
		}
	}
}

func basicReadLine() edit.LineRead {
	stdin := bufio.NewReaderSize(os.Stdin, 0)
	line, err := stdin.ReadString('\n')
	if err == nil {
		return edit.LineRead{Line: line}
	} else if err == io.EOF {
		return edit.LineRead{EOF: true}
	} else {
		return edit.LineRead{Err: err}
	}
}

func logSignals() {
	sigs := make(chan os.Signal)
	signal.Notify(sigs)
	for sig := range sigs {
		Logger.Println("signal", sig)
	}
}

func handleHupAndQuit() {
	sigs := make(chan os.Signal)
	signal.Notify(sigs, syscall.SIGHUP, syscall.SIGQUIT)
	for sig := range sigs {
		fmt.Print(sys.DumpStack())
		if sig == syscall.SIGQUIT {
			os.Exit(3)
		}
	}
}

func newEvalerAndStore() (*eval.Evaler, *store.Store) {
	dataDir, err := store.EnsureDataDir()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Warning: cannot create data dir ~/.elvish")
	}

	var st *store.Store
	if err == nil {
		db := *dbname
		if db == "" {
			db = dataDir + "/db"
		}
		st, err = store.NewStore(db)
		if err != nil {
			fmt.Fprintln(os.Stderr, "Warning: cannot connect to store:", err)
		}
	}

	return eval.NewEvaler(st), st
}

func printError(err error) {
	if err == nil {
		return
	}
	switch err := err.(type) {
	case *util.Errors:
		for _, e := range err.Errors {
			printError(e)
		}
	default:
		eval.PprintError(err)
		fmt.Println()
	}
}
<commit_msg>Use a cooldown in face of editor error.<commit_after>// Package run is the entry point of elvish.
package run

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"os/signal"
	"os/user"
	"syscall"
	"time"

	"github.com/elves/elvish/edit"
	"github.com/elves/elvish/eval"
	"github.com/elves/elvish/parse"
	"github.com/elves/elvish/store"
	"github.com/elves/elvish/stub"
	"github.com/elves/elvish/sys"
	"github.com/elves/elvish/util"
)

var Logger = util.GetLogger("[main] ")

var (
	log    = flag.String("log", "", "a file to write debug log to")
	dbname = flag.String("db", "", "path to the database")
	help   = flag.Bool("help", false, "show usage help and quit")
)

func usage() {
	fmt.Println("usage: elvish [flags] [script]")
	fmt.Println("flags:")
	flag.PrintDefaults()
}

// Main is the entry point of elvish.
func Main() {
	defer rescue()

	flag.Usage = usage
	flag.Parse()
	args := flag.Args()

	if len(args) > 1 {
		usage()
		os.Exit(2)
	}

	if *help {
		usage()
		os.Exit(0)
	}

	if *log != "" {
		err := util.SetOutputFile(*log)
		if err != nil {
			fmt.Println(err)
		}
	}

	go handleHupAndQuit()
	go logSignals()

	ev, st := newEvalerAndStore()
	defer func() {
		err := st.Close()
		if err != nil {
			fmt.Println("failed to close database:", err)
		}
	}()

	stub, err := stub.NewStub(os.Stderr)
	if err != nil {
		fmt.Println("failed to spawn stub:", err)
	} else {
		ev.Stub = stub
	}

	if len(args) == 1 {
		script(ev, args[0])
	} else if !sys.IsATTY(0) {
		script(ev, "/dev/stdin")
	} else {
		interact(ev, st)
	}
}

func rescue() {
	r := recover()
	if r != nil {
		println()
		fmt.Println(r)
		print(sys.DumpStack())
		println("\nexecing recovery shell /bin/sh")
		syscall.Exec("/bin/sh", []string{"/bin/sh"}, os.Environ())
	}
}

func script(ev *eval.Evaler, fname string) {
	err := ev.Source(fname)
	if err != nil {
		printError(err)
		os.Exit(1)
	}
}

func interact(ev *eval.Evaler, st *store.Store) {
	// Build Editor.
	sigch := make(chan os.Signal)
	signal.Notify(sigch)
	ed := edit.NewEditor(os.Stdin, sigch, ev, st)

	// Source rc.elv.
	datadir, err := store.EnsureDataDir()
	printError(err)
	if err == nil {
		// XXX
		err := ev.Source(datadir + "/rc.elv")
		if err != nil && !os.IsNotExist(err) {
			printError(err)
		}
	}

	// Build prompt and rprompt.
	username := "???"
	user, err := user.Current()
	if err == nil {
		username = user.Username
	}
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "???"
	}
	rpromptStr := username + "@" + hostname
	prompt := func() string {
		return util.Getwd() + "> "
	}
	rprompt := func() string {
		return rpromptStr
	}

	// Build readLine function.
	readLine := func() edit.LineRead {
		return ed.ReadLine(prompt, rprompt)
	}

	cooldown := time.Second
	usingBasic := false
	cmdNum := 0

	for {
		cmdNum++
		// name := fmt.Sprintf("<tty %d>", cmdNum)

		lr := readLine()

		if lr.EOF {
			break
		} else if lr.Err != nil {
			fmt.Println("Editor error:", lr.Err)
			if !usingBasic {
				fmt.Println("Falling back to basic line editor")
				readLine = basicReadLine
				usingBasic = true
			} else {
				fmt.Println("Don't know what to do, pid is", os.Getpid())
				fmt.Println("Restarting editor in", cooldown)
				time.Sleep(cooldown)
				if cooldown < time.Minute {
					cooldown *= 2
				}
			}
			continue
		}

		// No error; reset cooldown.
		cooldown = time.Second

		n, err := parse.Parse(lr.Line)
		printError(err)

		if err == nil {
			err := ev.EvalInteractive(lr.Line, n)
			printError(err)
		}
	}
}

func basicReadLine() edit.LineRead {
	stdin := bufio.NewReaderSize(os.Stdin, 0)
	line, err := stdin.ReadString('\n')
	if err == nil {
		return edit.LineRead{Line: line}
	} else if err == io.EOF {
		return edit.LineRead{EOF: true}
	} else {
		return edit.LineRead{Err: err}
	}
}

func logSignals() {
	sigs := make(chan os.Signal)
	signal.Notify(sigs)
	for sig := range sigs {
		Logger.Println("signal", sig)
	}
}

func handleHupAndQuit() {
	sigs := make(chan os.Signal)
	signal.Notify(sigs, syscall.SIGHUP, syscall.SIGQUIT)
	for sig := range sigs {
		fmt.Print(sys.DumpStack())
		if sig == syscall.SIGQUIT {
			os.Exit(3)
		}
	}
}

func newEvalerAndStore() (*eval.Evaler, *store.Store) {
	dataDir, err := store.EnsureDataDir()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Warning: cannot create data dir ~/.elvish")
	}

	var st *store.Store
	if err == nil {
		db := *dbname
		if db == "" {
			db = dataDir + "/db"
		}
		st, err = store.NewStore(db)
		if err != nil {
			fmt.Fprintln(os.Stderr, "Warning: cannot connect to store:", err)
		}
	}

	return eval.NewEvaler(st), st
}

func printError(err error) {
	if err == nil {
		return
	}
	switch err := err.(type) {
	case *util.Errors:
		for _, e := range err.Errors {
			printError(e)
		}
	default:
		eval.PprintError(err)
		fmt.Println()
	}
}
<|endoftext|>


<commit_before>package tpl_test

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

	"github.com/joefitzgerald/inductor/tpl"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Tpl", func() {
	var (
		err       error
		osName    string
		tmpDir    string
		templates *tpl.Templates
	)
	BeforeEach(func() {
		tmpDir, err = ioutil.TempDir("", "inductor")
		Expect(err).NotTo(HaveOccurred())
		createTemplateFile(tmpDir, "nano/Autounattend.xml.disks.ptpl")
		createTemplateFile(tmpDir, "nano/packer.json.provisioners.ptpl")
		createTemplateFile(tmpDir, "windowsxp/Autounattend.xml.tpl")
		createTemplateFile(tmpDir, "scripts/win-updates.ps1")
		createTemplateFile(tmpDir, "scripts/nano/SetupComplete.cmd")
		createTemplateFile(tmpDir, "inductor.json")
		createTemplateFile(tmpDir, "README.md")
		createTemplateFile(tmpDir, "Autounattend.xml.tpl")
		createTemplateFile(tmpDir, "Autounattend.xml.oobe.ptpl")
		createTemplateFile(tmpDir, "Autounattend.xml.disks.ptpl")
		createTemplateFile(tmpDir, "packer.json.tpl")
		createTemplateFile(tmpDir, "packer.json.builders.ptpl")
		createTemplateFile(tmpDir, "packer.json.provisioners.ptpl")
		createTemplateFile(tmpDir, "Vagrantfile.tpl")
	})
	JustBeforeEach(func() {
		templates = tpl.New(tmpDir, osName)
	})
	AfterEach(func() {
		os.RemoveAll(tmpDir)
	})

	Context("Windows2012r2", func() {
		BeforeEach(func() {
			osName = "windows2012r2"
		})
		Describe("Autounattend.xml.tpl root template", func() {
			var rootTemplate *tpl.Template
			JustBeforeEach(func() {
				rootTemplate = templates.FindTemplate(filepath.Join(tmpDir, "Autounattend.xml.tpl"))
			})
			It("should be found", func() {
				Expect(rootTemplate).ToNot(BeNil())
			})
			It("should have 2 partial templates", func() {
				Expect(rootTemplate.PartialTemplates).To(HaveLen(2))
			})
			It("should include c", func() {
				path := filepath.Join(tmpDir, "Autounattend.xml.oobe.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
			It("should include Autounattend.xml.disks.ptpl", func() {
				path := filepath.Join(tmpDir, "Autounattend.xml.disks.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
			It("should have template content", func() {
				expected := `Autounattend.xml.tpl
{{define "disks"}}
Autounattend.xml.disks.ptpl
{{end}}
{{define "oobe"}}
Autounattend.xml.oobe.ptpl
{{end}}`
				var buffer bytes.Buffer
				Expect(rootTemplate.Content(&buffer)).NotTo(HaveOccurred())
				Expect(buffer.String()).To(Equal(expected))
			})
		})
	})

	Context("Nano", func() {
		BeforeEach(func() {
			osName = "nano"
		})
		Describe("Autounattend.xml.tpl root template", func() {
			var rootTemplate *tpl.Template
			JustBeforeEach(func() {
				rootTemplate = templates.FindTemplate(filepath.Join(tmpDir, "Autounattend.xml.tpl"))
			})
			It("should be found", func() {
				Expect(rootTemplate).ToNot(BeNil())
			})
			It("should have 2 partial templates", func() {
				Expect(rootTemplate.PartialTemplates).To(HaveLen(2))
			})
			It("should include Autounattend.xml.oobe.ptpl", func() {
				path := filepath.Join(tmpDir, "Autounattend.xml.oobe.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
			It("should include nano/Autounattend.xml.disks.ptpl", func() {
				path := filepath.Join(tmpDir, "nano/Autounattend.xml.disks.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
		})

		Describe("packer.json.tpl root template", func() {
			var rootTemplate *tpl.Template
			JustBeforeEach(func() {
				rootTemplate = templates.FindTemplate(filepath.Join(tmpDir, "packer.json.tpl"))
			})
			It("should be found", func() {
				Expect(rootTemplate).ToNot(BeNil())
			})
			It("should have 2 partial templates", func() {
				Expect(rootTemplate.PartialTemplates).To(HaveLen(2))
			})
			It("should include packer.json.builders.ptpl", func() {
				path := filepath.Join(tmpDir, "packer.json.builders.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
			It("should include nano/packer.json.provisioners.ptpl", func() {
				path := filepath.Join(tmpDir, "nano/packer.json.provisioners.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
		})

		Describe("Vagrantfile.tpl root template", func() {
			var rootTemplate *tpl.Template
			JustBeforeEach(func() {
				rootTemplate = templates.FindTemplate(filepath.Join(tmpDir, "Vagrantfile.tpl"))
			})
			It("should be found", func() {
				Expect(rootTemplate).ToNot(BeNil())
			})
			It("should have zero partial templates", func() {
				Expect(rootTemplate.PartialTemplates).To(BeEmpty())
			})
		})
	})
})

func createTemplateFile(baseDir, path string) {
	fullPath := filepath.Join(baseDir, path)
	fileName := filepath.Base(fullPath)
	fullDir := strings.TrimSuffix(fullPath, fileName)
	os.MkdirAll(fullDir, 0744)
	ioutil.WriteFile(fullPath, []byte(fileName), 0644)
}
<commit_msg>Fix typo in test description<commit_after>package tpl_test

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

	"github.com/joefitzgerald/inductor/tpl"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Tpl", func() {
	var (
		err       error
		osName    string
		tmpDir    string
		templates *tpl.Templates
	)
	BeforeEach(func() {
		tmpDir, err = ioutil.TempDir("", "inductor")
		Expect(err).NotTo(HaveOccurred())
		createTemplateFile(tmpDir, "nano/Autounattend.xml.disks.ptpl")
		createTemplateFile(tmpDir, "nano/packer.json.provisioners.ptpl")
		createTemplateFile(tmpDir, "windowsxp/Autounattend.xml.tpl")
		createTemplateFile(tmpDir, "scripts/win-updates.ps1")
		createTemplateFile(tmpDir, "scripts/nano/SetupComplete.cmd")
		createTemplateFile(tmpDir, "inductor.json")
		createTemplateFile(tmpDir, "README.md")
		createTemplateFile(tmpDir, "Autounattend.xml.tpl")
		createTemplateFile(tmpDir, "Autounattend.xml.oobe.ptpl")
		createTemplateFile(tmpDir, "Autounattend.xml.disks.ptpl")
		createTemplateFile(tmpDir, "packer.json.tpl")
		createTemplateFile(tmpDir, "packer.json.builders.ptpl")
		createTemplateFile(tmpDir, "packer.json.provisioners.ptpl")
		createTemplateFile(tmpDir, "Vagrantfile.tpl")
	})
	JustBeforeEach(func() {
		templates = tpl.New(tmpDir, osName)
	})
	AfterEach(func() {
		os.RemoveAll(tmpDir)
	})

	Context("Windows2012r2", func() {
		BeforeEach(func() {
			osName = "windows2012r2"
		})
		Describe("Autounattend.xml.tpl root template", func() {
			var rootTemplate *tpl.Template
			JustBeforeEach(func() {
				rootTemplate = templates.FindTemplate(filepath.Join(tmpDir, "Autounattend.xml.tpl"))
			})
			It("should be found", func() {
				Expect(rootTemplate).ToNot(BeNil())
			})
			It("should have 2 partial templates", func() {
				Expect(rootTemplate.PartialTemplates).To(HaveLen(2))
			})
			It("should include Autounattend.xml.oobe.ptpl", func() {
				path := filepath.Join(tmpDir, "Autounattend.xml.oobe.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
			It("should include Autounattend.xml.disks.ptpl", func() {
				path := filepath.Join(tmpDir, "Autounattend.xml.disks.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
			It("should have template content", func() {
				expected := `Autounattend.xml.tpl
{{define "disks"}}
Autounattend.xml.disks.ptpl
{{end}}
{{define "oobe"}}
Autounattend.xml.oobe.ptpl
{{end}}`
				var buffer bytes.Buffer
				Expect(rootTemplate.Content(&buffer)).NotTo(HaveOccurred())
				Expect(buffer.String()).To(Equal(expected))
			})
		})
	})

	Context("Nano", func() {
		BeforeEach(func() {
			osName = "nano"
		})
		Describe("Autounattend.xml.tpl root template", func() {
			var rootTemplate *tpl.Template
			JustBeforeEach(func() {
				rootTemplate = templates.FindTemplate(filepath.Join(tmpDir, "Autounattend.xml.tpl"))
			})
			It("should be found", func() {
				Expect(rootTemplate).ToNot(BeNil())
			})
			It("should have 2 partial templates", func() {
				Expect(rootTemplate.PartialTemplates).To(HaveLen(2))
			})
			It("should include Autounattend.xml.oobe.ptpl", func() {
				path := filepath.Join(tmpDir, "Autounattend.xml.oobe.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
			It("should include nano/Autounattend.xml.disks.ptpl", func() {
				path := filepath.Join(tmpDir, "nano/Autounattend.xml.disks.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
		})

		Describe("packer.json.tpl root template", func() {
			var rootTemplate *tpl.Template
			JustBeforeEach(func() {
				rootTemplate = templates.FindTemplate(filepath.Join(tmpDir, "packer.json.tpl"))
			})
			It("should be found", func() {
				Expect(rootTemplate).ToNot(BeNil())
			})
			It("should have 2 partial templates", func() {
				Expect(rootTemplate.PartialTemplates).To(HaveLen(2))
			})
			It("should include packer.json.builders.ptpl", func() {
				path := filepath.Join(tmpDir, "packer.json.builders.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
			It("should include nano/packer.json.provisioners.ptpl", func() {
				path := filepath.Join(tmpDir, "nano/packer.json.provisioners.ptpl")
				Expect(rootTemplate.FindPartialTemplate(path)).ToNot(BeNil())
			})
		})

		Describe("Vagrantfile.tpl root template", func() {
			var rootTemplate *tpl.Template
			JustBeforeEach(func() {
				rootTemplate = templates.FindTemplate(filepath.Join(tmpDir, "Vagrantfile.tpl"))
			})
			It("should be found", func() {
				Expect(rootTemplate).ToNot(BeNil())
			})
			It("should have zero partial templates", func() {
				Expect(rootTemplate.PartialTemplates).To(BeEmpty())
			})
		})
	})
})

func createTemplateFile(baseDir, path string) {
	fullPath := filepath.Join(baseDir, path)
	fileName := filepath.Base(fullPath)
	fullDir := strings.TrimSuffix(fullPath, fileName)
	os.MkdirAll(fullDir, 0744)
	ioutil.WriteFile(fullPath, []byte(fileName), 0644)
}
<|endoftext|>


<commit_before>// Command safesql is a tool for performing static analysis on programs to
// ensure that SQL injection attacks are not possible. It does this by ensuring
// package database/sql is only used with compile-time constant queries.
package main

import (
	"flag"
	"fmt"
	"os"

	"golang.org/x/tools/go/callgraph"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/pointer"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/types"
)

func main() {
	var verbose, quiet bool
	flag.BoolVar(&verbose, "v", false, "Verbose mode")
	flag.BoolVar(&quiet, "q", false, "Only print on failure")

	flag.Parse()
	pkgs := flag.Args()

	c := loader.Config{SourceImports: true}
	c.Import("database/sql")
	for _, pkg := range pkgs {
		c.Import(pkg)
	}
	p, err := c.Load()
	if err != nil {
		fmt.Printf("error loading packages %v: %v\n", pkgs, err)
		os.Exit(2)
	}
	s := ssa.Create(p, 0)
	s.BuildAll()

	qms := FindQueryMethods(p.ImportMap["database/sql"], s)
	if verbose {
		fmt.Println("database/sql functions that accept queries:")
		for _, m := range qms {
			fmt.Printf("- %s (param %d)\n", m.Func, m.Param)
		}
		fmt.Println()
	}

	res, err := pointer.Analyze(&pointer.Config{
		Mains:          FindMains(p, s),
		BuildCallGraph: true,
	})
	if err != nil {
		fmt.Printf("error performing pointer analysis: %v\n", err)
		os.Exit(2)
	}

	bad := FindNonConstCalls(res.CallGraph, qms)
	if len(bad) == 0 {
		if !quiet {
			fmt.Println(`You're safe from SQL injection! Yay \o/`)
		}
		return
	}

	fmt.Printf("Found %d potentially unsafe SQL statements:\n", len(bad))
	for _, ci := range bad {
		pos := p.Fset.Position(ci.Pos())
		fmt.Printf("- %s\n", pos)
	}
	fmt.Println("Please ensure that all SQL queries you use are compile-time constants.")
	fmt.Println("You should always use parametrized queries or prepared statements")
	fmt.Println("instead of building queries from strings.")
	os.Exit(1)
}

// QueryMethod represents a method on a type which has a string parameter named
// "query".
type QueryMethod struct {
	Func     *types.Func
	SSA      *ssa.Function
	ArgCount int
	Param    int
}

// FindQueryMethods locates all methods in the given package (assumed to be
// package database/sql) with a string parameter named "query".
func FindQueryMethods(sql *types.Package, ssa *ssa.Program) []*QueryMethod {
	methods := make([]*QueryMethod, 0)
	scope := sql.Scope()
	for _, name := range scope.Names() {
		o := scope.Lookup(name)
		if !o.Exported() {
			continue
		}
		if _, ok := o.(*types.TypeName); !ok {
			continue
		}
		n := o.Type().(*types.Named)
		for i := 0; i < n.NumMethods(); i++ {
			m := n.Method(i)
			if !m.Exported() {
				continue
			}
			s := m.Type().(*types.Signature)
			if num, ok := FuncHasQuery(s); ok {
				methods = append(methods, &QueryMethod{
					Func:     m,
					SSA:      ssa.FuncValue(m),
					ArgCount: s.Params().Len(),
					Param:    num,
				})
			}
		}
	}
	return methods
}

var stringType = types.New("string")

// FuncHasQuery returns the offset of the string parameter named "query", or
// none if no such parameter exists.
func FuncHasQuery(s *types.Signature) (offset int, ok bool) {
	params := s.Params()
	for i := 0; i < params.Len(); i++ {
		v := params.At(i)
		if v.Name() == "query" && v.Type() == stringType {
			return i, true
		}
	}
	return 0, false
}

// FindMains returns the set of all packages loaded into the given
// loader.Program which contain main functions
func FindMains(p *loader.Program, s *ssa.Program) []*ssa.Package {
	ips := p.InitialPackages()
	mains := make([]*ssa.Package, 0, len(ips))
	for _, info := range ips {
		ssaPkg := s.Package(info.Pkg)
		if ssaPkg.Func("main") != nil {
			mains = append(mains, ssaPkg)
		}
	}
	return mains
}

// FindNonConstCalls returns the set of callsites of the given set of methods
// for which the "query" parameter is not a compile-time constant.
func FindNonConstCalls(cg *callgraph.Graph, qms []*QueryMethod) []ssa.CallInstruction {
	cg.DeleteSyntheticNodes()

	// package database/sql has a couple helper functions which are thin
	// wrappers around other sensitive functions. Instead of handling the
	// general case by tracing down callsites of wrapper functions
	// recursively, let's just whitelist the functions we're already
	// tracking, since it happens to be good enough for our use case.
	okFuncs := make(map[*ssa.Function]struct{}, len(qms))
	for _, m := range qms {
		okFuncs[m.SSA] = struct{}{}
	}

	bad := make([]ssa.CallInstruction, 0)
	for _, m := range qms {
		node := cg.CreateNode(m.SSA)
		for _, edge := range node.In {
			if _, ok := okFuncs[edge.Site.Parent()]; ok {
				continue
			}
			cc := edge.Site.Common()
			args := cc.Args
			// The first parameter is occasionally the receiver.
			if len(args) == m.ArgCount+1 {
				args = args[1:]
			} else if len(args) != m.ArgCount {
				panic("arg count mismatch")
			}
			v := args[m.Param]
			if _, ok := v.(*ssa.Const); !ok {
				bad = append(bad, edge.Site)
			}
		}
	}

	return bad
}
<commit_msg>Better usage instructions<commit_after>// Command safesql is a tool for performing static analysis on programs to
// ensure that SQL injection attacks are not possible. It does this by ensuring
// package database/sql is only used with compile-time constant queries.
package main

import (
	"flag"
	"fmt"
	"os"

	"golang.org/x/tools/go/callgraph"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/pointer"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/types"
)

func main() {
	var verbose, quiet bool
	flag.BoolVar(&verbose, "v", false, "Verbose mode")
	flag.BoolVar(&quiet, "q", false, "Only print on failure")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [-q] [-v] package1 [package2 ...]\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.Parse()
	pkgs := flag.Args()
	if len(pkgs) == 0 {
		flag.Usage()
		os.Exit(2)
	}

	c := loader.Config{SourceImports: true}
	c.Import("database/sql")
	for _, pkg := range pkgs {
		c.Import(pkg)
	}
	p, err := c.Load()
	if err != nil {
		fmt.Printf("error loading packages %v: %v\n", pkgs, err)
		os.Exit(2)
	}
	s := ssa.Create(p, 0)
	s.BuildAll()

	qms := FindQueryMethods(p.ImportMap["database/sql"], s)
	if verbose {
		fmt.Println("database/sql functions that accept queries:")
		for _, m := range qms {
			fmt.Printf("- %s (param %d)\n", m.Func, m.Param)
		}
		fmt.Println()
	}

	res, err := pointer.Analyze(&pointer.Config{
		Mains:          FindMains(p, s),
		BuildCallGraph: true,
	})
	if err != nil {
		fmt.Printf("error performing pointer analysis: %v\n", err)
		os.Exit(2)
	}

	bad := FindNonConstCalls(res.CallGraph, qms)
	if len(bad) == 0 {
		if !quiet {
			fmt.Println(`You're safe from SQL injection! Yay \o/`)
		}
		return
	}

	fmt.Printf("Found %d potentially unsafe SQL statements:\n", len(bad))
	for _, ci := range bad {
		pos := p.Fset.Position(ci.Pos())
		fmt.Printf("- %s\n", pos)
	}
	fmt.Println("Please ensure that all SQL queries you use are compile-time constants.")
	fmt.Println("You should always use parametrized queries or prepared statements")
	fmt.Println("instead of building queries from strings.")
	os.Exit(1)
}

// QueryMethod represents a method on a type which has a string parameter named
// "query".
type QueryMethod struct {
	Func     *types.Func
	SSA      *ssa.Function
	ArgCount int
	Param    int
}

// FindQueryMethods locates all methods in the given package (assumed to be
// package database/sql) with a string parameter named "query".
func FindQueryMethods(sql *types.Package, ssa *ssa.Program) []*QueryMethod {
	methods := make([]*QueryMethod, 0)
	scope := sql.Scope()
	for _, name := range scope.Names() {
		o := scope.Lookup(name)
		if !o.Exported() {
			continue
		}
		if _, ok := o.(*types.TypeName); !ok {
			continue
		}
		n := o.Type().(*types.Named)
		for i := 0; i < n.NumMethods(); i++ {
			m := n.Method(i)
			if !m.Exported() {
				continue
			}
			s := m.Type().(*types.Signature)
			if num, ok := FuncHasQuery(s); ok {
				methods = append(methods, &QueryMethod{
					Func:     m,
					SSA:      ssa.FuncValue(m),
					ArgCount: s.Params().Len(),
					Param:    num,
				})
			}
		}
	}
	return methods
}

var stringType = types.New("string")

// FuncHasQuery returns the offset of the string parameter named "query", or
// none if no such parameter exists.
func FuncHasQuery(s *types.Signature) (offset int, ok bool) {
	params := s.Params()
	for i := 0; i < params.Len(); i++ {
		v := params.At(i)
		if v.Name() == "query" && v.Type() == stringType {
			return i, true
		}
	}
	return 0, false
}

// FindMains returns the set of all packages loaded into the given
// loader.Program which contain main functions
func FindMains(p *loader.Program, s *ssa.Program) []*ssa.Package {
	ips := p.InitialPackages()
	mains := make([]*ssa.Package, 0, len(ips))
	for _, info := range ips {
		ssaPkg := s.Package(info.Pkg)
		if ssaPkg.Func("main") != nil {
			mains = append(mains, ssaPkg)
		}
	}
	return mains
}

// FindNonConstCalls returns the set of callsites of the given set of methods
// for which the "query" parameter is not a compile-time constant.
func FindNonConstCalls(cg *callgraph.Graph, qms []*QueryMethod) []ssa.CallInstruction {
	cg.DeleteSyntheticNodes()

	// package database/sql has a couple helper functions which are thin
	// wrappers around other sensitive functions. Instead of handling the
	// general case by tracing down callsites of wrapper functions
	// recursively, let's just whitelist the functions we're already
	// tracking, since it happens to be good enough for our use case.
	okFuncs := make(map[*ssa.Function]struct{}, len(qms))
	for _, m := range qms {
		okFuncs[m.SSA] = struct{}{}
	}

	bad := make([]ssa.CallInstruction, 0)
	for _, m := range qms {
		node := cg.CreateNode(m.SSA)
		for _, edge := range node.In {
			if _, ok := okFuncs[edge.Site.Parent()]; ok {
				continue
			}
			cc := edge.Site.Common()
			args := cc.Args
			// The first parameter is occasionally the receiver.
			if len(args) == m.ArgCount+1 {
				args = args[1:]
			} else if len(args) != m.ArgCount {
				panic("arg count mismatch")
			}
			v := args[m.Param]
			if _, ok := v.(*ssa.Const); !ok {
				bad = append(bad, edge.Site)
			}
		}
	}

	return bad
}
<|endoftext|>


<commit_before>package exec

import (
	"bytes"
	"fmt"
	"os/exec"
	"sync"

	"github.com/gonuts/go-shellquote"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
	"github.com/influxdata/telegraf/plugins/parsers"
)

const sampleConfig = `
  ### Commands array
  commands = ["/tmp/test.sh", "/usr/bin/mycollector --foo=bar"]

  ### measurement name suffix (for separating different commands)
  name_suffix = "_mycollector"

  ### Data format to consume. This can be "json", "influx" or "graphite"
  ### Each data format has it's own unique set of configuration options, read
  ### more about them here:
  ### https://github.com/influxdata/telegraf/blob/master/DATA_FORMATS_INPUT.md
  data_format = "influx"
`

type Exec struct {
	Commands []string
	Command  string

	parser parsers.Parser

	wg sync.WaitGroup

	runner  Runner
	errChan chan error
}

func NewExec() *Exec {
	return &Exec{
		runner: CommandRunner{},
	}
}

type Runner interface {
	Run(*Exec, string) ([]byte, error)
}

type CommandRunner struct{}

func (c CommandRunner) Run(e *Exec, command string) ([]byte, error) {
	split_cmd, err := shellquote.Split(command)
	if err != nil || len(split_cmd) == 0 {
		return nil, fmt.Errorf("exec: unable to parse command, %s", err)
	}

	cmd := exec.Command(split_cmd[0], split_cmd[1:]...)

	var out bytes.Buffer
	cmd.Stdout = &out

	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("exec: %s for command '%s'", err, command)
	}

	return out.Bytes(), nil
}

func (e *Exec) ProcessCommand(command string, acc telegraf.Accumulator) {
	defer e.wg.Done()

	out, err := e.runner.Run(e, command)
	if err != nil {
		e.errChan <- err
		return
	}

	metrics, err := e.parser.Parse(out)
	if err != nil {
		e.errChan <- err
	} else {
		for _, metric := range metrics {
			acc.AddFields(metric.Name(), metric.Fields(), metric.Tags(), metric.Time())
		}
	}
}

func (e *Exec) SampleConfig() string {
	return sampleConfig
}

func (e *Exec) Description() string {
	return "Read metrics from one or more commands that can output to stdout"
}

func (e *Exec) SetParser(parser parsers.Parser) {
	e.parser = parser
}

func (e *Exec) Gather(acc telegraf.Accumulator) error {
	e.errChan = make(chan error, len(e.Commands))

	e.wg.Add(len(e.Commands))
	for _, command := range e.Commands {
		go e.ProcessCommand(command, acc)
	}
	e.wg.Wait()

	select {
	default:
		close(e.errChan)
		return nil
	case err := <-e.errChan:
		close(e.errChan)
		return err
	}

}

func init() {
	inputs.Add("exec", func() telegraf.Input {
		return NewExec()
	})
}
<commit_msg>Fix exec input legacy behavior, command=''<commit_after>package exec

import (
	"bytes"
	"fmt"
	"os/exec"
	"sync"

	"github.com/gonuts/go-shellquote"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
	"github.com/influxdata/telegraf/plugins/parsers"
)

const sampleConfig = `
  ### Commands array
  commands = ["/tmp/test.sh", "/usr/bin/mycollector --foo=bar"]

  ### measurement name suffix (for separating different commands)
  name_suffix = "_mycollector"

  ### Data format to consume. This can be "json", "influx" or "graphite"
  ### Each data format has it's own unique set of configuration options, read
  ### more about them here:
  ### https://github.com/influxdata/telegraf/blob/master/DATA_FORMATS_INPUT.md
  data_format = "influx"
`

type Exec struct {
	Commands []string
	Command  string

	parser parsers.Parser

	wg sync.WaitGroup

	runner  Runner
	errChan chan error
}

func NewExec() *Exec {
	return &Exec{
		runner: CommandRunner{},
	}
}

type Runner interface {
	Run(*Exec, string) ([]byte, error)
}

type CommandRunner struct{}

func (c CommandRunner) Run(e *Exec, command string) ([]byte, error) {
	split_cmd, err := shellquote.Split(command)
	if err != nil || len(split_cmd) == 0 {
		return nil, fmt.Errorf("exec: unable to parse command, %s", err)
	}

	cmd := exec.Command(split_cmd[0], split_cmd[1:]...)

	var out bytes.Buffer
	cmd.Stdout = &out

	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("exec: %s for command '%s'", err, command)
	}

	return out.Bytes(), nil
}

func (e *Exec) ProcessCommand(command string, acc telegraf.Accumulator) {
	defer e.wg.Done()

	out, err := e.runner.Run(e, command)
	if err != nil {
		e.errChan <- err
		return
	}

	metrics, err := e.parser.Parse(out)
	if err != nil {
		e.errChan <- err
	} else {
		for _, metric := range metrics {
			acc.AddFields(metric.Name(), metric.Fields(), metric.Tags(), metric.Time())
		}
	}
}

func (e *Exec) SampleConfig() string {
	return sampleConfig
}

func (e *Exec) Description() string {
	return "Read metrics from one or more commands that can output to stdout"
}

func (e *Exec) SetParser(parser parsers.Parser) {
	e.parser = parser
}

func (e *Exec) Gather(acc telegraf.Accumulator) error {
	// Legacy single command support
	if e.Command != "" {
		e.Commands = append(e.Commands, e.Command)
		e.Command = ""
	}

	e.errChan = make(chan error, len(e.Commands))

	e.wg.Add(len(e.Commands))
	for _, command := range e.Commands {
		go e.ProcessCommand(command, acc)
	}
	e.wg.Wait()

	select {
	default:
		close(e.errChan)
		return nil
	case err := <-e.errChan:
		close(e.errChan)
		return err
	}

}

func init() {
	inputs.Add("exec", func() telegraf.Input {
		return NewExec()
	})
}
<|endoftext|>


<commit_before>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"errors"
	"github.com/getgauge/gauge/logger"
	"golang.org/x/tools/go/exact"
	"golang.org/x/tools/go/types"
	"os"
	"regexp"
	"strconv"
	"strings"
)

type scenarioIndexFilterToRetain struct {
	indexToNotFilter     int
	currentScenarioIndex int
}
type ScenarioFilterBasedOnTags struct {
	specTags       []string
	tagExpression  string
	isSpecSelected bool
}

func newScenarioIndexFilterToRetain(index int) *scenarioIndexFilterToRetain {
	return &scenarioIndexFilterToRetain{index, 0}
}

func (filter *scenarioIndexFilterToRetain) filter(item item) bool {
	if item.kind() == scenarioKind {
		if filter.currentScenarioIndex != filter.indexToNotFilter {
			filter.currentScenarioIndex++
			return true
		} else {
			filter.currentScenarioIndex++
			return false
		}
	}
	return false
}

func newScenarioFilterBasedOnTags(specTags []string, tagExp string) *ScenarioFilterBasedOnTags {
	return &ScenarioFilterBasedOnTags{specTags, tagExp, false}
}

func (filter *ScenarioFilterBasedOnTags) filter(item item) bool {
	if item.kind() == scenarioKind {
		if filter.isSpecSelected {
			return false
		}
		if filter.filterTags(filter.specTags) {
			filter.isSpecSelected = true
			return false
		}
		tags := item.(*scenario).tags
		if tags == nil {
			return !filter.filterTags(make([]string, 0))
		}
		return !filter.filterTags(tags.values)
	}
	return true
}

func (filter *ScenarioFilterBasedOnTags) filterTags(stags []string) bool {
	tagsMap := make(map[string]bool, 0)
	for _, tag := range stags {
		tagsMap[strings.Replace(tag, " ", "", -1)] = true
	}
	filter.replaceSpecialChar()
	value, _ := filter.formatAndEvaluateExpression(tagsMap, filter.isTagPresent)
	return value
}
func (filter *ScenarioFilterBasedOnTags) replaceSpecialChar() {
	filter.tagExpression = strings.Replace(strings.Replace(strings.Replace(strings.Replace(filter.tagExpression, " ", "", -1), ",", "&", -1), "&&", "&", -1), "||", "|", -1)
}

func (filter *ScenarioFilterBasedOnTags) formatAndEvaluateExpression(tagsMap map[string]bool, isTagQualified func(tagsMap map[string]bool, tagName string) bool) (bool, error) {
	_, tags := filter.getOperatorsAndOperands()
	expToBeEvaluated := filter.tagExpression
	for _, tag := range tags {
		expToBeEvaluated = strings.Replace(expToBeEvaluated, strings.TrimSpace(tag), strconv.FormatBool(isTagQualified(tagsMap, strings.TrimSpace(tag))), -1)
	}
	return filter.evaluateExp(filter.handleNegation(expToBeEvaluated))
}

func (filter *ScenarioFilterBasedOnTags) handleNegation(tagExpression string) string {
	tagExpression = strings.Replace(strings.Replace(tagExpression, "!true", "false", -1), "!false", "true", -1)
	for strings.Contains(tagExpression, "!(") {
		tagExpression = filter.evaluateBrackets(tagExpression)
	}
	return tagExpression
}

func (filter *ScenarioFilterBasedOnTags) evaluateBrackets(tagExpression string) string {
	if strings.Contains(tagExpression, "!(") {
		innerText := filter.resolveBracketExpression(tagExpression)
		return strings.Replace(tagExpression, "!("+innerText+")", filter.evaluateBrackets(innerText), -1)
	}
	value, _ := filter.evaluateExp(tagExpression)
	return strconv.FormatBool(!value)
}

func (filter *ScenarioFilterBasedOnTags) resolveBracketExpression(tagExpression string) string {
	indexOfOpenBracket := strings.Index(tagExpression, "!(") + 1
	bracketStack := make([]string, 0)
	i := indexOfOpenBracket
	for ; i < len(tagExpression); i++ {
		if tagExpression[i] == '(' {
			bracketStack = append(bracketStack, "(")
		} else if tagExpression[i] == ')' {
			bracketStack = append(bracketStack[:len(bracketStack)-1])
		}
		if len(bracketStack) == 0 {
			break
		}
	}
	return tagExpression[indexOfOpenBracket+1 : i]
}

func (filter *ScenarioFilterBasedOnTags) evaluateExp(tagExpression string) (bool, error) {
	tre := regexp.MustCompile("true")
	fre := regexp.MustCompile("false")

	s := fre.ReplaceAllString(tre.ReplaceAllString(tagExpression, "1"), "0")

	val, err := types.Eval(s, nil, nil)
	if err != nil {
		return false, errors.New("Invalid Expression.\n" + err.Error())
	}
	res, _ := exact.Uint64Val(val.Value)

	var final bool
	if res == 1 {
		final = true
	} else {
		final = false
	}

	return final, nil
}

func (filter *ScenarioFilterBasedOnTags) isTagPresent(tagsMap map[string]bool, tagName string) bool {
	_, ok := tagsMap[tagName]
	return ok
}

func (filter *ScenarioFilterBasedOnTags) getOperatorsAndOperands() ([]string, []string) {
	listOfOperators := make([]string, 0)
	listOfTags := strings.FieldsFunc(filter.tagExpression, func(r rune) bool {
		isValidOperator := r == '&' || r == '|' || r == '(' || r == ')' || r == '!'
		if isValidOperator {
			operator, _ := strconv.Unquote(strconv.QuoteRuneToASCII(r))
			listOfOperators = append(listOfOperators, operator)
			return isValidOperator
		}
		return false
	})
	return listOfOperators, listOfTags
}

func filterSpecsItems(specs []*specification, filter specFilter) []*specification {
	filteredSpecs := make([]*specification, 0)
	for _, spec := range specs {
		spec.filter(filter)
		if len(spec.scenarios) != 0 {
			filteredSpecs = append(filteredSpecs, spec)
		}
	}
	return filteredSpecs
}

func filterSpecsByTags(specs []*specification, tagExpression string) []*specification {
	filteredSpecs := make([]*specification, 0)
	for _, spec := range specs {
		tagValues := make([]string, 0)
		if spec.tags != nil {
			tagValues = spec.tags.values
		}
		spec.filter(newScenarioFilterBasedOnTags(tagValues, tagExpression))
		if len(spec.scenarios) != 0 {
			filteredSpecs = append(filteredSpecs, spec)
		}
	}
	return filteredSpecs
}

func validateTagExpression(tagExpression string) {
	filter := &ScenarioFilterBasedOnTags{tagExpression: tagExpression}
	filter.replaceSpecialChar()
	_, err := filter.formatAndEvaluateExpression(make(map[string]bool, 0), func(a map[string]bool, b string) bool { return true })
	if err != nil {
		logger.Log.Critical(err.Error())
		os.Exit(1)
	}
}
<commit_msg>using spec tags only as abstraction over scenario tags<commit_after>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"errors"
	"github.com/getgauge/gauge/logger"
	"golang.org/x/tools/go/exact"
	"golang.org/x/tools/go/types"
	"os"
	"regexp"
	"strconv"
	"strings"
)

type scenarioIndexFilterToRetain struct {
	indexToNotFilter     int
	currentScenarioIndex int
}
type ScenarioFilterBasedOnTags struct {
	specTags      []string
	tagExpression string
}

func newScenarioIndexFilterToRetain(index int) *scenarioIndexFilterToRetain {
	return &scenarioIndexFilterToRetain{index, 0}
}

func (filter *scenarioIndexFilterToRetain) filter(item item) bool {
	if item.kind() == scenarioKind {
		if filter.currentScenarioIndex != filter.indexToNotFilter {
			filter.currentScenarioIndex++
			return true
		} else {
			filter.currentScenarioIndex++
			return false
		}
	}
	return false
}

func newScenarioFilterBasedOnTags(specTags []string, tagExp string) *ScenarioFilterBasedOnTags {
	return &ScenarioFilterBasedOnTags{specTags, tagExp}
}

func (filter *ScenarioFilterBasedOnTags) filter(item item) bool {
	if item.kind() == scenarioKind {
		tags := item.(*scenario).tags
		if tags == nil {
			return !filter.filterTags(filter.specTags)
		}
		return !filter.filterTags(append(tags.values, filter.specTags...))
	}
	return true
}

func (filter *ScenarioFilterBasedOnTags) filterTags(stags []string) bool {
	tagsMap := make(map[string]bool, 0)
	for _, tag := range stags {
		tagsMap[strings.Replace(tag, " ", "", -1)] = true
	}
	filter.replaceSpecialChar()
	value, _ := filter.formatAndEvaluateExpression(tagsMap, filter.isTagPresent)
	return value
}
func (filter *ScenarioFilterBasedOnTags) replaceSpecialChar() {
	filter.tagExpression = strings.Replace(strings.Replace(strings.Replace(strings.Replace(filter.tagExpression, " ", "", -1), ",", "&", -1), "&&", "&", -1), "||", "|", -1)
}

func (filter *ScenarioFilterBasedOnTags) formatAndEvaluateExpression(tagsMap map[string]bool, isTagQualified func(tagsMap map[string]bool, tagName string) bool) (bool, error) {
	_, tags := filter.getOperatorsAndOperands()
	expToBeEvaluated := filter.tagExpression
	for _, tag := range tags {
		expToBeEvaluated = strings.Replace(expToBeEvaluated, strings.TrimSpace(tag), strconv.FormatBool(isTagQualified(tagsMap, strings.TrimSpace(tag))), -1)
	}
	return filter.evaluateExp(filter.handleNegation(expToBeEvaluated))
}

func (filter *ScenarioFilterBasedOnTags) handleNegation(tagExpression string) string {
	tagExpression = strings.Replace(strings.Replace(tagExpression, "!true", "false", -1), "!false", "true", -1)
	for strings.Contains(tagExpression, "!(") {
		tagExpression = filter.evaluateBrackets(tagExpression)
	}
	return tagExpression
}

func (filter *ScenarioFilterBasedOnTags) evaluateBrackets(tagExpression string) string {
	if strings.Contains(tagExpression, "!(") {
		innerText := filter.resolveBracketExpression(tagExpression)
		return strings.Replace(tagExpression, "!("+innerText+")", filter.evaluateBrackets(innerText), -1)
	}
	value, _ := filter.evaluateExp(tagExpression)
	return strconv.FormatBool(!value)
}

func (filter *ScenarioFilterBasedOnTags) resolveBracketExpression(tagExpression string) string {
	indexOfOpenBracket := strings.Index(tagExpression, "!(") + 1
	bracketStack := make([]string, 0)
	i := indexOfOpenBracket
	for ; i < len(tagExpression); i++ {
		if tagExpression[i] == '(' {
			bracketStack = append(bracketStack, "(")
		} else if tagExpression[i] == ')' {
			bracketStack = append(bracketStack[:len(bracketStack)-1])
		}
		if len(bracketStack) == 0 {
			break
		}
	}
	return tagExpression[indexOfOpenBracket+1 : i]
}

func (filter *ScenarioFilterBasedOnTags) evaluateExp(tagExpression string) (bool, error) {
	tre := regexp.MustCompile("true")
	fre := regexp.MustCompile("false")

	s := fre.ReplaceAllString(tre.ReplaceAllString(tagExpression, "1"), "0")

	val, err := types.Eval(s, nil, nil)
	if err != nil {
		return false, errors.New("Invalid Expression.\n" + err.Error())
	}
	res, _ := exact.Uint64Val(val.Value)

	var final bool
	if res == 1 {
		final = true
	} else {
		final = false
	}

	return final, nil
}

func (filter *ScenarioFilterBasedOnTags) isTagPresent(tagsMap map[string]bool, tagName string) bool {
	_, ok := tagsMap[tagName]
	return ok
}

func (filter *ScenarioFilterBasedOnTags) getOperatorsAndOperands() ([]string, []string) {
	listOfOperators := make([]string, 0)
	listOfTags := strings.FieldsFunc(filter.tagExpression, func(r rune) bool {
		isValidOperator := r == '&' || r == '|' || r == '(' || r == ')' || r == '!'
		if isValidOperator {
			operator, _ := strconv.Unquote(strconv.QuoteRuneToASCII(r))
			listOfOperators = append(listOfOperators, operator)
			return isValidOperator
		}
		return false
	})
	return listOfOperators, listOfTags
}

func filterSpecsItems(specs []*specification, filter specFilter) []*specification {
	filteredSpecs := make([]*specification, 0)
	for _, spec := range specs {
		spec.filter(filter)
		if len(spec.scenarios) != 0 {
			filteredSpecs = append(filteredSpecs, spec)
		}
	}
	return filteredSpecs
}

func filterSpecsByTags(specs []*specification, tagExpression string) []*specification {
	filteredSpecs := make([]*specification, 0)
	for _, spec := range specs {
		tagValues := make([]string, 0)
		if spec.tags != nil {
			tagValues = spec.tags.values
		}
		spec.filter(newScenarioFilterBasedOnTags(tagValues, tagExpression))
		if len(spec.scenarios) != 0 {
			filteredSpecs = append(filteredSpecs, spec)
		}
	}
	return filteredSpecs
}

func validateTagExpression(tagExpression string) {
	filter := &ScenarioFilterBasedOnTags{tagExpression: tagExpression}
	filter.replaceSpecialChar()
	_, err := filter.formatAndEvaluateExpression(make(map[string]bool, 0), func(a map[string]bool, b string) bool { return true })
	if err != nil {
		logger.Log.Critical(err.Error())
		os.Exit(1)
	}
}
<|endoftext|>


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

package track

import (
	"context"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// Dialer represents a net.Dialer that can track bytes read/written
type Dialer interface {
	ByteTracker

	Dial(network, address string) (net.Conn, error)
	DialContext(ctx context.Context, network, address string) (net.Conn, error)

	// Accessors
	Timeout() time.Duration
	Deadline() time.Time
	LocalAddr() net.Addr
	DualStack() bool
	FallbackDelay() time.Duration
	KeepAlive() time.Duration
	Resolver() *net.Resolver

	// Mutators
	SetTimeout(newTimeout time.Duration)
	SetDeadline(newDeadline time.Time)
	SetLocalAddr(newLocalAddr net.Addr)
	SetDualStack(newDualStack bool)
	SetFallbackDelay(newFallbackDelay time.Duration)
	SetKeepAlive(newKeepAlive time.Duration)
	SetResolver(newResolver *net.Resolver)
}

type basicDialer struct {
	bytesRead    uint64
	bytesWritten uint64
	connCounter  uint64

	*net.Dialer
	conns    map[uint64]Conn
	connsMut sync.RWMutex

	// flushMut is for critical sections that can
	// affect totals
	flushMut sync.RWMutex
}

// NewDefaultDialer returns a Dialer based on
// a default net.Dialer
func NewDefaultDialer() Dialer {
	return &basicDialer{
		Dialer: &net.Dialer{
			Timeout:   5 * time.Second,
			KeepAlive: 5 * time.Second,
			DualStack: true,
		},
	}
}

func (dialer *basicDialer) addConn(conn Conn, num uint64) {
	dialer.connsMut.Lock()
	if dialer.conns == nil {
		dialer.conns = map[uint64]Conn{}
	}
	dialer.conns[num] = conn
	dialer.connsMut.Unlock()
}

func (dialer *basicDialer) Dial(network, address string) (net.Conn, error) {
	conn, err := dialer.Dialer.Dial(network, address)
	if err != nil {
		return nil, err
	}
	connNum := dialer.nextConnNum()
	tConn := newConn(conn)
	tConn.OnClose = dialer.makeOnConnClose(connNum)
	dialer.addConn(tConn, connNum)

	return tConn, nil
}

func (dialer *basicDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	conn, err := dialer.Dialer.DialContext(ctx, network, address)
	if err != nil {
		return nil, err
	}
	connNum := dialer.nextConnNum()
	tConn := newConn(conn)
	tConn.OnClose = dialer.makeOnConnClose(connNum)
	dialer.addConn(tConn, connNum)

	return tConn, nil
}

func (dialer *basicDialer) makeOnConnClose(connNum uint64) func() {
	return func() {
		dialer.flushMut.Lock()
		defer dialer.flushMut.Unlock()

		dialer.connsMut.Lock()

		// Ensure connection exists for duration of this
		// call until we delete the connection
		conn, ok := dialer.conns[connNum]

		if !ok {
			dialer.connsMut.Unlock()
			return
		}
		delete(dialer.conns, connNum)
		dialer.connsMut.Unlock()

		read, written := conn.BytesReadWritten()
		atomic.AddUint64(&dialer.bytesRead, read)
		atomic.AddUint64(&dialer.bytesWritten, written)
	}
}

func (dialer *basicDialer) nextConnNum() uint64 {
	return atomic.AddUint64(&dialer.connCounter, 1)
}

func (dialer *basicDialer) BytesRead() uint64 {
	read, _ := dialer.BytesReadWritten()
	return read
}

func (dialer *basicDialer) BytesWritten() uint64 {
	_, written := dialer.BytesReadWritten()
	return written
}

func (dialer *basicDialer) BytesReadWritten() (uint64, uint64) {
	dialer.flushMut.RLock()
	defer dialer.flushMut.RUnlock()

	var totalRead, totalWritten uint64

	dialer.connsMut.RLock()
	for _, conn := range dialer.conns {
		read, written := conn.BytesReadWritten()
		totalRead += read
		totalWritten += written
	}
	dialer.connsMut.RUnlock()

	totalRead += atomic.LoadUint64(&dialer.bytesRead)
	totalWritten += atomic.LoadUint64(&dialer.bytesWritten)

	return totalRead, totalWritten
}

func (dialer *basicDialer) ResetBytes() {
	dialer.flushMut.Lock()
	defer dialer.flushMut.Unlock()

	dialer.connsMut.Lock()
	for _, conn := range dialer.conns {
		conn.ResetBytes()
	}
	dialer.connsMut.Unlock()

	atomic.StoreUint64(&dialer.bytesRead, 0)
	atomic.StoreUint64(&dialer.bytesWritten, 0)
}

func (dialer *basicDialer) Timeout() time.Duration {
	return dialer.Dialer.Timeout
}

func (dialer *basicDialer) Deadline() time.Time {
	return dialer.Dialer.Deadline
}

func (dialer *basicDialer) LocalAddr() net.Addr {
	return dialer.Dialer.LocalAddr
}

func (dialer *basicDialer) DualStack() bool {
	return dialer.Dialer.DualStack
}

func (dialer *basicDialer) FallbackDelay() time.Duration {
	return dialer.Dialer.FallbackDelay
}

func (dialer *basicDialer) KeepAlive() time.Duration {
	return dialer.Dialer.KeepAlive
}

func (dialer *basicDialer) Resolver() *net.Resolver {
	return dialer.Dialer.Resolver
}

func (dialer *basicDialer) SetTimeout(newTimeout time.Duration) {
	dialer.Dialer.Timeout = newTimeout
}

func (dialer *basicDialer) SetDeadline(newDeadline time.Time) {
	dialer.Dialer.Deadline = newDeadline
}

func (dialer *basicDialer) SetLocalAddr(newLocalAddr net.Addr) {
	dialer.Dialer.LocalAddr = newLocalAddr
}

func (dialer *basicDialer) SetDualStack(newDualStack bool) {
	dialer.Dialer.DualStack = newDualStack
}

func (dialer *basicDialer) SetFallbackDelay(newFallbackDelay time.Duration) {
	dialer.Dialer.FallbackDelay = newFallbackDelay
}

func (dialer *basicDialer) SetKeepAlive(newKeepAlive time.Duration) {
	dialer.Dialer.KeepAlive = newKeepAlive
}

func (dialer *basicDialer) SetResolver(newResolver *net.Resolver) {
	dialer.Dialer.Resolver = newResolver
}
<commit_msg>Lock around dialer settings<commit_after>// Copyright 2017 Eric Daniels
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package track

import (
	"context"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// Dialer represents a net.Dialer that can track bytes read/written
type Dialer interface {
	ByteTracker

	Dial(network, address string) (net.Conn, error)
	DialContext(ctx context.Context, network, address string) (net.Conn, error)

	// Accessors
	Timeout() time.Duration
	Deadline() time.Time
	LocalAddr() net.Addr
	DualStack() bool
	FallbackDelay() time.Duration
	KeepAlive() time.Duration
	Resolver() *net.Resolver

	// Mutators
	SetTimeout(newTimeout time.Duration)
	SetDeadline(newDeadline time.Time)
	SetLocalAddr(newLocalAddr net.Addr)
	SetDualStack(newDualStack bool)
	SetFallbackDelay(newFallbackDelay time.Duration)
	SetKeepAlive(newKeepAlive time.Duration)
	SetResolver(newResolver *net.Resolver)
}

type basicDialer struct {
	bytesRead    uint64
	bytesWritten uint64
	connCounter  uint64

	*net.Dialer
	conns    map[uint64]Conn
	connsMut sync.RWMutex

	// flushMut is for critical sections that can
	// affect totals
	flushMut sync.RWMutex

	stgsMut sync.Mutex
}

// NewDefaultDialer returns a Dialer based on
// a default net.Dialer
func NewDefaultDialer() Dialer {
	return &basicDialer{
		Dialer: &net.Dialer{
			Timeout:   5 * time.Second,
			KeepAlive: 5 * time.Second,
			DualStack: true,
		},
	}
}

func (dialer *basicDialer) addConn(conn Conn, num uint64) {
	dialer.connsMut.Lock()
	if dialer.conns == nil {
		dialer.conns = map[uint64]Conn{}
	}
	dialer.conns[num] = conn
	dialer.connsMut.Unlock()
}

func (dialer *basicDialer) Dial(network, address string) (net.Conn, error) {
	conn, err := dialer.Dialer.Dial(network, address)
	if err != nil {
		return nil, err
	}
	connNum := dialer.nextConnNum()
	tConn := newConn(conn)
	tConn.OnClose = dialer.makeOnConnClose(connNum)
	dialer.addConn(tConn, connNum)

	return tConn, nil
}

func (dialer *basicDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	conn, err := dialer.Dialer.DialContext(ctx, network, address)
	if err != nil {
		return nil, err
	}
	connNum := dialer.nextConnNum()
	tConn := newConn(conn)
	tConn.OnClose = dialer.makeOnConnClose(connNum)
	dialer.addConn(tConn, connNum)

	return tConn, nil
}

func (dialer *basicDialer) makeOnConnClose(connNum uint64) func() {
	return func() {
		dialer.flushMut.Lock()
		defer dialer.flushMut.Unlock()

		dialer.connsMut.Lock()

		// Ensure connection exists for duration of this
		// call until we delete the connection
		conn, ok := dialer.conns[connNum]

		if !ok {
			dialer.connsMut.Unlock()
			return
		}
		delete(dialer.conns, connNum)
		dialer.connsMut.Unlock()

		read, written := conn.BytesReadWritten()
		atomic.AddUint64(&dialer.bytesRead, read)
		atomic.AddUint64(&dialer.bytesWritten, written)
	}
}

func (dialer *basicDialer) nextConnNum() uint64 {
	return atomic.AddUint64(&dialer.connCounter, 1)
}

func (dialer *basicDialer) BytesRead() uint64 {
	read, _ := dialer.BytesReadWritten()
	return read
}

func (dialer *basicDialer) BytesWritten() uint64 {
	_, written := dialer.BytesReadWritten()
	return written
}

func (dialer *basicDialer) BytesReadWritten() (uint64, uint64) {
	dialer.flushMut.RLock()
	defer dialer.flushMut.RUnlock()

	var totalRead, totalWritten uint64

	dialer.connsMut.RLock()
	for _, conn := range dialer.conns {
		read, written := conn.BytesReadWritten()
		totalRead += read
		totalWritten += written
	}
	dialer.connsMut.RUnlock()

	totalRead += atomic.LoadUint64(&dialer.bytesRead)
	totalWritten += atomic.LoadUint64(&dialer.bytesWritten)

	return totalRead, totalWritten
}

func (dialer *basicDialer) ResetBytes() {
	dialer.flushMut.Lock()
	defer dialer.flushMut.Unlock()

	dialer.connsMut.Lock()
	for _, conn := range dialer.conns {
		conn.ResetBytes()
	}
	dialer.connsMut.Unlock()

	atomic.StoreUint64(&dialer.bytesRead, 0)
	atomic.StoreUint64(&dialer.bytesWritten, 0)
}

func (dialer *basicDialer) Timeout() time.Duration {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	return dialer.Dialer.Timeout
}

func (dialer *basicDialer) Deadline() time.Time {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	return dialer.Dialer.Deadline
}

func (dialer *basicDialer) LocalAddr() net.Addr {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	return dialer.Dialer.LocalAddr
}

func (dialer *basicDialer) DualStack() bool {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	return dialer.Dialer.DualStack
}

func (dialer *basicDialer) FallbackDelay() time.Duration {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	return dialer.Dialer.FallbackDelay
}

func (dialer *basicDialer) KeepAlive() time.Duration {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	return dialer.Dialer.KeepAlive
}

func (dialer *basicDialer) Resolver() *net.Resolver {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	return dialer.Dialer.Resolver
}

func (dialer *basicDialer) SetTimeout(newTimeout time.Duration) {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	dialer.Dialer.Timeout = newTimeout
}

func (dialer *basicDialer) SetDeadline(newDeadline time.Time) {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	dialer.Dialer.Deadline = newDeadline
}

func (dialer *basicDialer) SetLocalAddr(newLocalAddr net.Addr) {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	dialer.Dialer.LocalAddr = newLocalAddr
}

func (dialer *basicDialer) SetDualStack(newDualStack bool) {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	dialer.Dialer.DualStack = newDualStack
}

func (dialer *basicDialer) SetFallbackDelay(newFallbackDelay time.Duration) {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	dialer.Dialer.FallbackDelay = newFallbackDelay
}

func (dialer *basicDialer) SetKeepAlive(newKeepAlive time.Duration) {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	dialer.Dialer.KeepAlive = newKeepAlive
}

func (dialer *basicDialer) SetResolver(newResolver *net.Resolver) {
	dialer.stgsMut.Lock()
	defer dialer.stgsMut.Unlock()
	dialer.Dialer.Resolver = newResolver
}
<|endoftext|>


<commit_before>package backup

import (
	"fmt"
	"github.com/pkg/errors"
	"github.com/pkg/sftp"
	"github.com/stefanprodan/mgob/config"
	"golang.org/x/crypto/ssh"
	"io"
	"net"
	"os"
	"path/filepath"
)

func sftpUpload(file string, plan config.Plan) (string, error) {
	sshConf := &ssh.ClientConfig{
		User: plan.SFTP.Username,
		Auth: []ssh.AuthMethod{ssh.Password(plan.SFTP.Password)},
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}

	sshCon, err := ssh.Dial("tcp", fmt.Sprintf("%v:%v", plan.SFTP.Host, plan.SFTP.Port), sshConf)
	if err != nil {
		return "", errors.Wrapf(err, "SSH dial to %v:%v failed", plan.SFTP.Host, plan.SFTP.Port)
	}
	defer sshCon.Close()

	sftpClient, err := sftp.NewClient(sshCon)
	if err != nil {
		return "", errors.Wrapf(err, "SFTP client init %v:%v failed", plan.SFTP.Host, plan.SFTP.Port)
	}
	defer sftpClient.Close()

	f, err := os.Open(file)
	if err != nil {
		return "", errors.Wrapf(err, "Opening file %v failed", file)
	}
	defer f.Close()

	_, fname := filepath.Split(file)
	dstPath := filepath.Join(plan.SFTP.Dir, fname)
	sf, err := sftpClient.Create(dstPath)
	if err != nil {
		return "", errors.Wrapf(err, "SFTP %v:%v creating file %v failed", plan.SFTP.Host, plan.SFTP.Port, dstPath)
	}

	_, err = io.Copy(sf, f)
	if err != nil {
		return "", errors.Wrapf(err, "SFTP %v:%v upload file %v failed", plan.SFTP.Host, plan.SFTP.Port, dstPath)
	}
	sf.Close()

	listSftpBackups(sftpClient, plan.SFTP.Dir)

	return fmt.Sprintf("SFTP %v:%v upload %v finished", plan.SFTP.Host, plan.SFTP.Port, dstPath), nil
}

func listSftpBackups(client *sftp.Client, dir string) error {
	list, err := client.ReadDir(fmt.Sprintf("/%v", dir))
	if err != nil {
		return errors.Wrapf(err, "SFTP reading %v dir failed", dir)
	}

	for _, item := range list {
		fmt.Println(item.Name())
	}

	return nil
}
<commit_msg>don't list SFTP files<commit_after>package backup

import (
	"fmt"
	"github.com/pkg/errors"
	"github.com/pkg/sftp"
	"github.com/stefanprodan/mgob/config"
	"golang.org/x/crypto/ssh"
	"io"
	"net"
	"os"
	"path/filepath"
)

func sftpUpload(file string, plan config.Plan) (string, error) {
	sshConf := &ssh.ClientConfig{
		User: plan.SFTP.Username,
		Auth: []ssh.AuthMethod{ssh.Password(plan.SFTP.Password)},
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}

	sshCon, err := ssh.Dial("tcp", fmt.Sprintf("%v:%v", plan.SFTP.Host, plan.SFTP.Port), sshConf)
	if err != nil {
		return "", errors.Wrapf(err, "SSH dial to %v:%v failed", plan.SFTP.Host, plan.SFTP.Port)
	}
	defer sshCon.Close()

	sftpClient, err := sftp.NewClient(sshCon)
	if err != nil {
		return "", errors.Wrapf(err, "SFTP client init %v:%v failed", plan.SFTP.Host, plan.SFTP.Port)
	}
	defer sftpClient.Close()

	f, err := os.Open(file)
	if err != nil {
		return "", errors.Wrapf(err, "Opening file %v failed", file)
	}
	defer f.Close()

	_, fname := filepath.Split(file)
	dstPath := filepath.Join(plan.SFTP.Dir, fname)
	sf, err := sftpClient.Create(dstPath)
	if err != nil {
		return "", errors.Wrapf(err, "SFTP %v:%v creating file %v failed", plan.SFTP.Host, plan.SFTP.Port, dstPath)
	}

	_, err = io.Copy(sf, f)
	if err != nil {
		return "", errors.Wrapf(err, "SFTP %v:%v upload file %v failed", plan.SFTP.Host, plan.SFTP.Port, dstPath)
	}
	sf.Close()

	//listSftpBackups(sftpClient, plan.SFTP.Dir)

	return fmt.Sprintf("SFTP %v:%v upload %v finished", plan.SFTP.Host, plan.SFTP.Port, dstPath), nil
}

func listSftpBackups(client *sftp.Client, dir string) error {
	list, err := client.ReadDir(fmt.Sprintf("/%v", dir))
	if err != nil {
		return errors.Wrapf(err, "SFTP reading %v dir failed", dir)
	}

	for _, item := range list {
		fmt.Println(item.Name())
	}

	return nil
}
<|endoftext|>


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

package gin

import (
	"net/http"
	"os"
)

type onlyFilesFS struct {
	fs http.FileSystem
}

type neuteredReaddirFile struct {
	http.File
}

// Dir returns a http.Filesystem that can be used by http.FileServer(). It is used internally
// in router.Static().
// if listDirectory == true, then it works the same as http.Dir() otherwise it returns
// a filesystem that prevents http.FileServer() to list the directory files.
func Dir(root string, listDirectory bool) http.FileSystem {
	fs := http.Dir(root)
	if listDirectory {
		return fs
	}
	return &onlyFilesFS{fs}
}

// Open conforms to http.Filesystem.
func (fs onlyFilesFS) Open(name string) (http.File, error) {
	f, err := fs.fs.Open(name)
	if err != nil {
		return nil, err
	}
	return neuteredReaddirFile{f}, nil
}

// Readdir overrides the http.File default implementation.
func (f neuteredReaddirFile) Readdir(count int) ([]os.FileInfo, error) {
	// this disables directory listing
	return nil, nil
}
<commit_msg>Fixed typo in documentation (#2733)<commit_after>// Copyright 2017 Manu Martinez-Almeida.  All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file.

package gin

import (
	"net/http"
	"os"
)

type onlyFilesFS struct {
	fs http.FileSystem
}

type neuteredReaddirFile struct {
	http.File
}

// Dir returns a http.FileSystem that can be used by http.FileServer(). It is used internally
// in router.Static().
// if listDirectory == true, then it works the same as http.Dir() otherwise it returns
// a filesystem that prevents http.FileServer() to list the directory files.
func Dir(root string, listDirectory bool) http.FileSystem {
	fs := http.Dir(root)
	if listDirectory {
		return fs
	}
	return &onlyFilesFS{fs}
}

// Open conforms to http.Filesystem.
func (fs onlyFilesFS) Open(name string) (http.File, error) {
	f, err := fs.fs.Open(name)
	if err != nil {
		return nil, err
	}
	return neuteredReaddirFile{f}, nil
}

// Readdir overrides the http.File default implementation.
func (f neuteredReaddirFile) Readdir(count int) ([]os.FileInfo, error) {
	// this disables directory listing
	return nil, nil
}
<|endoftext|>


<commit_before>// +build codegen

package api

import (
	"fmt"
	"regexp"
	"strings"
)

// updateTopLevelShapeReferences moves resultWrapper, locationName, and
// xmlNamespace traits from toplevel shape references to the toplevel
// shapes for easier code generation
func (a *API) updateTopLevelShapeReferences() {
	for _, o := range a.Operations {
		// these are for REST-XML services
		if o.InputRef.LocationName != "" {
			o.InputRef.Shape.LocationName = o.InputRef.LocationName
		}
		if o.InputRef.Location != "" {
			o.InputRef.Shape.Location = o.InputRef.Location
		}
		if o.InputRef.Payload != "" {
			o.InputRef.Shape.Payload = o.InputRef.Payload
		}
		if o.InputRef.XMLNamespace.Prefix != "" {
			o.InputRef.Shape.XMLNamespace.Prefix = o.InputRef.XMLNamespace.Prefix
		}
		if o.InputRef.XMLNamespace.URI != "" {
			o.InputRef.Shape.XMLNamespace.URI = o.InputRef.XMLNamespace.URI
		}
	}

}

// writeShapeNames sets each shape's API and shape name values. Binding the
// shape to its parent API.
func (a *API) writeShapeNames() {
	for n, s := range a.Shapes {
		s.API = a
		s.ShapeName = n
	}
}

func (a *API) resolveReferences() {
	resolver := referenceResolver{API: a, visited: map[*ShapeRef]bool{}}

	for _, s := range a.Shapes {
		resolver.resolveShape(s)
	}

	for _, o := range a.Operations {
		o.API = a // resolve parent reference

		resolver.resolveReference(&o.InputRef)
		resolver.resolveReference(&o.OutputRef)

		// Resolve references for errors also
		for i := range o.ErrorRefs {
			resolver.resolveReference(&o.ErrorRefs[i])
			o.ErrorRefs[i].Shape.IsError = true
		}
	}
}

// A referenceResolver provides a way to resolve shape references to
// shape definitions.
type referenceResolver struct {
	*API
	visited map[*ShapeRef]bool
}

var jsonvalueShape = &Shape{
	ShapeName: "JSONValue",
	Type:      "jsonvalue",
	ValueRef: ShapeRef{
		JSONValue: true,
	},
}

// resolveReference updates a shape reference to reference the API and
// its shape definition. All other nested references are also resolved.
func (r *referenceResolver) resolveReference(ref *ShapeRef) {
	if ref.ShapeName == "" {
		return
	}

	if shape, ok := r.API.Shapes[ref.ShapeName]; ok {
		if ref.JSONValue {
			ref.ShapeName = "JSONValue"
			r.API.Shapes[ref.ShapeName] = jsonvalueShape
		}

		ref.API = r.API   // resolve reference back to API
		ref.Shape = shape // resolve shape reference

		if r.visited[ref] {
			return
		}
		r.visited[ref] = true

		shape.refs = append(shape.refs, ref) // register the ref

		// resolve shape's references, if it has any
		r.resolveShape(shape)
	}
}

// resolveShape resolves a shape's Member Key Value, and nested member
// shape references.
func (r *referenceResolver) resolveShape(shape *Shape) {
	r.resolveReference(&shape.MemberRef)
	r.resolveReference(&shape.KeyRef)
	r.resolveReference(&shape.ValueRef)
	for _, m := range shape.MemberRefs {
		r.resolveReference(m)
	}
}

// renameToplevelShapes renames all top level shapes of an API to their
// exportable variant. The shapes are also updated to include notations
// if they are Input or Outputs.
func (a *API) renameToplevelShapes() {
	for _, v := range a.OperationList() {
		if v.HasInput() {
			name := v.ExportedName + "Input"
			switch {
			case a.Shapes[name] == nil:
				if service, ok := shamelist[a.name]; ok {
					if check, ok := service[v.Name]; ok && check.input {
						break
					}
				}
				v.InputRef.Shape.Rename(name)
			}
		}
		if v.HasOutput() {
			name := v.ExportedName + "Output"
			switch {
			case a.Shapes[name] == nil:
				if service, ok := shamelist[a.name]; ok {
					if check, ok := service[v.Name]; ok && check.output {
						break
					}
				}
				v.OutputRef.Shape.Rename(name)
			}
		}
		v.InputRef.Payload = a.ExportableName(v.InputRef.Payload)
		v.OutputRef.Payload = a.ExportableName(v.OutputRef.Payload)
	}
}

// fixStutterNames fixes all name struttering based on Go naming conventions.
// "Stuttering" is when the prefix of a structure or function matches the
// package name (case insensitive).
func (a *API) fixStutterNames() {
	str, end := a.StructName(), ""
	if len(str) > 1 {
		l := len(str) - 1
		str, end = str[0:l], str[l:]
	}
	re := regexp.MustCompile(fmt.Sprintf(`\A(?i:%s)%s`, str, end))

	for name, op := range a.Operations {
		newName := re.ReplaceAllString(name, "")
		if newName != name {
			delete(a.Operations, name)
			a.Operations[newName] = op
		}
		op.ExportedName = newName
	}

	for k, s := range a.Shapes {
		newName := re.ReplaceAllString(k, "")
		if newName != s.ShapeName {
			s.Rename(newName)
		}
	}
}

// renameExportable renames all operation names to be exportable names.
// All nested Shape names are also updated to the exportable variant.
func (a *API) renameExportable() {
	for name, op := range a.Operations {
		newName := a.ExportableName(name)
		if newName != name {
			delete(a.Operations, name)
			a.Operations[newName] = op
		}
		op.ExportedName = newName
	}

	for k, s := range a.Shapes {
		// FIXME SNS has lower and uppercased shape names with the same name,
		// except the lowercased variant is used exclusively for string and
		// other primitive types. Renaming both would cause a collision.
		// We work around this by only renaming the structure shapes.
		if s.Type == "string" {
			continue
		}

		for mName, member := range s.MemberRefs {
			newName := a.ExportableName(mName)
			if newName != mName {
				delete(s.MemberRefs, mName)
				s.MemberRefs[newName] = member

				// also apply locationName trait so we keep the old one
				// but only if there's no locationName trait on ref or shape
				if member.LocationName == "" && member.Shape.LocationName == "" {
					member.LocationName = mName
				}
			}

			if newName == "_" {
				panic("Shape " + s.ShapeName + " uses reserved member name '_'")
			}
		}

		newName := a.ExportableName(k)
		if newName != s.ShapeName {
			s.Rename(newName)
		}

		s.Payload = a.ExportableName(s.Payload)

		// fix required trait names
		for i, n := range s.Required {
			s.Required[i] = a.ExportableName(n)
		}
	}

	for _, s := range a.Shapes {
		// fix enum names
		if s.IsEnum() {
			s.EnumConsts = make([]string, len(s.Enum))
			for i := range s.Enum {
				shape := s.ShapeName
				shape = strings.ToUpper(shape[0:1]) + shape[1:]
				s.EnumConsts[i] = shape + s.EnumName(i)
			}
		}
	}
}

// createInputOutputShapes creates toplevel input/output shapes if they
// have not been defined in the API. This normalizes all APIs to always
// have an input and output structure in the signature.
func (a *API) createInputOutputShapes() {
	for _, op := range a.Operations {
		if !op.HasInput() {
			setAsPlacholderShape(&op.InputRef, op.ExportedName+"Input", a)
		}
		if !op.HasOutput() {
			setAsPlacholderShape(&op.OutputRef, op.ExportedName+"Output", a)
		}
	}
}

func setAsPlacholderShape(tgtShapeRef *ShapeRef, name string, a *API) {
	shape := a.makeIOShape(name)
	shape.Placeholder = true
	*tgtShapeRef = ShapeRef{API: a, ShapeName: shape.ShapeName, Shape: shape}
	shape.refs = append(shape.refs, tgtShapeRef)
}

// makeIOShape returns a pointer to a new Shape initialized by the name provided.
func (a *API) makeIOShape(name string) *Shape {
	shape := &Shape{
		API: a, ShapeName: name, Type: "structure",
		MemberRefs: map[string]*ShapeRef{},
	}
	a.Shapes[name] = shape
	return shape
}

// removeUnusedShapes removes shapes from the API which are not referenced by any
// other shape in the API.
func (a *API) removeUnusedShapes() {
	for n, s := range a.Shapes {
		if len(s.refs) == 0 {
			delete(a.Shapes, n)
		}
	}
}

// Represents the service package name to EndpointsID mapping
var custEndpointsKey = map[string]string{
	"applicationautoscaling": "application-autoscaling",
}

// Sents the EndpointsID field of Metadata  with the value of the
// EndpointPrefix if EndpointsID is not set. Also adds
// customizations for services if EndpointPrefix is not a valid key.
func (a *API) setMetadataEndpointsKey() {
	if len(a.Metadata.EndpointsID) != 0 {
		return
	}

	if v, ok := custEndpointsKey[a.PackageName()]; ok {
		a.Metadata.EndpointsID = v
	} else {
		a.Metadata.EndpointsID = a.Metadata.EndpointPrefix
	}
}
<commit_msg>codegen: Dont rename API type sutters if will cause empty name (#1643)<commit_after>// +build codegen

package api

import (
	"fmt"
	"regexp"
	"strings"
)

// updateTopLevelShapeReferences moves resultWrapper, locationName, and
// xmlNamespace traits from toplevel shape references to the toplevel
// shapes for easier code generation
func (a *API) updateTopLevelShapeReferences() {
	for _, o := range a.Operations {
		// these are for REST-XML services
		if o.InputRef.LocationName != "" {
			o.InputRef.Shape.LocationName = o.InputRef.LocationName
		}
		if o.InputRef.Location != "" {
			o.InputRef.Shape.Location = o.InputRef.Location
		}
		if o.InputRef.Payload != "" {
			o.InputRef.Shape.Payload = o.InputRef.Payload
		}
		if o.InputRef.XMLNamespace.Prefix != "" {
			o.InputRef.Shape.XMLNamespace.Prefix = o.InputRef.XMLNamespace.Prefix
		}
		if o.InputRef.XMLNamespace.URI != "" {
			o.InputRef.Shape.XMLNamespace.URI = o.InputRef.XMLNamespace.URI
		}
	}

}

// writeShapeNames sets each shape's API and shape name values. Binding the
// shape to its parent API.
func (a *API) writeShapeNames() {
	for n, s := range a.Shapes {
		s.API = a
		s.ShapeName = n
	}
}

func (a *API) resolveReferences() {
	resolver := referenceResolver{API: a, visited: map[*ShapeRef]bool{}}

	for _, s := range a.Shapes {
		resolver.resolveShape(s)
	}

	for _, o := range a.Operations {
		o.API = a // resolve parent reference

		resolver.resolveReference(&o.InputRef)
		resolver.resolveReference(&o.OutputRef)

		// Resolve references for errors also
		for i := range o.ErrorRefs {
			resolver.resolveReference(&o.ErrorRefs[i])
			o.ErrorRefs[i].Shape.IsError = true
		}
	}
}

// A referenceResolver provides a way to resolve shape references to
// shape definitions.
type referenceResolver struct {
	*API
	visited map[*ShapeRef]bool
}

var jsonvalueShape = &Shape{
	ShapeName: "JSONValue",
	Type:      "jsonvalue",
	ValueRef: ShapeRef{
		JSONValue: true,
	},
}

// resolveReference updates a shape reference to reference the API and
// its shape definition. All other nested references are also resolved.
func (r *referenceResolver) resolveReference(ref *ShapeRef) {
	if ref.ShapeName == "" {
		return
	}

	if shape, ok := r.API.Shapes[ref.ShapeName]; ok {
		if ref.JSONValue {
			ref.ShapeName = "JSONValue"
			r.API.Shapes[ref.ShapeName] = jsonvalueShape
		}

		ref.API = r.API   // resolve reference back to API
		ref.Shape = shape // resolve shape reference

		if r.visited[ref] {
			return
		}
		r.visited[ref] = true

		shape.refs = append(shape.refs, ref) // register the ref

		// resolve shape's references, if it has any
		r.resolveShape(shape)
	}
}

// resolveShape resolves a shape's Member Key Value, and nested member
// shape references.
func (r *referenceResolver) resolveShape(shape *Shape) {
	r.resolveReference(&shape.MemberRef)
	r.resolveReference(&shape.KeyRef)
	r.resolveReference(&shape.ValueRef)
	for _, m := range shape.MemberRefs {
		r.resolveReference(m)
	}
}

// renameToplevelShapes renames all top level shapes of an API to their
// exportable variant. The shapes are also updated to include notations
// if they are Input or Outputs.
func (a *API) renameToplevelShapes() {
	for _, v := range a.OperationList() {
		if v.HasInput() {
			name := v.ExportedName + "Input"
			switch {
			case a.Shapes[name] == nil:
				if service, ok := shamelist[a.name]; ok {
					if check, ok := service[v.Name]; ok && check.input {
						break
					}
				}
				v.InputRef.Shape.Rename(name)
			}
		}
		if v.HasOutput() {
			name := v.ExportedName + "Output"
			switch {
			case a.Shapes[name] == nil:
				if service, ok := shamelist[a.name]; ok {
					if check, ok := service[v.Name]; ok && check.output {
						break
					}
				}
				v.OutputRef.Shape.Rename(name)
			}
		}
		v.InputRef.Payload = a.ExportableName(v.InputRef.Payload)
		v.OutputRef.Payload = a.ExportableName(v.OutputRef.Payload)
	}
}

// fixStutterNames fixes all name struttering based on Go naming conventions.
// "Stuttering" is when the prefix of a structure or function matches the
// package name (case insensitive).
func (a *API) fixStutterNames() {
	str, end := a.StructName(), ""
	if len(str) > 1 {
		l := len(str) - 1
		str, end = str[0:l], str[l:]
	}
	re := regexp.MustCompile(fmt.Sprintf(`\A(?i:%s)%s`, str, end))

	for name, op := range a.Operations {
		newName := re.ReplaceAllString(name, "")
		if newName != name && len(newName) > 0 {
			delete(a.Operations, name)
			a.Operations[newName] = op
		}
		op.ExportedName = newName
	}

	for k, s := range a.Shapes {
		newName := re.ReplaceAllString(k, "")
		if newName != s.ShapeName && len(newName) > 0 {
			s.Rename(newName)
		}
	}
}

// renameExportable renames all operation names to be exportable names.
// All nested Shape names are also updated to the exportable variant.
func (a *API) renameExportable() {
	for name, op := range a.Operations {
		newName := a.ExportableName(name)
		if newName != name {
			delete(a.Operations, name)
			a.Operations[newName] = op
		}
		op.ExportedName = newName
	}

	for k, s := range a.Shapes {
		// FIXME SNS has lower and uppercased shape names with the same name,
		// except the lowercased variant is used exclusively for string and
		// other primitive types. Renaming both would cause a collision.
		// We work around this by only renaming the structure shapes.
		if s.Type == "string" {
			continue
		}

		for mName, member := range s.MemberRefs {
			newName := a.ExportableName(mName)
			if newName != mName {
				delete(s.MemberRefs, mName)
				s.MemberRefs[newName] = member

				// also apply locationName trait so we keep the old one
				// but only if there's no locationName trait on ref or shape
				if member.LocationName == "" && member.Shape.LocationName == "" {
					member.LocationName = mName
				}
			}

			if newName == "_" {
				panic("Shape " + s.ShapeName + " uses reserved member name '_'")
			}
		}

		newName := a.ExportableName(k)
		if newName != s.ShapeName {
			s.Rename(newName)
		}

		s.Payload = a.ExportableName(s.Payload)

		// fix required trait names
		for i, n := range s.Required {
			s.Required[i] = a.ExportableName(n)
		}
	}

	for _, s := range a.Shapes {
		// fix enum names
		if s.IsEnum() {
			s.EnumConsts = make([]string, len(s.Enum))
			for i := range s.Enum {
				shape := s.ShapeName
				shape = strings.ToUpper(shape[0:1]) + shape[1:]
				s.EnumConsts[i] = shape + s.EnumName(i)
			}
		}
	}
}

// createInputOutputShapes creates toplevel input/output shapes if they
// have not been defined in the API. This normalizes all APIs to always
// have an input and output structure in the signature.
func (a *API) createInputOutputShapes() {
	for _, op := range a.Operations {
		if !op.HasInput() {
			setAsPlacholderShape(&op.InputRef, op.ExportedName+"Input", a)
		}
		if !op.HasOutput() {
			setAsPlacholderShape(&op.OutputRef, op.ExportedName+"Output", a)
		}
	}
}

func setAsPlacholderShape(tgtShapeRef *ShapeRef, name string, a *API) {
	shape := a.makeIOShape(name)
	shape.Placeholder = true
	*tgtShapeRef = ShapeRef{API: a, ShapeName: shape.ShapeName, Shape: shape}
	shape.refs = append(shape.refs, tgtShapeRef)
}

// makeIOShape returns a pointer to a new Shape initialized by the name provided.
func (a *API) makeIOShape(name string) *Shape {
	shape := &Shape{
		API: a, ShapeName: name, Type: "structure",
		MemberRefs: map[string]*ShapeRef{},
	}
	a.Shapes[name] = shape
	return shape
}

// removeUnusedShapes removes shapes from the API which are not referenced by any
// other shape in the API.
func (a *API) removeUnusedShapes() {
	for n, s := range a.Shapes {
		if len(s.refs) == 0 {
			delete(a.Shapes, n)
		}
	}
}

// Represents the service package name to EndpointsID mapping
var custEndpointsKey = map[string]string{
	"applicationautoscaling": "application-autoscaling",
}

// Sents the EndpointsID field of Metadata  with the value of the
// EndpointPrefix if EndpointsID is not set. Also adds
// customizations for services if EndpointPrefix is not a valid key.
func (a *API) setMetadataEndpointsKey() {
	if len(a.Metadata.EndpointsID) != 0 {
		return
	}

	if v, ok := custEndpointsKey[a.PackageName()]; ok {
		a.Metadata.EndpointsID = v
	} else {
		a.Metadata.EndpointsID = a.Metadata.EndpointPrefix
	}
}
<|endoftext|>