<commit_before>// Copyright © 2016 Wei Shen <shenwei356@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package taxon

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/boltdb/bolt"
	"github.com/gin-gonic/gin"
	"github.com/parnurzeal/gorequest"
	"github.com/shenwei356/gtaxon/taxon/nodes"
)

var threadNum int

// StartServer runs a web server for query
func StartServer(dbFilePath string, port int, timeout int, threads int) {
	runtime.GOMAXPROCS(threads)
	threadNum = threads
	pool = NewDBPool(dbFilePath, threads)

	done := make(chan int)
	go func() {
		db := pool.GetDB()
		defer pool.ReleaseDB(db)

		log.Info("load all names ...")
		names, err := LoadAllNames(db, "names")
		checkError(err)
		nodes.SetNames(names)
		log.Info("load all names ... done")
		done <- 1
	}()

	done1 := make(chan int)
	go func() {
		db := pool.GetDB()
		defer pool.ReleaseDB(db)

		log.Info("load all nodes ...")
		nods, err := LoadAllNodes(db, "nodes")
		nodes.SetNodes(nods)
		checkError(err)
		log.Info("load all nodes ... done")

		done1 <- 1
	}()

	done2 := make(chan int)
	go func() {
		db := pool.GetDB()
		defer pool.ReleaseDB(db)

		log.Info("load all divisions ...")
		divisions, err := LoadAllDivisions(db, "divisions")
		nodes.SetDivisions(divisions)
		checkError(err)
		log.Info("load all divisions ... done")

		done2 <- 1
	}()

	done3 := make(chan int)
	go func() {
		db := pool.GetDB()
		defer pool.ReleaseDB(db)

		log.Info("load all gencodes ...")
		gencodes, err := LoadAllGenCodes(db, "gencodes")
		nodes.SetGenCodes(gencodes)
		checkError(err)
		log.Info("load all gencodes ... done")

		done3 <- 1
	}()

	<-done2
	<-done3
	<-done
	<-done1

	// gin.SetMode(gin.ReleaseMode)
	router := gin.Default()

	router.GET("/gi2taxid", gi2taxid)
	router.GET("/taxid2taxon", taxid2taxon)
	router.GET("/name2taxid", name2taxid)
	router.GET("/lca", lca)

	// router.Run(fmt.Sprintf(":%d", port))
	s := &http.Server{
		Addr:           fmt.Sprintf(":%d", port),
		Handler:        router,
		ReadTimeout:    2 * time.Second,
		WriteTimeout:   2 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}
	s.ListenAndServe()
}

// --------------------------------------------------------------------------

// MessageLCAMap is
type MessageLCAMap struct {
	Status  string `json:"status"`
	Message string `json:"message"`

	LCA map[string]nodes.Taxon `json:"taxids2taxon"`
}

func lca(c *gin.Context) {
	var msg MessageLCAMap

	c.Request.ParseForm()
	queries := c.Request.Form["taxids"]

	if queries == nil {
		msg.Status = "FAILED"
		msg.Message = "no Taxids given"
		c.JSON(http.StatusOK, msg)
		return
	}

	db := pool.GetDB()
	defer pool.ReleaseDB(db)

	result := make(map[string]nodes.Node)
	for _, query := range queries {
		taxids := strings.Split(query, ",")
		lca, err := nodes.LCA(nodes.Nodes, taxids)
		if err != nil {
			msg.Status = "FAILED"
			msg.Message = fmt.Sprintf("error: %s", err)
			c.JSON(http.StatusOK, msg)
			return
		}
		result[query] = lca
	}

	msg.Status = "OK"
	msg.Message = fmt.Sprintf("sum: %d", len(queries))
	msg.LCA = make(map[string]nodes.Taxon, len(result))
	for k, node := range result {
		msg.LCA[k], _ = nodes.GetTaxonByTaxID(node.TaxID)
	}

	c.JSON(http.StatusOK, msg)
}

// RemoteQueryLCA is
func RemoteQueryLCA(host string, port int, queries []string) MessageLCAMap {
	host = strings.TrimSpace(host)
	var url string
	if regexp.MustCompile("^http://").MatchString(host) {
		url = fmt.Sprintf("%s:%d/lca", host, port)
	} else {
		url = fmt.Sprintf("http://%s:%d/lca", host, port)
	}

	request := gorequest.New().Get(url)

	for _, query := range queries {
		request = request.Param("taxids", query)
	}

	_, body, errs := request.End()
	if errs != nil {
		log.Error(errs)
		os.Exit(-1)
	}

	var result MessageLCAMap
	err := json.Unmarshal([]byte(body), &result)
	checkError(err)

	return result
}

// --------------------------------------------------------------------------

// MessageTaxid2TaxonMap is
type MessageTaxid2TaxonMap struct {
	Status  string `json:"status"`
	Message string `json:"message"`

	Taxons map[string]nodes.Taxon `json:"taxid2taxon"`
}

func taxid2taxon(c *gin.Context) {
	var msg MessageTaxid2TaxonMap

	c.Request.ParseForm()
	taxids := c.Request.Form["taxid"]

	if taxids == nil {
		msg.Status = "FAILED"
		msg.Message = "no Taxids given"
		c.JSON(http.StatusOK, msg)
		return
	}

	db := pool.GetDB()
	defer pool.ReleaseDB(db)

	bucket := "nodes"
	nods, err := QueryNodeByTaxID(db, bucket, taxids)
	if err != nil {
		msg.Status = "FAILED"
		msg.Message = fmt.Sprintf("error: %s", err)
		c.JSON(http.StatusOK, msg)
		return
	}
	msg.Status = "OK"
	msg.Message = fmt.Sprintf("sum: %d", len(nods))
	msg.Taxons = make(map[string]nodes.Taxon)
	for _, node := range nods {
		taxon, _ := nodes.GetTaxonByTaxID(node.TaxID)
		msg.Taxons[node.TaxID] = taxon
	}
	c.JSON(http.StatusOK, msg)
}

// RemoteQueryTaxid2Taxon is
func RemoteQueryTaxid2Taxon(host string, port int, taxids []string) MessageTaxid2TaxonMap {
	host = strings.TrimSpace(host)
	var url string
	if regexp.MustCompile("^http://").MatchString(host) {
		url = fmt.Sprintf("%s:%d/taxid2taxon", host, port)
	} else {
		url = fmt.Sprintf("http://%s:%d/taxid2taxon", host, port)
	}

	request := gorequest.New().Get(url)

	for _, taxid := range taxids {
		request = request.Param("taxid", taxid)
	}

	_, body, errs := request.End()
	if errs != nil {
		log.Error(errs)
		os.Exit(-1)
	}

	var result MessageTaxid2TaxonMap
	err := json.Unmarshal([]byte(body), &result)
	checkError(err)

	return result
}

// --------------------------------------------------------------------------

// MssageName2TaxIDMap is
type MssageName2TaxIDMap struct {
	Status  string `json:"status"`
	Message string `json:"message"`

	TaxIDs map[string][]TaxIDSciNameItem `json:"name2taxid"`
}

// TaxIDSciNameItem is
type TaxIDSciNameItem struct {
	TaxID          int
	ScientificName string
}

func name2taxid(c *gin.Context) {
	var msg MssageName2TaxIDMap

	c.Request.ParseForm()
	useRegexp := false
	if c.Query("regexp") != "" {
		useRegexp = true
	}
	nameClass := ""
	if c.Query("class") != "" {
		nameClass = c.Query("class")
	}
	names := c.Request.Form["name"]

	if names == nil {
		msg.Status = "FAILED"
		msg.Message = "no names given"
		c.JSON(http.StatusOK, msg)
		return
	}

	db := pool.GetDB()
	defer pool.ReleaseDB(db)

	bucket := "names"
	results, err := QueryTaxIDByName(db, bucket, useRegexp, nameClass, threadNum*2, names)
	if err != nil {
		msg.Status = "FAILED"
		msg.Message = fmt.Sprintf("error: %s", err)
		c.JSON(http.StatusOK, msg)
		return
	}
	msg.Status = "OK"
	msg.Message = fmt.Sprintf("sum: %d", len(names))

	name2taxidResults := make(map[string][]TaxIDSciNameItem, len(results))
	for name, taxids := range results {
		taxidsSciNameItems := make([]TaxIDSciNameItem, len(taxids))
		for i, taxid := range taxids {
			node := nodes.Nodes[taxid]
			taxidInt, _ := strconv.Atoi(node.TaxID)
			scientificName := ""
			for _, nameItem := range nodes.Names[node.TaxID].Names {
				if nameItem.NameClass == "scientific name" {
					scientificName = nameItem.Name
				}
			}
			taxidsSciNameItems[i] = TaxIDSciNameItem{
				TaxID:          taxidInt,
				ScientificName: scientificName,
			}
		}
		name2taxidResults[name] = taxidsSciNameItems
	}
	msg.TaxIDs = name2taxidResults
	c.JSON(http.StatusOK, msg)
}

// RemoteQueryName2TaxID is
func RemoteQueryName2TaxID(host string, port int, useRegexp bool, nameClass string, names []string) MssageName2TaxIDMap {
	host = strings.TrimSpace(host)
	var url string
	if regexp.MustCompile("^http://").MatchString(host) {
		url = fmt.Sprintf("%s:%d/name2taxid", host, port)
	} else {
		url = fmt.Sprintf("http://%s:%d/name2taxid", host, port)
	}

	request := gorequest.New().Get(url)
	if useRegexp {
		request = request.Param("regexp", "1")
	}
	if nameClass != "" {
		request = request.Param("class", nameClass)
	}

	for _, name := range names {
		request = request.Param("name", name)
	}

	_, body, errs := request.End()
	if errs != nil {
		log.Error(errs)
		os.Exit(-1)
	}

	var result MssageName2TaxIDMap
	err := json.Unmarshal([]byte(body), &result)
	checkError(err)
	return result
}

// --------------------------------------------------------------------------

// MessageGI2TaxidMap is
type MessageGI2TaxidMap struct {
	Status  string `json:"status"`
	Message string `json:"message"`

	Taxids map[string]string `json:"gi2taxid"`
}

func gi2taxid(c *gin.Context) {
	var msg MessageGI2TaxidMap

	// gi := c.Query("gi")  // single value
	// multiple values
	c.Request.ParseForm()
	gis := c.Request.Form["gi"]

	if gis == nil {
		msg.Status = "FAILED"
		msg.Message = "no GIs given"
		c.JSON(http.StatusOK, msg)
		return
	}

	bucket := c.DefaultQuery("db", "gi_taxid_prot")
	if bucket != "gi_taxid_prot" && bucket != "gi_taxid_nucl" {
		msg.Status = "FAILED"
		msg.Message = fmt.Sprintf("invalid db: %s. valid: gi_taxid_prot or gi_taxid_nucl", bucket)
		c.JSON(http.StatusOK, msg)
		return
	}

	db := pool.GetDB()
	defer pool.ReleaseDB(db)

	taxids := make(map[string]string, len(gis))
	n := 0 // counter of seccessful query
	err := db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if b == nil {
			return fmt.Errorf("db not found: %s", bucket)
		}
		for _, gi := range gis {
			taxid := string(b.Get([]byte(gi)))
			if taxid != "" {
				n++
			}
			taxids[gi] = taxid
		}
		return nil
	})
	if err != nil {
		msg.Status = "FAILED"
		msg.Message = fmt.Sprintf("error: %s", err)
		c.JSON(http.StatusOK, msg)
		return
	}
	msg.Status = "OK"
	msg.Message = fmt.Sprintf("sum: %d, found: %d", len(gis), n)
	msg.Taxids = taxids
	c.JSON(http.StatusOK, msg)
}

// RemoteQueryGi2Taxid query from remote server
func RemoteQueryGi2Taxid(host string, port int, dbType string, gis []string) MessageGI2TaxidMap {
	host = strings.TrimSpace(host)
	var url string
	if regexp.MustCompile("^http://").MatchString(host) {
		url = fmt.Sprintf("%s:%d/gi2taxid", host, port)
	} else {
		url = fmt.Sprintf("http://%s:%d/gi2taxid", host, port)
	}

	request := gorequest.New().Get(url).Param("db", dbType)

	for _, gi := range gis {
		request = request.Param("gi", gi)
	}

	_, body, errs := request.End()
	if errs != nil {
		log.Error(errs)
		os.Exit(-1)
	}

	var result MessageGI2TaxidMap
	err := json.Unmarshal([]byte(body), &result)
	checkError(err)

	return result
}
<commit_msg>change gin web server mode to release<commit_after>// Copyright © 2016 Wei Shen <shenwei356@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package taxon

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/boltdb/bolt"
	"github.com/gin-gonic/gin"
	"github.com/parnurzeal/gorequest"
	"github.com/shenwei356/gtaxon/taxon/nodes"
)

var threadNum int

// StartServer runs a web server for query
func StartServer(dbFilePath string, port int, timeout int, threads int) {
	runtime.GOMAXPROCS(threads)
	threadNum = threads
	pool = NewDBPool(dbFilePath, threads)

	done := make(chan int)
	go func() {
		db := pool.GetDB()
		defer pool.ReleaseDB(db)

		log.Info("load all names ...")
		names, err := LoadAllNames(db, "names")
		checkError(err)
		nodes.SetNames(names)
		log.Info("load all names ... done")
		done <- 1
	}()

	done1 := make(chan int)
	go func() {
		db := pool.GetDB()
		defer pool.ReleaseDB(db)

		log.Info("load all nodes ...")
		nods, err := LoadAllNodes(db, "nodes")
		nodes.SetNodes(nods)
		checkError(err)
		log.Info("load all nodes ... done")

		done1 <- 1
	}()

	done2 := make(chan int)
	go func() {
		db := pool.GetDB()
		defer pool.ReleaseDB(db)

		log.Info("load all divisions ...")
		divisions, err := LoadAllDivisions(db, "divisions")
		nodes.SetDivisions(divisions)
		checkError(err)
		log.Info("load all divisions ... done")

		done2 <- 1
	}()

	done3 := make(chan int)
	go func() {
		db := pool.GetDB()
		defer pool.ReleaseDB(db)

		log.Info("load all gencodes ...")
		gencodes, err := LoadAllGenCodes(db, "gencodes")
		nodes.SetGenCodes(gencodes)
		checkError(err)
		log.Info("load all gencodes ... done")

		done3 <- 1
	}()

	<-done2
	<-done3
	<-done
	<-done1

	gin.SetMode(gin.ReleaseMode)
	router := gin.Default()

	router.GET("/gi2taxid", gi2taxid)
	router.GET("/taxid2taxon", taxid2taxon)
	router.GET("/name2taxid", name2taxid)
	router.GET("/lca", lca)

	// router.Run(fmt.Sprintf(":%d", port))
	s := &http.Server{
		Addr:           fmt.Sprintf(":%d", port),
		Handler:        router,
		ReadTimeout:    2 * time.Second,
		WriteTimeout:   2 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}
	s.ListenAndServe()
}

// --------------------------------------------------------------------------

// MessageLCAMap is
type MessageLCAMap struct {
	Status  string `json:"status"`
	Message string `json:"message"`

	LCA map[string]nodes.Taxon `json:"taxids2taxon"`
}

func lca(c *gin.Context) {
	var msg MessageLCAMap

	c.Request.ParseForm()
	queries := c.Request.Form["taxids"]

	if queries == nil {
		msg.Status = "FAILED"
		msg.Message = "no Taxids given"
		c.JSON(http.StatusOK, msg)
		return
	}

	db := pool.GetDB()
	defer pool.ReleaseDB(db)

	result := make(map[string]nodes.Node)
	for _, query := range queries {
		taxids := strings.Split(query, ",")
		lca, err := nodes.LCA(nodes.Nodes, taxids)
		if err != nil {
			msg.Status = "FAILED"
			msg.Message = fmt.Sprintf("error: %s", err)
			c.JSON(http.StatusOK, msg)
			return
		}
		result[query] = lca
	}

	msg.Status = "OK"
	msg.Message = fmt.Sprintf("sum: %d", len(queries))
	msg.LCA = make(map[string]nodes.Taxon, len(result))
	for k, node := range result {
		msg.LCA[k], _ = nodes.GetTaxonByTaxID(node.TaxID)
	}

	c.JSON(http.StatusOK, msg)
}

// RemoteQueryLCA is
func RemoteQueryLCA(host string, port int, queries []string) MessageLCAMap {
	host = strings.TrimSpace(host)
	var url string
	if regexp.MustCompile("^http://").MatchString(host) {
		url = fmt.Sprintf("%s:%d/lca", host, port)
	} else {
		url = fmt.Sprintf("http://%s:%d/lca", host, port)
	}

	request := gorequest.New().Get(url)

	for _, query := range queries {
		request = request.Param("taxids", query)
	}

	_, body, errs := request.End()
	if errs != nil {
		log.Error(errs)
		os.Exit(-1)
	}

	var result MessageLCAMap
	err := json.Unmarshal([]byte(body), &result)
	checkError(err)

	return result
}

// --------------------------------------------------------------------------

// MessageTaxid2TaxonMap is
type MessageTaxid2TaxonMap struct {
	Status  string `json:"status"`
	Message string `json:"message"`

	Taxons map[string]nodes.Taxon `json:"taxid2taxon"`
}

func taxid2taxon(c *gin.Context) {
	var msg MessageTaxid2TaxonMap

	c.Request.ParseForm()
	taxids := c.Request.Form["taxid"]

	if taxids == nil {
		msg.Status = "FAILED"
		msg.Message = "no Taxids given"
		c.JSON(http.StatusOK, msg)
		return
	}

	db := pool.GetDB()
	defer pool.ReleaseDB(db)

	bucket := "nodes"
	nods, err := QueryNodeByTaxID(db, bucket, taxids)
	if err != nil {
		msg.Status = "FAILED"
		msg.Message = fmt.Sprintf("error: %s", err)
		c.JSON(http.StatusOK, msg)
		return
	}
	msg.Status = "OK"
	msg.Message = fmt.Sprintf("sum: %d", len(nods))
	msg.Taxons = make(map[string]nodes.Taxon)
	for _, node := range nods {
		taxon, _ := nodes.GetTaxonByTaxID(node.TaxID)
		msg.Taxons[node.TaxID] = taxon
	}
	c.JSON(http.StatusOK, msg)
}

// RemoteQueryTaxid2Taxon is
func RemoteQueryTaxid2Taxon(host string, port int, taxids []string) MessageTaxid2TaxonMap {
	host = strings.TrimSpace(host)
	var url string
	if regexp.MustCompile("^http://").MatchString(host) {
		url = fmt.Sprintf("%s:%d/taxid2taxon", host, port)
	} else {
		url = fmt.Sprintf("http://%s:%d/taxid2taxon", host, port)
	}

	request := gorequest.New().Get(url)

	for _, taxid := range taxids {
		request = request.Param("taxid", taxid)
	}

	_, body, errs := request.End()
	if errs != nil {
		log.Error(errs)
		os.Exit(-1)
	}

	var result MessageTaxid2TaxonMap
	err := json.Unmarshal([]byte(body), &result)
	checkError(err)

	return result
}

// --------------------------------------------------------------------------

// MssageName2TaxIDMap is
type MssageName2TaxIDMap struct {
	Status  string `json:"status"`
	Message string `json:"message"`

	TaxIDs map[string][]TaxIDSciNameItem `json:"name2taxid"`
}

// TaxIDSciNameItem is
type TaxIDSciNameItem struct {
	TaxID          int
	ScientificName string
}

func name2taxid(c *gin.Context) {
	var msg MssageName2TaxIDMap

	c.Request.ParseForm()
	useRegexp := false
	if c.Query("regexp") != "" {
		useRegexp = true
	}
	nameClass := ""
	if c.Query("class") != "" {
		nameClass = c.Query("class")
	}
	names := c.Request.Form["name"]

	if names == nil {
		msg.Status = "FAILED"
		msg.Message = "no names given"
		c.JSON(http.StatusOK, msg)
		return
	}

	db := pool.GetDB()
	defer pool.ReleaseDB(db)

	bucket := "names"
	results, err := QueryTaxIDByName(db, bucket, useRegexp, nameClass, threadNum*2, names)
	if err != nil {
		msg.Status = "FAILED"
		msg.Message = fmt.Sprintf("error: %s", err)
		c.JSON(http.StatusOK, msg)
		return
	}
	msg.Status = "OK"
	msg.Message = fmt.Sprintf("sum: %d", len(names))

	name2taxidResults := make(map[string][]TaxIDSciNameItem, len(results))
	for name, taxids := range results {
		taxidsSciNameItems := make([]TaxIDSciNameItem, len(taxids))
		for i, taxid := range taxids {
			node := nodes.Nodes[taxid]
			taxidInt, _ := strconv.Atoi(node.TaxID)
			scientificName := ""
			for _, nameItem := range nodes.Names[node.TaxID].Names {
				if nameItem.NameClass == "scientific name" {
					scientificName = nameItem.Name
				}
			}
			taxidsSciNameItems[i] = TaxIDSciNameItem{
				TaxID:          taxidInt,
				ScientificName: scientificName,
			}
		}
		name2taxidResults[name] = taxidsSciNameItems
	}
	msg.TaxIDs = name2taxidResults
	c.JSON(http.StatusOK, msg)
}

// RemoteQueryName2TaxID is
func RemoteQueryName2TaxID(host string, port int, useRegexp bool, nameClass string, names []string) MssageName2TaxIDMap {
	host = strings.TrimSpace(host)
	var url string
	if regexp.MustCompile("^http://").MatchString(host) {
		url = fmt.Sprintf("%s:%d/name2taxid", host, port)
	} else {
		url = fmt.Sprintf("http://%s:%d/name2taxid", host, port)
	}

	request := gorequest.New().Get(url)
	if useRegexp {
		request = request.Param("regexp", "1")
	}
	if nameClass != "" {
		request = request.Param("class", nameClass)
	}

	for _, name := range names {
		request = request.Param("name", name)
	}

	_, body, errs := request.End()
	if errs != nil {
		log.Error(errs)
		os.Exit(-1)
	}

	var result MssageName2TaxIDMap
	err := json.Unmarshal([]byte(body), &result)
	checkError(err)
	return result
}

// --------------------------------------------------------------------------

// MessageGI2TaxidMap is
type MessageGI2TaxidMap struct {
	Status  string `json:"status"`
	Message string `json:"message"`

	Taxids map[string]string `json:"gi2taxid"`
}

func gi2taxid(c *gin.Context) {
	var msg MessageGI2TaxidMap

	// gi := c.Query("gi")  // single value
	// multiple values
	c.Request.ParseForm()
	gis := c.Request.Form["gi"]

	if gis == nil {
		msg.Status = "FAILED"
		msg.Message = "no GIs given"
		c.JSON(http.StatusOK, msg)
		return
	}

	bucket := c.DefaultQuery("db", "gi_taxid_prot")
	if bucket != "gi_taxid_prot" && bucket != "gi_taxid_nucl" {
		msg.Status = "FAILED"
		msg.Message = fmt.Sprintf("invalid db: %s. valid: gi_taxid_prot or gi_taxid_nucl", bucket)
		c.JSON(http.StatusOK, msg)
		return
	}

	db := pool.GetDB()
	defer pool.ReleaseDB(db)

	taxids := make(map[string]string, len(gis))
	n := 0 // counter of seccessful query
	err := db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if b == nil {
			return fmt.Errorf("db not found: %s", bucket)
		}
		for _, gi := range gis {
			taxid := string(b.Get([]byte(gi)))
			if taxid != "" {
				n++
			}
			taxids[gi] = taxid
		}
		return nil
	})
	if err != nil {
		msg.Status = "FAILED"
		msg.Message = fmt.Sprintf("error: %s", err)
		c.JSON(http.StatusOK, msg)
		return
	}
	msg.Status = "OK"
	msg.Message = fmt.Sprintf("sum: %d, found: %d", len(gis), n)
	msg.Taxids = taxids
	c.JSON(http.StatusOK, msg)
}

// RemoteQueryGi2Taxid query from remote server
func RemoteQueryGi2Taxid(host string, port int, dbType string, gis []string) MessageGI2TaxidMap {
	host = strings.TrimSpace(host)
	var url string
	if regexp.MustCompile("^http://").MatchString(host) {
		url = fmt.Sprintf("%s:%d/gi2taxid", host, port)
	} else {
		url = fmt.Sprintf("http://%s:%d/gi2taxid", host, port)
	}

	request := gorequest.New().Get(url).Param("db", dbType)

	for _, gi := range gis {
		request = request.Param("gi", gi)
	}

	_, body, errs := request.End()
	if errs != nil {
		log.Error(errs)
		os.Exit(-1)
	}

	var result MessageGI2TaxidMap
	err := json.Unmarshal([]byte(body), &result)
	checkError(err)

	return result
}
<|endoftext|>


<commit_before>// Package auth is a plugin for verifying peer at the first time.
//
// Copyright 2017 HenryLee. All Rights Reserved.
//
// 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 auth

import (
	"fmt"
	"net"
	"sync/atomic"

	"github.com/henrylee2cn/goutil"
	tp "github.com/henrylee2cn/teleport"
)

// NewBearerPlugin creates a auth bearer plugin for client.
func NewBearerPlugin(fn Bearer, infoSetting ...tp.MessageSetting) tp.Plugin {
	return &authBearerPlugin{
		bearerFunc: fn,
		msgSetting: infoSetting,
	}
}

// NewCheckerPlugin creates a auth checker plugin for server.
func NewCheckerPlugin(fn Checker, retSetting ...tp.MessageSetting) tp.Plugin {
	return &authCheckerPlugin{
		checkerFunc: fn,
		msgSetting:  retSetting,
	}
}

type (
	// Bearer initiates an authorization request and handles the response.
	Bearer func(sess Session, fn SendOnce) *tp.Status
	// SendOnce sends authorization request once.
	SendOnce func(info, retRecv interface{}) *tp.Status

	// Checker checks the authorization request.
	Checker func(sess Session, fn RecvOnce) (ret interface{}, stat *tp.Status)
	// RecvOnce receives authorization request once.
	RecvOnce func(infoRecv interface{}) *tp.Status

	// Session auth session provides SetID, RemoteAddr and Swap methods in base session
	Session interface {
		// Peer returns the peer.
		Peer() tp.Peer
		// SetID sets the session id.
		SetID(newID string)
		// RemoteAddr returns the remote network address.
		RemoteAddr() net.Addr
		// Swap returns custom data swap of the session(socket).
		Swap() goutil.Map
	}
)

type authBearerPlugin struct {
	bearerFunc Bearer
	msgSetting []tp.MessageSetting
}

type authCheckerPlugin struct {
	checkerFunc Checker
	msgSetting  []tp.MessageSetting
}

var (
	_ tp.PostDialPlugin   = new(authBearerPlugin)
	_ tp.PostAcceptPlugin = new(authCheckerPlugin)
)

func (a *authBearerPlugin) Name() string {
	return "auth-bearer"
}

func (a *authCheckerPlugin) Name() string {
	return "auth-checker"
}

// MultiSendErr the error of multiple call SendOnce function
var MultiSendErr = tp.NewStatus(
	tp.CodeWriteFailed,
	"auth-bearer plugin usage is incorrect",
	"multiple call SendOnce function",
)

// MultiRecvErr the error of multiple call RecvOnce function
var MultiRecvErr = tp.NewStatus(
	tp.CodeInternalServerError,
	"auth-checker plugin usage is incorrect",
	"multiple call RecvOnce function",
)

func (a *authBearerPlugin) PostDial(sess tp.PreSession, _ bool) *tp.Status {
	if a.bearerFunc == nil {
		return nil
	}
	var called int32
	return a.bearerFunc(sess, func(info, retRecv interface{}) *tp.Status {
		if !atomic.CompareAndSwapInt32(&called, 0, 1) {
			return MultiSendErr
		}
		stat := sess.PreSend(tp.TypeAuthCall, "", info, nil, a.msgSetting...)
		if !stat.OK() {
			return stat
		}
		retMsg := sess.PreReceive(func(header tp.Header) interface{} {
			if header.Mtype() != tp.TypeAuthReply {
				return nil
			}
			return retRecv
		})
		if !retMsg.StatusOK() {
			return retMsg.Status()
		}
		if retMsg.Mtype() != tp.TypeAuthReply {
			return tp.NewStatus(
				tp.CodeUnauthorized,
				tp.CodeText(tp.CodeUnauthorized),
				fmt.Sprintf("auth message(1st) expect: AUTH_REPLY, but received: %s",
					tp.TypeText(retMsg.Mtype())),
			)
		}
		return nil
	})
}

func (a *authCheckerPlugin) PostAccept(sess tp.PreSession) *tp.Status {
	if a.checkerFunc == nil {
		return nil
	}
	var called int32
	ret, stat := a.checkerFunc(sess, func(infoRecv interface{}) *tp.Status {
		if !atomic.CompareAndSwapInt32(&called, 0, 1) {
			return MultiRecvErr
		}
		infoMsg := sess.PreReceive(func(header tp.Header) interface{} {
			if header.Mtype() != tp.TypeAuthCall {
				return nil
			}
			return infoRecv
		})
		if !infoMsg.StatusOK() {
			return infoMsg.Status()
		}
		if infoMsg.Mtype() != tp.TypeAuthCall {
			return tp.NewStatus(
				tp.CodeUnauthorized,
				tp.CodeText(tp.CodeUnauthorized),
				fmt.Sprintf("auth message(1st) expect: AUTH_CALL, but received: %s",
					tp.TypeText(infoMsg.Mtype())),
			)
		}
		return nil
	})
	if stat == MultiRecvErr {
		sess.PreSend(tp.TypeAuthReply, "", nil, stat, a.msgSetting...)
		return stat
	}
	stat2 := sess.PreSend(tp.TypeAuthReply, "", ret, stat, a.msgSetting...)
	if !stat2.OK() {
		return stat2
	}
	return stat
}
<commit_msg>feat(auth): Add Session.LocalAddr #70<commit_after>// Package auth is a plugin for verifying peer at the first time.
//
// Copyright 2017 HenryLee. All Rights Reserved.
//
// 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 auth

import (
	"fmt"
	"net"
	"sync/atomic"

	"github.com/henrylee2cn/goutil"
	tp "github.com/henrylee2cn/teleport"
)

// NewBearerPlugin creates a auth bearer plugin for client.
func NewBearerPlugin(fn Bearer, infoSetting ...tp.MessageSetting) tp.Plugin {
	return &authBearerPlugin{
		bearerFunc: fn,
		msgSetting: infoSetting,
	}
}

// NewCheckerPlugin creates a auth checker plugin for server.
func NewCheckerPlugin(fn Checker, retSetting ...tp.MessageSetting) tp.Plugin {
	return &authCheckerPlugin{
		checkerFunc: fn,
		msgSetting:  retSetting,
	}
}

type (
	// Bearer initiates an authorization request and handles the response.
	Bearer func(sess Session, fn SendOnce) *tp.Status
	// SendOnce sends authorization request once.
	SendOnce func(info, retRecv interface{}) *tp.Status

	// Checker checks the authorization request.
	Checker func(sess Session, fn RecvOnce) (ret interface{}, stat *tp.Status)
	// RecvOnce receives authorization request once.
	RecvOnce func(infoRecv interface{}) *tp.Status

	// Session auth session provides SetID, RemoteAddr and Swap methods in base session
	Session interface {
		// Peer returns the peer.
		Peer() tp.Peer
		// SetID sets the session id.
		SetID(newID string)
		// LocalAddr returns the local network address.
		LocalAddr() net.Addr
		// RemoteAddr returns the remote network address.
		RemoteAddr() net.Addr
		// Swap returns custom data swap of the session(socket).
		Swap() goutil.Map
	}
)

type authBearerPlugin struct {
	bearerFunc Bearer
	msgSetting []tp.MessageSetting
}

type authCheckerPlugin struct {
	checkerFunc Checker
	msgSetting  []tp.MessageSetting
}

var (
	_ tp.PostDialPlugin   = new(authBearerPlugin)
	_ tp.PostAcceptPlugin = new(authCheckerPlugin)
)

func (a *authBearerPlugin) Name() string {
	return "auth-bearer"
}

func (a *authCheckerPlugin) Name() string {
	return "auth-checker"
}

// MultiSendErr the error of multiple call SendOnce function
var MultiSendErr = tp.NewStatus(
	tp.CodeWriteFailed,
	"auth-bearer plugin usage is incorrect",
	"multiple call SendOnce function",
)

// MultiRecvErr the error of multiple call RecvOnce function
var MultiRecvErr = tp.NewStatus(
	tp.CodeInternalServerError,
	"auth-checker plugin usage is incorrect",
	"multiple call RecvOnce function",
)

func (a *authBearerPlugin) PostDial(sess tp.PreSession, _ bool) *tp.Status {
	if a.bearerFunc == nil {
		return nil
	}
	var called int32
	return a.bearerFunc(sess, func(info, retRecv interface{}) *tp.Status {
		if !atomic.CompareAndSwapInt32(&called, 0, 1) {
			return MultiSendErr
		}
		stat := sess.PreSend(tp.TypeAuthCall, "", info, nil, a.msgSetting...)
		if !stat.OK() {
			return stat
		}
		retMsg := sess.PreReceive(func(header tp.Header) interface{} {
			if header.Mtype() != tp.TypeAuthReply {
				return nil
			}
			return retRecv
		})
		if !retMsg.StatusOK() {
			return retMsg.Status()
		}
		if retMsg.Mtype() != tp.TypeAuthReply {
			return tp.NewStatus(
				tp.CodeUnauthorized,
				tp.CodeText(tp.CodeUnauthorized),
				fmt.Sprintf("auth message(1st) expect: AUTH_REPLY, but received: %s",
					tp.TypeText(retMsg.Mtype())),
			)
		}
		return nil
	})
}

func (a *authCheckerPlugin) PostAccept(sess tp.PreSession) *tp.Status {
	if a.checkerFunc == nil {
		return nil
	}
	var called int32
	ret, stat := a.checkerFunc(sess, func(infoRecv interface{}) *tp.Status {
		if !atomic.CompareAndSwapInt32(&called, 0, 1) {
			return MultiRecvErr
		}
		infoMsg := sess.PreReceive(func(header tp.Header) interface{} {
			if header.Mtype() != tp.TypeAuthCall {
				return nil
			}
			return infoRecv
		})
		if !infoMsg.StatusOK() {
			return infoMsg.Status()
		}
		if infoMsg.Mtype() != tp.TypeAuthCall {
			return tp.NewStatus(
				tp.CodeUnauthorized,
				tp.CodeText(tp.CodeUnauthorized),
				fmt.Sprintf("auth message(1st) expect: AUTH_CALL, but received: %s",
					tp.TypeText(infoMsg.Mtype())),
			)
		}
		return nil
	})
	if stat == MultiRecvErr {
		sess.PreSend(tp.TypeAuthReply, "", nil, stat, a.msgSetting...)
		return stat
	}
	stat2 := sess.PreSend(tp.TypeAuthReply, "", ret, stat, a.msgSetting...)
	if !stat2.OK() {
		return stat2
	}
	return stat
}
<|endoftext|>


<commit_before>// passvault_test: tests for passvault.go
//
// Copyright (c) 2013 CloudFlare, Inc.

package passvault

import (
	"os"
	"testing"
)

func TestStaticVault(t *testing.T) {
	// Creates a temporary on-disk database to test if passvault can read and
	// write from/to disk.  It's deleted at the bottom of the function--this
	// should be the only test that requires touching disk.
	records, err := InitFrom("/tmp/redoctober.json")
	if err != nil {
		t.Fatalf("Error reading record")
	}

	_, err = records.AddNewRecord("test", "bad pass", true, DefaultRecordType)
	if err != nil {
		t.Fatalf("Error creating record")
	}

	// Reads data written last time.
	_, err = InitFrom("/tmp/redoctober.json")
	if err != nil {
		t.Fatalf("Error reading record")
	}

	os.Remove("/tmp/redoctober.json")
}

func TestRSAEncryptDecrypt(t *testing.T) {
	records, err := InitFrom("memory")
	if err != nil {
		t.Fatalf("%v", err)
	}

	myRec, err := records.AddNewRecord("user", "weakpassword", true, RSARecord)
	if err != nil {
		t.Fatalf("%v", err)
	}

	_, err = myRec.GetKeyRSAPub()
	if err != nil {
		t.Fatalf("Error extracting RSA Pub")
	}

	rsaPriv, err := myRec.GetKeyRSA("mypasswordiswrong")
	if err == nil {
		t.Fatalf("Incorrect password did not fail")
	}

	rsaPriv, err = myRec.GetKeyRSA("weakpassword")
	if err != nil {
		t.Fatalf("Error decrypting RSA key")
	}

	err = rsaPriv.Validate()
	if err != nil {
		t.Fatalf("Error validating RSA key")
	}
}

func TestECCEncryptDecrypt(t *testing.T) {
	records, err := InitFrom("memory")
	if err != nil {
		t.Fatalf("%v", err)
	}

	myRec, err := records.AddNewRecord("user", "weakpassword", true, ECCRecord)
	if err != nil {
		t.Fatalf("%v", err)
	}

	_, err = myRec.GetKeyECCPub()
	if err != nil {
		t.Fatalf("%v", err)
	}

	_, err = myRec.GetKeyECC("mypasswordiswrong")
	if err == nil {
		t.Fatalf("Incorrect password did not fail")
	}

	_, err = myRec.GetKeyECC("weakpassword")
	if err != nil {
		t.Fatalf("%v", err)
	}
}
<commit_msg>Adding TestChangePassword to tests<commit_after>// passvault_test: tests for passvault.go
//
// Copyright (c) 2013 CloudFlare, Inc.

package passvault

import (
	"os"
	"testing"
)

func TestStaticVault(t *testing.T) {
	// Creates a temporary on-disk database to test if passvault can read and
	// write from/to disk.  It's deleted at the bottom of the function--this
	// should be the only test that requires touching disk.
	records, err := InitFrom("/tmp/redoctober.json")
	if err != nil {
		t.Fatalf("Error reading record")
	}

	_, err = records.AddNewRecord("test", "bad pass", true, DefaultRecordType)
	if err != nil {
		t.Fatalf("Error creating record")
	}

	// Reads data written last time.
	_, err = InitFrom("/tmp/redoctober.json")
	if err != nil {
		t.Fatalf("Error reading record")
	}

	os.Remove("/tmp/redoctober.json")
}

func TestRSAEncryptDecrypt(t *testing.T) {
	records, err := InitFrom("memory")
	if err != nil {
		t.Fatalf("%v", err)
	}

	myRec, err := records.AddNewRecord("user", "weakpassword", true, RSARecord)
	if err != nil {
		t.Fatalf("%v", err)
	}

	_, err = myRec.GetKeyRSAPub()
	if err != nil {
		t.Fatalf("Error extracting RSA Pub")
	}

	rsaPriv, err := myRec.GetKeyRSA("mypasswordiswrong")
	if err == nil {
		t.Fatalf("Incorrect password did not fail")
	}

	rsaPriv, err = myRec.GetKeyRSA("weakpassword")
	if err != nil {
		t.Fatalf("Error decrypting RSA key")
	}

	err = rsaPriv.Validate()
	if err != nil {
		t.Fatalf("Error validating RSA key")
	}
}

func TestECCEncryptDecrypt(t *testing.T) {
	records, err := InitFrom("memory")
	if err != nil {
		t.Fatalf("%v", err)
	}

	myRec, err := records.AddNewRecord("user", "weakpassword", true, ECCRecord)
	if err != nil {
		t.Fatalf("%v", err)
	}

	_, err = myRec.GetKeyECCPub()
	if err != nil {
		t.Fatalf("%v", err)
	}

	_, err = myRec.GetKeyECC("mypasswordiswrong")
	if err == nil {
		t.Fatalf("Incorrect password did not fail")
	}

	_, err = myRec.GetKeyECC("weakpassword")
	if err != nil {
		t.Fatalf("%v", err)
	}
}

func TestChangePassword(t *testing.T) {
	records, err := InitFrom("memory")
	if err != nil {
		t.Fatalf("%v", err)
	}

	_, err = records.AddNewRecord("user", "weakpassword", true, ECCRecord)
	if err != nil {
		t.Fatalf("%v", err)
	}

	err = records.ChangePassword("user", "weakpassword", "newpassword")
	if err != nil {
		t.Fatalf("%v", err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All rights reserved.
//
// 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 zoekt

import (
	"context"
	"fmt"
	"log"
	"sort"
	"strings"

	"golang.org/x/net/trace"

	"github.com/google/zoekt/query"
)

func (d *indexData) simplify(in query.Q) query.Q {
	eval := query.Map(in, func(q query.Q) query.Q {
		if r, ok := q.(*query.Repo); ok {
			return &query.Const{Value: strings.Contains(d.repoMetaData.Name, r.Pattern)}
		}
		if l, ok := q.(*query.Language); ok {
			_, has := d.metaData.LanguageMap[l.Language]
			if !has {
				return &query.Const{Value: false}
			}
		}
		return q
	})
	return query.Simplify(eval)
}

func (o *SearchOptions) SetDefaults() {
	if o.ShardMaxMatchCount == 0 {
		// We cap the total number of matches, so overly broad
		// searches don't crash the machine.
		o.ShardMaxMatchCount = 100000
	}
	if o.TotalMaxMatchCount == 0 {
		o.TotalMaxMatchCount = 10 * o.ShardMaxMatchCount
	}
	if o.ShardMaxImportantMatch == 0 {
		o.ShardMaxImportantMatch = 10
	}
	if o.TotalMaxImportantMatch == 0 {
		o.TotalMaxImportantMatch = 10 * o.ShardMaxImportantMatch
	}
}

func (d *indexData) Search(ctx context.Context, q query.Q, opts *SearchOptions) (sr *SearchResult, err error) {
	copyOpts := *opts
	opts = &copyOpts
	opts.SetDefaults()
	importantMatchCount := 0

	var res SearchResult
	if len(d.fileNameIndex) == 0 {
		return &res, nil
	}

	tr := trace.New("indexData.Search", d.file.Name())
	tr.LazyPrintf("opts: %+v", opts)
	defer func() {
		if sr != nil {
			tr.LazyPrintf("num files: %d", len(sr.Files))
			tr.LazyPrintf("stats: %+v", sr.Stats)
		}
		if err != nil {
			tr.LazyPrintf("error: %v", err)
			tr.SetError()
		}
		tr.Finish()
	}()

	q = d.simplify(q)
	tr.LazyLog(q, true)
	if c, ok := q.(*query.Const); ok && !c.Value {
		return &res, nil
	}

	if opts.EstimateDocCount {
		res.Stats.ShardFilesConsidered = len(d.fileBranchMasks)
		return &res, nil
	}

	select {
	case <-ctx.Done():
		res.Stats.ShardsSkipped++
		return &res, nil
	default:
	}

	q = query.Map(q, query.ExpandFileContent)

	mt, err := d.newMatchTree(q, &res.Stats)
	if err != nil {
		return nil, err
	}

	totalAtomCount := 0
	visitMatchTree(mt, func(t matchTree) {
		totalAtomCount++
	})

	cp := &contentProvider{
		id:    d,
		stats: &res.Stats,
	}

	docCount := uint32(len(d.fileBranchMasks))
	lastDoc := int(-1)

nextFileMatch:
	for {
		canceled := false
		select {
		case <-ctx.Done():
			canceled = true
		default:
		}

		nextDoc := mt.nextDoc()
		if int(nextDoc) <= lastDoc {
			nextDoc = uint32(lastDoc + 1)
		}
		if nextDoc >= docCount {
			break
		}
		lastDoc = int(nextDoc)

		if canceled || res.Stats.MatchCount >= opts.ShardMaxMatchCount ||
			importantMatchCount >= opts.ShardMaxImportantMatch {
			res.Stats.FilesSkipped += d.repoListEntry.Stats.Documents - lastDoc
			break
		}

		res.Stats.FilesConsidered++
		mt.prepare(nextDoc)

		cp.setDocument(nextDoc)

		known := make(map[matchTree]bool)
		for cost := costMin; cost <= costMax; cost++ {
			v, ok := mt.matches(cp, cost, known)
			if ok && !v {
				continue nextFileMatch
			}

			if cost == costMax && !ok {
				log.Panicf("did not decide. Repo %s, doc %d, known %v",
					d.repoMetaData.Name, nextDoc, known)
			}
		}

		fileMatch := FileMatch{
			Repository: d.repoMetaData.Name,
			FileName:   string(d.fileName(nextDoc)),
			// Maintain ordering of input files. This
			// strictly dominates the in-file ordering of
			// the matches.
			Score:    10 * float64(nextDoc) / float64(len(d.boundaries)),
			Checksum: d.getChecksum(nextDoc),
			Language: d.languageMap[d.languages[nextDoc]],
		}

		if s := d.subRepos[nextDoc]; s > 0 {
			if s >= uint32(len(d.subRepoPaths)) {
				log.Panicf("corrupt index: subrepo %d beyond %v", s, d.subRepoPaths)
			}
			path := d.subRepoPaths[s]
			fileMatch.SubRepositoryPath = path
			sr := d.repoMetaData.SubRepoMap[path]
			fileMatch.SubRepositoryName = sr.Name
			if idx := d.branchIndex(nextDoc); idx >= 0 {
				fileMatch.Version = sr.Branches[idx].Version
			}
		} else {
			idx := d.branchIndex(nextDoc)
			if idx >= 0 {
				fileMatch.Version = d.repoMetaData.Branches[idx].Version
			}
		}

		atomMatchCount := 0
		visitMatches(mt, known, func(mt matchTree) {
			atomMatchCount++
		})
		fileMatch.Score += float64(atomMatchCount) / float64(totalAtomCount) * scoreFactorAtomMatch
		finalCands := gatherMatches(mt, known)

		if len(finalCands) == 0 {
			nm := d.fileName(nextDoc)
			finalCands = append(finalCands,
				&candidateMatch{
					caseSensitive: false,
					fileName:      true,
					substrBytes:   nm,
					substrLowered: nm,
					file:          nextDoc,
					runeOffset:    0,
					byteOffset:    0,
					byteMatchSz:   uint32(len(nm)),
				})
		}
		fileMatch.LineMatches = cp.fillMatches(finalCands)

		maxFileScore := 0.0
		for i := range fileMatch.LineMatches {
			if maxFileScore < fileMatch.LineMatches[i].Score {
				maxFileScore = fileMatch.LineMatches[i].Score

			}

			// Order by ordering in file.
			fileMatch.LineMatches[i].Score += 1.0 - (float64(i) / float64(len(fileMatch.LineMatches)))
		}
		fileMatch.Score += maxFileScore

		if fileMatch.Score > scoreImportantThreshold {
			importantMatchCount++
		}
		fileMatch.Branches = d.gatherBranches(nextDoc, mt, known)

		sortMatchesByScore(fileMatch.LineMatches)
		if opts.Whole {
			fileMatch.Content = cp.data(false)
		}

		res.Files = append(res.Files, fileMatch)
		res.Stats.MatchCount += len(fileMatch.LineMatches)
		res.Stats.FileCount++
	}
	SortFilesByScore(res.Files)

	addRepo(&res, &d.repoMetaData)
	for _, v := range d.repoMetaData.SubRepoMap {
		addRepo(&res, v)
	}

	mt.updateStats(&res.Stats)
	return &res, nil
}

func addRepo(res *SearchResult, repo *Repository) {
	if res.RepoURLs == nil {
		res.RepoURLs = map[string]string{}
	}
	res.RepoURLs[repo.Name] = repo.FileURLTemplate

	if res.LineFragments == nil {
		res.LineFragments = map[string]string{}
	}
	res.LineFragments[repo.Name] = repo.LineFragmentTemplate
}

type sortByOffsetSlice []*candidateMatch

func (m sortByOffsetSlice) Len() int      { return len(m) }
func (m sortByOffsetSlice) Swap(i, j int) { m[i], m[j] = m[j], m[i] }
func (m sortByOffsetSlice) Less(i, j int) bool {
	return m[i].byteOffset < m[j].byteOffset
}

// Gather matches from this document. This never returns a mixture of
// filename/content matches: if there are content matches, all
// filename matches are trimmed from the result. The matches are
// returned in document order and are non-overlapping.
func gatherMatches(mt matchTree, known map[matchTree]bool) []*candidateMatch {
	var cands []*candidateMatch
	visitMatches(mt, known, func(mt matchTree) {
		if smt, ok := mt.(*substrMatchTree); ok {
			cands = append(cands, smt.current...)
		}
		if rmt, ok := mt.(*regexpMatchTree); ok {
			cands = append(cands, rmt.found...)
		}
	})

	foundContentMatch := false
	for _, c := range cands {
		if !c.fileName {
			foundContentMatch = true
			break
		}
	}

	res := cands[:0]
	for _, c := range cands {
		if !foundContentMatch || !c.fileName {
			res = append(res, c)
		}
	}
	cands = res

	// Merge adjacent candidates. This guarantees that the matches
	// are non-overlapping.
	sort.Sort((sortByOffsetSlice)(cands))
	res = cands[:0]
	for i, c := range cands {
		if i == 0 {
			res = append(res, c)
			continue
		}
		last := res[len(res)-1]
		lastEnd := last.byteOffset + last.byteMatchSz
		end := c.byteOffset + c.byteMatchSz
		if lastEnd >= c.byteOffset {
			if end > lastEnd {
				last.byteMatchSz = end - last.byteOffset
			}
			continue
		}

		res = append(res, c)
	}

	return res
}

func (d *indexData) branchIndex(docID uint32) int {
	mask := d.fileBranchMasks[docID]
	idx := 0
	for mask != 0 {
		if mask&0x1 != 0 {
			return idx
		}
		idx++
		mask >>= 1
	}
	return -1
}

// gatherBranches returns a list of branch names.
func (d *indexData) gatherBranches(docID uint32, mt matchTree, known map[matchTree]bool) []string {
	foundBranchQuery := false
	var branches []string

	visitMatches(mt, known, func(mt matchTree) {
		bq, ok := mt.(*branchQueryMatchTree)
		if ok {
			foundBranchQuery = true
			branches = append(branches,
				d.branchNames[uint(bq.mask)])
		}
	})

	if !foundBranchQuery {
		mask := d.fileBranchMasks[docID]
		id := uint32(1)
		for mask != 0 {
			if mask&0x1 != 0 {
				branches = append(branches, d.branchNames[uint(id)])
			}
			id <<= 1
			mask >>= 1
		}
	}
	return branches
}

func (d *indexData) List(ctx context.Context, q query.Q) (rl *RepoList, err error) {
	tr := trace.New("indexData.List", d.file.Name())
	defer func() {
		if rl != nil {
			tr.LazyPrintf("repos size: %d", len(rl.Repos))
			tr.LazyPrintf("crashes: %d", rl.Crashes)
		}
		if err != nil {
			tr.LazyPrintf("error: %v", err)
			tr.SetError()
		}
		tr.Finish()
	}()

	q = d.simplify(q)
	tr.LazyLog(q, true)
	c, ok := q.(*query.Const)

	if !ok {
		return nil, fmt.Errorf("List should receive Repo-only query.")
	}

	l := &RepoList{}
	if c.Value {
		l.Repos = append(l.Repos, &d.repoListEntry)
	}
	return l, nil
}
<commit_msg>Bail earlier if the context was canceled<commit_after>// Copyright 2016 Google Inc. All rights reserved.
//
// 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 zoekt

import (
	"context"
	"fmt"
	"log"
	"sort"
	"strings"

	"golang.org/x/net/trace"

	"github.com/google/zoekt/query"
)

func (d *indexData) simplify(in query.Q) query.Q {
	eval := query.Map(in, func(q query.Q) query.Q {
		if r, ok := q.(*query.Repo); ok {
			return &query.Const{Value: strings.Contains(d.repoMetaData.Name, r.Pattern)}
		}
		if l, ok := q.(*query.Language); ok {
			_, has := d.metaData.LanguageMap[l.Language]
			if !has {
				return &query.Const{Value: false}
			}
		}
		return q
	})
	return query.Simplify(eval)
}

func (o *SearchOptions) SetDefaults() {
	if o.ShardMaxMatchCount == 0 {
		// We cap the total number of matches, so overly broad
		// searches don't crash the machine.
		o.ShardMaxMatchCount = 100000
	}
	if o.TotalMaxMatchCount == 0 {
		o.TotalMaxMatchCount = 10 * o.ShardMaxMatchCount
	}
	if o.ShardMaxImportantMatch == 0 {
		o.ShardMaxImportantMatch = 10
	}
	if o.TotalMaxImportantMatch == 0 {
		o.TotalMaxImportantMatch = 10 * o.ShardMaxImportantMatch
	}
}

func (d *indexData) Search(ctx context.Context, q query.Q, opts *SearchOptions) (sr *SearchResult, err error) {
	copyOpts := *opts
	opts = &copyOpts
	opts.SetDefaults()
	importantMatchCount := 0

	var res SearchResult
	if len(d.fileNameIndex) == 0 {
		return &res, nil
	}

	select {
	case <-ctx.Done():
		res.Stats.ShardsSkipped++
		return &res, nil
	default:
	}

	tr := trace.New("indexData.Search", d.file.Name())
	tr.LazyPrintf("opts: %+v", opts)
	defer func() {
		if sr != nil {
			tr.LazyPrintf("num files: %d", len(sr.Files))
			tr.LazyPrintf("stats: %+v", sr.Stats)
		}
		if err != nil {
			tr.LazyPrintf("error: %v", err)
			tr.SetError()
		}
		tr.Finish()
	}()

	q = d.simplify(q)
	tr.LazyLog(q, true)
	if c, ok := q.(*query.Const); ok && !c.Value {
		return &res, nil
	}

	if opts.EstimateDocCount {
		res.Stats.ShardFilesConsidered = len(d.fileBranchMasks)
		return &res, nil
	}

	q = query.Map(q, query.ExpandFileContent)

	mt, err := d.newMatchTree(q, &res.Stats)
	if err != nil {
		return nil, err
	}

	totalAtomCount := 0
	visitMatchTree(mt, func(t matchTree) {
		totalAtomCount++
	})

	cp := &contentProvider{
		id:    d,
		stats: &res.Stats,
	}

	docCount := uint32(len(d.fileBranchMasks))
	lastDoc := int(-1)

nextFileMatch:
	for {
		canceled := false
		select {
		case <-ctx.Done():
			canceled = true
		default:
		}

		nextDoc := mt.nextDoc()
		if int(nextDoc) <= lastDoc {
			nextDoc = uint32(lastDoc + 1)
		}
		if nextDoc >= docCount {
			break
		}
		lastDoc = int(nextDoc)

		if canceled || res.Stats.MatchCount >= opts.ShardMaxMatchCount ||
			importantMatchCount >= opts.ShardMaxImportantMatch {
			res.Stats.FilesSkipped += d.repoListEntry.Stats.Documents - lastDoc
			break
		}

		res.Stats.FilesConsidered++
		mt.prepare(nextDoc)

		cp.setDocument(nextDoc)

		known := make(map[matchTree]bool)
		for cost := costMin; cost <= costMax; cost++ {
			v, ok := mt.matches(cp, cost, known)
			if ok && !v {
				continue nextFileMatch
			}

			if cost == costMax && !ok {
				log.Panicf("did not decide. Repo %s, doc %d, known %v",
					d.repoMetaData.Name, nextDoc, known)
			}
		}

		fileMatch := FileMatch{
			Repository: d.repoMetaData.Name,
			FileName:   string(d.fileName(nextDoc)),
			// Maintain ordering of input files. This
			// strictly dominates the in-file ordering of
			// the matches.
			Score:    10 * float64(nextDoc) / float64(len(d.boundaries)),
			Checksum: d.getChecksum(nextDoc),
			Language: d.languageMap[d.languages[nextDoc]],
		}

		if s := d.subRepos[nextDoc]; s > 0 {
			if s >= uint32(len(d.subRepoPaths)) {
				log.Panicf("corrupt index: subrepo %d beyond %v", s, d.subRepoPaths)
			}
			path := d.subRepoPaths[s]
			fileMatch.SubRepositoryPath = path
			sr := d.repoMetaData.SubRepoMap[path]
			fileMatch.SubRepositoryName = sr.Name
			if idx := d.branchIndex(nextDoc); idx >= 0 {
				fileMatch.Version = sr.Branches[idx].Version
			}
		} else {
			idx := d.branchIndex(nextDoc)
			if idx >= 0 {
				fileMatch.Version = d.repoMetaData.Branches[idx].Version
			}
		}

		atomMatchCount := 0
		visitMatches(mt, known, func(mt matchTree) {
			atomMatchCount++
		})
		fileMatch.Score += float64(atomMatchCount) / float64(totalAtomCount) * scoreFactorAtomMatch
		finalCands := gatherMatches(mt, known)

		if len(finalCands) == 0 {
			nm := d.fileName(nextDoc)
			finalCands = append(finalCands,
				&candidateMatch{
					caseSensitive: false,
					fileName:      true,
					substrBytes:   nm,
					substrLowered: nm,
					file:          nextDoc,
					runeOffset:    0,
					byteOffset:    0,
					byteMatchSz:   uint32(len(nm)),
				})
		}
		fileMatch.LineMatches = cp.fillMatches(finalCands)

		maxFileScore := 0.0
		for i := range fileMatch.LineMatches {
			if maxFileScore < fileMatch.LineMatches[i].Score {
				maxFileScore = fileMatch.LineMatches[i].Score

			}

			// Order by ordering in file.
			fileMatch.LineMatches[i].Score += 1.0 - (float64(i) / float64(len(fileMatch.LineMatches)))
		}
		fileMatch.Score += maxFileScore

		if fileMatch.Score > scoreImportantThreshold {
			importantMatchCount++
		}
		fileMatch.Branches = d.gatherBranches(nextDoc, mt, known)

		sortMatchesByScore(fileMatch.LineMatches)
		if opts.Whole {
			fileMatch.Content = cp.data(false)
		}

		res.Files = append(res.Files, fileMatch)
		res.Stats.MatchCount += len(fileMatch.LineMatches)
		res.Stats.FileCount++
	}
	SortFilesByScore(res.Files)

	addRepo(&res, &d.repoMetaData)
	for _, v := range d.repoMetaData.SubRepoMap {
		addRepo(&res, v)
	}

	mt.updateStats(&res.Stats)
	return &res, nil
}

func addRepo(res *SearchResult, repo *Repository) {
	if res.RepoURLs == nil {
		res.RepoURLs = map[string]string{}
	}
	res.RepoURLs[repo.Name] = repo.FileURLTemplate

	if res.LineFragments == nil {
		res.LineFragments = map[string]string{}
	}
	res.LineFragments[repo.Name] = repo.LineFragmentTemplate
}

type sortByOffsetSlice []*candidateMatch

func (m sortByOffsetSlice) Len() int      { return len(m) }
func (m sortByOffsetSlice) Swap(i, j int) { m[i], m[j] = m[j], m[i] }
func (m sortByOffsetSlice) Less(i, j int) bool {
	return m[i].byteOffset < m[j].byteOffset
}

// Gather matches from this document. This never returns a mixture of
// filename/content matches: if there are content matches, all
// filename matches are trimmed from the result. The matches are
// returned in document order and are non-overlapping.
func gatherMatches(mt matchTree, known map[matchTree]bool) []*candidateMatch {
	var cands []*candidateMatch
	visitMatches(mt, known, func(mt matchTree) {
		if smt, ok := mt.(*substrMatchTree); ok {
			cands = append(cands, smt.current...)
		}
		if rmt, ok := mt.(*regexpMatchTree); ok {
			cands = append(cands, rmt.found...)
		}
	})

	foundContentMatch := false
	for _, c := range cands {
		if !c.fileName {
			foundContentMatch = true
			break
		}
	}

	res := cands[:0]
	for _, c := range cands {
		if !foundContentMatch || !c.fileName {
			res = append(res, c)
		}
	}
	cands = res

	// Merge adjacent candidates. This guarantees that the matches
	// are non-overlapping.
	sort.Sort((sortByOffsetSlice)(cands))
	res = cands[:0]
	for i, c := range cands {
		if i == 0 {
			res = append(res, c)
			continue
		}
		last := res[len(res)-1]
		lastEnd := last.byteOffset + last.byteMatchSz
		end := c.byteOffset + c.byteMatchSz
		if lastEnd >= c.byteOffset {
			if end > lastEnd {
				last.byteMatchSz = end - last.byteOffset
			}
			continue
		}

		res = append(res, c)
	}

	return res
}

func (d *indexData) branchIndex(docID uint32) int {
	mask := d.fileBranchMasks[docID]
	idx := 0
	for mask != 0 {
		if mask&0x1 != 0 {
			return idx
		}
		idx++
		mask >>= 1
	}
	return -1
}

// gatherBranches returns a list of branch names.
func (d *indexData) gatherBranches(docID uint32, mt matchTree, known map[matchTree]bool) []string {
	foundBranchQuery := false
	var branches []string

	visitMatches(mt, known, func(mt matchTree) {
		bq, ok := mt.(*branchQueryMatchTree)
		if ok {
			foundBranchQuery = true
			branches = append(branches,
				d.branchNames[uint(bq.mask)])
		}
	})

	if !foundBranchQuery {
		mask := d.fileBranchMasks[docID]
		id := uint32(1)
		for mask != 0 {
			if mask&0x1 != 0 {
				branches = append(branches, d.branchNames[uint(id)])
			}
			id <<= 1
			mask >>= 1
		}
	}
	return branches
}

func (d *indexData) List(ctx context.Context, q query.Q) (rl *RepoList, err error) {
	tr := trace.New("indexData.List", d.file.Name())
	defer func() {
		if rl != nil {
			tr.LazyPrintf("repos size: %d", len(rl.Repos))
			tr.LazyPrintf("crashes: %d", rl.Crashes)
		}
		if err != nil {
			tr.LazyPrintf("error: %v", err)
			tr.SetError()
		}
		tr.Finish()
	}()

	q = d.simplify(q)
	tr.LazyLog(q, true)
	c, ok := q.(*query.Const)

	if !ok {
		return nil, fmt.Errorf("List should receive Repo-only query.")
	}

	l := &RepoList{}
	if c.Value {
		l.Repos = append(l.Repos, &d.repoListEntry)
	}
	return l, nil
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"

	"github.com/99designs/aws-vault/keyring"
	"github.com/aws/aws-sdk-go/aws/awserr"
)

type ExecCommandInput struct {
	Profile  string
	Command  string
	Args     []string
	Keyring  keyring.Keyring
	Duration time.Duration
	WriteEnv bool
	MfaToken string
	Signals  chan os.Signal
}

func ExecCommand(ui Ui, input ExecCommandInput) {
	creds, err := NewVaultCredentials(input.Keyring, input.Profile, VaultOptions{
		SessionDuration: input.Duration,
		WriteEnv:        input.WriteEnv,
		MfaToken:        input.MfaToken,
	})
	if err != nil {
		ui.Error.Fatal(err)
	}

	val, err := creds.Get()
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "NoCredentialProviders" {
			ui.Error.Fatalf("No credentials found for profile %q", input.Profile)
		} else {
			ui.Error.Fatal(err)
		}
	}

	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		ui.Error.Fatal(err)
	}

	go func() {
		log.Printf("Metadata server listening on %s", l.Addr().String())
		ui.Error.Fatal(http.Serve(l, proxyHandler(NewMetadataHandler(creds))))
	}()

	profs, err := parseProfiles()
	if err != nil {
		ui.Error.Fatal(err)
	}

	cfg, err := writeTempConfig(input.Profile, profs)
	if err != nil {
		ui.Error.Fatal(err)
	}

	env := os.Environ()
	env = overwriteEnv(env, "http_proxy", "http://"+l.Addr().String())
	env = overwriteEnv(env, "no_proxy", "amazonaws.com")
	env = overwriteEnv(env, "AWS_CONFIG_FILE", cfg.Name())
	env = overwriteEnv(env, "AWS_PROFILE", input.Profile)
	env = unsetEnv(env, "AWS_CREDENTIAL_FILE")

	if region, ok := profs[input.Profile]["region"]; ok {
		env = overwriteEnv(env, "AWS_REGION", region)
	}

	if input.WriteEnv {
		env = overwriteEnv(env, "AWS_ACCESS_KEY_ID", val.AccessKeyID)
		env = overwriteEnv(env, "AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)

		if val.SessionToken != "" {
			env = overwriteEnv(env, "AWS_SESSION_TOKEN", val.SessionToken)
			env = overwriteEnv(env, "AWS_SECURITY_TOKEN", val.SessionToken)
		}
	}

	cmd := exec.Command(input.Command, input.Args...)
	cmd.Env = env
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	go func() {
		sig := <-input.Signals
		if cmd.Process != nil {
			cmd.Process.Signal(sig)
		}
	}()

	var waitStatus syscall.WaitStatus
	if err := cmd.Run(); err != nil {
		if err != nil {
			ui.Error.Println(err)
		}
		if exitError, ok := err.(*exec.ExitError); ok {
			waitStatus = exitError.Sys().(syscall.WaitStatus)
			os.Exit(waitStatus.ExitStatus())
		}
	}

}

// write out a config excluding role switching keys
func writeTempConfig(profile string, conf profiles) (*os.File, error) {
	tmpConfig, err := ioutil.TempFile(os.TempDir(), "aws-vault")
	if err != nil {
		return nil, err
	}

	// allow some time for keychain prompt
	newConfig := map[string]string{
		"metadata_service_timeout":      "15",
		"metadata_service_num_attempts": "2",
	}

	for k, v := range conf[profile] {
		if k != "source_profile" && k != "role_arn" {
			newConfig[k] = v
		}
	}

	return tmpConfig, writeProfiles(tmpConfig, profiles{profile: newConfig})
}

func unsetEnv(env []string, key string) []string {
	envCopy := []string{}

	for _, e := range env {
		if !strings.HasPrefix(key+"=", e) {
			envCopy = append(envCopy, e)
		}
	}

	return envCopy
}

func overwriteEnv(env []string, key, val string) []string {
	var found bool

	for idx, e := range env {
		if strings.HasPrefix(key+"=", e) {
			env[idx] = key + "=" + val
			found = true
		} else {
			env[idx] = e
		}
	}

	if !found {
		env = append(env, key+"="+val)
	}

	return env
}

func connectProxy(w http.ResponseWriter, r *http.Request) error {
	d, err := net.Dial("tcp", r.RequestURI)
	if err != nil {
		return err
	}
	hj, ok := w.(http.Hijacker)
	if !ok {
		return err
	}
	conn, bufrw, err := hj.Hijack()
	if err != nil {
		return err
	}
	defer conn.Close()
	defer d.Close()

	bufrw.WriteString("HTTP/1.0 200 Connection established\r\n\r\n")
	bufrw.Flush()

	go io.Copy(d, conn)
	io.Copy(conn, d)

	return nil
}

func proxyHandler(upstream http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%s %s", r.Method, r.RequestURI)
		if r.Method == "GET" && strings.HasPrefix(r.URL.Path, "/latest/meta-data") {
			upstream.ServeHTTP(w, r)
		} else if r.Method == "CONNECT" {
			if err := connectProxy(w, r); err != nil {
				http.Error(w, "Error connecting to proxy: "+err.Error(), http.StatusGatewayTimeout)
			}
		} else {
			proxy := &httputil.ReverseProxy{Director: func(req *http.Request) {
				req.Method = "GET"
				req.URL.Scheme = "http"
				req.URL.Host = r.Host
			}}
			proxy.ServeHTTP(w, r)
		}
	})
}
<commit_msg>Unset AWS envs that interfer with operation<commit_after>package main

import (
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"

	"github.com/99designs/aws-vault/keyring"
	"github.com/aws/aws-sdk-go/aws/awserr"
)

type ExecCommandInput struct {
	Profile  string
	Command  string
	Args     []string
	Keyring  keyring.Keyring
	Duration time.Duration
	WriteEnv bool
	MfaToken string
	Signals  chan os.Signal
}

func ExecCommand(ui Ui, input ExecCommandInput) {
	creds, err := NewVaultCredentials(input.Keyring, input.Profile, VaultOptions{
		SessionDuration: input.Duration,
		WriteEnv:        input.WriteEnv,
		MfaToken:        input.MfaToken,
	})
	if err != nil {
		ui.Error.Fatal(err)
	}

	val, err := creds.Get()
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "NoCredentialProviders" {
			ui.Error.Fatalf("No credentials found for profile %q", input.Profile)
		} else {
			ui.Error.Fatal(err)
		}
	}

	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		ui.Error.Fatal(err)
	}

	go func() {
		log.Printf("Metadata server listening on %s", l.Addr().String())
		ui.Error.Fatal(http.Serve(l, proxyHandler(NewMetadataHandler(creds))))
	}()

	profs, err := parseProfiles()
	if err != nil {
		ui.Error.Fatal(err)
	}

	cfg, err := writeTempConfig(input.Profile, profs)
	if err != nil {
		ui.Error.Fatal(err)
	}

	env := os.Environ()
	env = overwriteEnv(env, "http_proxy", "http://"+l.Addr().String())
	env = overwriteEnv(env, "no_proxy", "amazonaws.com")
	env = overwriteEnv(env, "AWS_CONFIG_FILE", cfg.Name())
	env = overwriteEnv(env, "AWS_PROFILE", input.Profile)

	env = unsetEnv(env, "AWS_ACCESS_KEY_ID")
	env = unsetEnv(env, "AWS_SECRET_ACCESS_KEY")
	env = unsetEnv(env, "AWS_CREDENTIAL_FILE")

	if region, ok := profs[input.Profile]["region"]; ok {
		env = overwriteEnv(env, "AWS_REGION", region)
	}

	if input.WriteEnv {
		env = overwriteEnv(env, "AWS_ACCESS_KEY_ID", val.AccessKeyID)
		env = overwriteEnv(env, "AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)

		if val.SessionToken != "" {
			env = overwriteEnv(env, "AWS_SESSION_TOKEN", val.SessionToken)
			env = overwriteEnv(env, "AWS_SECURITY_TOKEN", val.SessionToken)
		}
	}

	cmd := exec.Command(input.Command, input.Args...)
	cmd.Env = env
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	go func() {
		sig := <-input.Signals
		if cmd.Process != nil {
			cmd.Process.Signal(sig)
		}
	}()

	var waitStatus syscall.WaitStatus
	if err := cmd.Run(); err != nil {
		if err != nil {
			ui.Error.Println(err)
		}
		if exitError, ok := err.(*exec.ExitError); ok {
			waitStatus = exitError.Sys().(syscall.WaitStatus)
			os.Exit(waitStatus.ExitStatus())
		}
	}

}

// write out a config excluding role switching keys
func writeTempConfig(profile string, conf profiles) (*os.File, error) {
	tmpConfig, err := ioutil.TempFile(os.TempDir(), "aws-vault")
	if err != nil {
		return nil, err
	}

	// allow some time for keychain prompt
	newConfig := map[string]string{
		"metadata_service_timeout":      "15",
		"metadata_service_num_attempts": "2",
	}

	for k, v := range conf[profile] {
		if k != "source_profile" && k != "role_arn" {
			newConfig[k] = v
		}
	}

	return tmpConfig, writeProfiles(tmpConfig, profiles{profile: newConfig})
}

func unsetEnv(env []string, key string) []string {
	envCopy := []string{}

	for _, e := range env {
		if !strings.HasPrefix(key+"=", e) {
			envCopy = append(envCopy, e)
		}
	}

	return envCopy
}

func overwriteEnv(env []string, key, val string) []string {
	var found bool

	for idx, e := range env {
		if strings.HasPrefix(key+"=", e) {
			env[idx] = key + "=" + val
			found = true
		} else {
			env[idx] = e
		}
	}

	if !found {
		env = append(env, key+"="+val)
	}

	return env
}

func connectProxy(w http.ResponseWriter, r *http.Request) error {
	d, err := net.Dial("tcp", r.RequestURI)
	if err != nil {
		return err
	}
	hj, ok := w.(http.Hijacker)
	if !ok {
		return err
	}
	conn, bufrw, err := hj.Hijack()
	if err != nil {
		return err
	}
	defer conn.Close()
	defer d.Close()

	bufrw.WriteString("HTTP/1.0 200 Connection established\r\n\r\n")
	bufrw.Flush()

	go io.Copy(d, conn)
	io.Copy(conn, d)

	return nil
}

func proxyHandler(upstream http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%s %s", r.Method, r.RequestURI)
		if r.Method == "GET" && strings.HasPrefix(r.URL.Path, "/latest/meta-data") {
			upstream.ServeHTTP(w, r)
		} else if r.Method == "CONNECT" {
			if err := connectProxy(w, r); err != nil {
				http.Error(w, "Error connecting to proxy: "+err.Error(), http.StatusGatewayTimeout)
			}
		} else {
			proxy := &httputil.ReverseProxy{Director: func(req *http.Request) {
				req.Method = "GET"
				req.URL.Scheme = "http"
				req.URL.Host = r.Host
			}}
			proxy.ServeHTTP(w, r)
		}
	})
}
<|endoftext|>


<commit_before>package lib

import (
	"fmt"
	"plaid/lang"
	"plaid/lang/types"
)

func IO() *lang.Library {
	lib := lang.MakeLibrary("io")

	lib.Function("print", types.Function{
		Params: types.Tuple{[]types.Type{
			types.Any{},
		}},
		Ret: types.Void{},
	}, func(args []lang.Object) (lang.Object, error) {
		if len(args) != 1 {
			err := fmt.Errorf("wanted 1 argument, got %d", len(args))
			return lang.ObjectNone{}, err
		}

		fmt.Println(args[0].String())
		return lang.ObjectNone{}, nil
	})

	return lib
}
<commit_msg>use formatted value instead of debugging representation<commit_after>package lib

import (
	"fmt"
	"plaid/lang"
	"plaid/lang/types"
)

func IO() *lang.Library {
	lib := lang.MakeLibrary("io")

	lib.Function("print", types.Function{
		Params: types.Tuple{[]types.Type{
			types.Any{},
		}},
		Ret: types.Void{},
	}, func(args []lang.Object) (lang.Object, error) {
		if len(args) != 1 {
			err := fmt.Errorf("wanted 1 argument, got %d", len(args))
			return lang.ObjectNone{}, err
		}

		fmt.Println(args[0].Value())
		return lang.ObjectNone{}, nil
	})

	return lib
}
<|endoftext|>


<commit_before>package utp

/*
#include "utp.h"
*/
import "C"
import (
	"errors"
	"sync"
)

type Option = C.int

const (
	LogNormal   Option = C.UTP_LOG_NORMAL
	LogMtu      Option = C.UTP_LOG_MTU
	LogDebug    Option = C.UTP_LOG_DEBUG
	SendBuffer  Option = C.UTP_SNDBUF
	RecvBuffer  Option = C.UTP_RCVBUF
	TargetDelay Option = C.UTP_TARGET_DELAY

	TimedOut = C.UTP_ETIMEDOUT
)

var (
	mu                 sync.Mutex
	libContextToSocket = map[*C.utp_context]*Socket{}
)

func getSocketForLibContext(uc *C.utp_context) *Socket {
	return libContextToSocket[uc]
}

func errorForCode(code C.int) error {
	return errors.New(libErrorCodeNames(code))
}
<commit_msg>Use anacrolix/sync.Mutex<commit_after>package utp

/*
#include "utp.h"
*/
import "C"
import (
	"errors"

	"github.com/anacrolix/sync"
)

type Option = C.int

const (
	LogNormal   Option = C.UTP_LOG_NORMAL
	LogMtu      Option = C.UTP_LOG_MTU
	LogDebug    Option = C.UTP_LOG_DEBUG
	SendBuffer  Option = C.UTP_SNDBUF
	RecvBuffer  Option = C.UTP_RCVBUF
	TargetDelay Option = C.UTP_TARGET_DELAY

	TimedOut = C.UTP_ETIMEDOUT
)

var (
	mu                 sync.Mutex
	libContextToSocket = map[*C.utp_context]*Socket{}
)

func getSocketForLibContext(uc *C.utp_context) *Socket {
	return libContextToSocket[uc]
}

func errorForCode(code C.int) error {
	return errors.New(libErrorCodeNames(code))
}
<|endoftext|>


<commit_before>package web

import (
    "bytes"
    "bufio"
    "encoding/binary"
    "log"
    "net"
    "os"
)

const (
    FcgiBeginRequest = iota + 1
    FcgiAbortRequest
    FcgiEndRequest
    FcgiParams
    FcgiStdin
    FcgiStdout
    FcgiStderr
    FcgiData
    FcgiGetValues
    FcgiGetValuesResult
    FcgiUnknownType
    FcgiMaxType = FcgiUnknownType
)

const (
    FcgiRequestComplete = iota
    FcgiCantMpxConn
    FcgiOverloaded
    FcgiUnknownRole
)

type fcgiHeader struct {
    Version       uint8
    Type          uint8
    RequestId     uint16
    ContentLength uint16
    PaddingLength uint8
    Reserved      uint8
}

func (h fcgiHeader) bytes() []byte {
    order := binary.BigEndian
    buf := make([]byte, 8)
    buf[0] = h.Version
    buf[1] = h.Type
    order.PutUint16(buf[2:4], h.RequestId)
    order.PutUint16(buf[4:6], h.ContentLength)
    buf[6] = h.PaddingLength
    buf[7] = h.Reserved
    return buf
}

type fcgiEndRequest struct {
    appStatus      uint32
    protocolStatus uint8
    reserved       [3]uint8
}

func (er fcgiEndRequest) bytes() []byte {
    buf := make([]byte, 8)
    binary.BigEndian.PutUint32(buf, er.appStatus)
    buf[4] = er.protocolStatus
    return buf
}

type fcgiConn struct {
    requestId uint16
    fd        net.Conn
}

func (conn *fcgiConn) WriteString(s string) os.Error {
    content := bytes.NewBufferString(s).Bytes()
    l := len(content)
    // round to the nearest 8
    padding := make([]byte, uint8(-l&7))
    hdr := fcgiHeader{
        Version: 1,
        Type: FcgiStdout,
        RequestId: conn.requestId,
        ContentLength: uint16(l),
        PaddingLength: uint8(len(padding)),
    }

    conn.fd.Write(hdr.bytes())
    conn.fd.Write(content)
    conn.fd.Write(padding)

    return nil
}

func (conn *fcgiConn) StartResponse(status int) os.Error {
    return nil
}

func (conn *fcgiConn) Close() os.Error {
    content := fcgiEndRequest{appStatus: 200, protocolStatus: FcgiRequestComplete}.bytes()
    l := len(content)

    hdr := fcgiHeader{
        Version: 1,
        Type: FcgiEndRequest,
        RequestId: uint16(conn.requestId),
        ContentLength: uint16(l),
        PaddingLength: 0,
    }

    conn.fd.Write(hdr.bytes())
    conn.fd.Write(content)

    conn.fd.Close()

    return nil
}

func readFcgiParams(data []byte) {
    var params = make(map[string]string)

    for idx := 0; len(data) > idx; {
        var keySize int = int(data[idx])
        if keySize>>7 == 0 {
            idx += 1
        } else {
            binary.Read(bytes.NewBuffer(data[idx:idx+4]), binary.BigEndian, &keySize)
            idx += 4
        }

        var valSize int = int(data[idx])
        if valSize>>7 == 0 {
            idx += 1
        } else {
            binary.Read(bytes.NewBuffer(data[idx:idx+4]), binary.BigEndian, &valSize)
            idx += 4
        }

        key := data[idx : idx+keySize]
        idx += keySize
        val := data[idx : idx+valSize]
        idx += valSize

        println(string(key), ":", string(val))
        params[string(key)] = string(val)
    }
}

func handleFcgiRequest(fd net.Conn) {

    br := bufio.NewReader(fd)
    for {
        var h fcgiHeader
        err := binary.Read(br, binary.BigEndian, &h)
        if err != nil {
            log.Stderrf(err.String())
        }
        content := make([]byte, h.ContentLength)
        n, err := br.Read(content)
        //read padding
        if h.PaddingLength > 0 {
            padding := make([]byte, h.PaddingLength)
            n, err = br.Read(padding)
            println("read padding", n)
        }

        switch h.Type {
        case FcgiBeginRequest:
            println("begin request!")
        case FcgiParams:
            readFcgiParams(content)
        case FcgiStdin:
            println("stdin")
        case FcgiData:
            println("data!")
        case FcgiAbortRequest:
            println("abort!")
        }

    }
}

func listenAndServeFcgi(addr string) {
    println("listening and serving scgi!")
    l, err := net.Listen("tcp", addr)
    if err != nil {
        log.Stderrf(err.String())
        return
    }

    for {
        fd, err := l.Accept()
        if err != nil {
            log.Stderrf(err.String())
            break
        }
        go handleFcgiRequest(fd)

    }
}
<commit_msg>test fcgi return values<commit_after>package web

import (
    "bytes"
    "bufio"
    "encoding/binary"
    "log"
    "net"
    "os"
)

const (
    FcgiBeginRequest = iota + 1
    FcgiAbortRequest
    FcgiEndRequest
    FcgiParams
    FcgiStdin
    FcgiStdout
    FcgiStderr
    FcgiData
    FcgiGetValues
    FcgiGetValuesResult
    FcgiUnknownType
    FcgiMaxType = FcgiUnknownType
)

const (
    FcgiRequestComplete = iota
    FcgiCantMpxConn
    FcgiOverloaded
    FcgiUnknownRole
)

type fcgiHeader struct {
    Version       uint8
    Type          uint8
    RequestId     uint16
    ContentLength uint16
    PaddingLength uint8
    Reserved      uint8
}

func (h fcgiHeader) bytes() []byte {
    order := binary.BigEndian
    buf := make([]byte, 8)
    buf[0] = h.Version
    buf[1] = h.Type
    order.PutUint16(buf[2:4], h.RequestId)
    order.PutUint16(buf[4:6], h.ContentLength)
    buf[6] = h.PaddingLength
    buf[7] = h.Reserved
    return buf
}

type fcgiEndRequest struct {
    appStatus      uint32
    protocolStatus uint8
    reserved       [3]uint8
}

func (er fcgiEndRequest) bytes() []byte {
    buf := make([]byte, 8)
    binary.BigEndian.PutUint32(buf, er.appStatus)
    buf[4] = er.protocolStatus
    return buf
}

type fcgiConn struct {
    requestId uint16
    fd        net.Conn
}

func (conn *fcgiConn) WriteString(s string) os.Error {
    content := bytes.NewBufferString(s).Bytes()
    l := len(content)
    // round to the nearest 8
    padding := make([]byte, uint8(-l&7))
    hdr := fcgiHeader{
        Version: 1,
        Type: FcgiStdout,
        RequestId: conn.requestId,
        ContentLength: uint16(l),
        PaddingLength: uint8(len(padding)),
    }

    conn.fd.Write(hdr.bytes())
    conn.fd.Write(content)
    conn.fd.Write(padding)

    return nil
}

func (conn *fcgiConn) StartResponse(status int) os.Error {
    return nil
}

func (conn *fcgiConn) Close() os.Error {
    content := fcgiEndRequest{appStatus: 200, protocolStatus: FcgiRequestComplete}.bytes()
    l := len(content)

    hdr := fcgiHeader{
        Version: 1,
        Type: FcgiEndRequest,
        RequestId: uint16(conn.requestId),
        ContentLength: uint16(l),
        PaddingLength: 0,
    }

    conn.fd.Write(hdr.bytes())
    conn.fd.Write(content)

    conn.fd.Close()

    return nil
}

func readFcgiParams(data []byte) {
    var params = make(map[string]string)

    for idx := 0; len(data) > idx; {
        var keySize int = int(data[idx])
        if keySize>>7 == 0 {
            idx += 1
        } else {
            binary.Read(bytes.NewBuffer(data[idx:idx+4]), binary.BigEndian, &keySize)
            idx += 4
        }

        var valSize int = int(data[idx])
        if valSize>>7 == 0 {
            idx += 1
        } else {
            binary.Read(bytes.NewBuffer(data[idx:idx+4]), binary.BigEndian, &valSize)
            idx += 4
        }

        key := data[idx : idx+keySize]
        idx += keySize
        val := data[idx : idx+valSize]
        idx += valSize

        println(string(key), ":", string(val))
        params[string(key)] = string(val)
    }
}

func handleFcgiRequest(fd net.Conn) {

    br := bufio.NewReader(fd)

    var fc *fcgiConn
    for {
        var h fcgiHeader
        err := binary.Read(br, binary.BigEndian, &h)
        if err != nil {
            log.Stderrf(err.String())
            break
        }
        content := make([]byte, h.ContentLength)
        n, err := br.Read(content)
        //read padding
        if h.PaddingLength > 0 {
            padding := make([]byte, h.PaddingLength)
            n, err = br.Read(padding)
            println("read padding", n)
        }

        switch h.Type {
        case FcgiBeginRequest:
            fc = &fcgiConn{h.RequestId, fd}
        case FcgiParams:
            readFcgiParams(content)
        case FcgiStdin:
            fc.WriteString("HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\nHello")
            fc.Close()
        case FcgiData:
            println("data!")
        case FcgiAbortRequest:
            println("abort!")
        }
    }
}

func listenAndServeFcgi(addr string) {
    println("listening and serving scgi!")
    l, err := net.Listen("tcp", addr)
    if err != nil {
        log.Stderrf(err.String())
        return
    }

    for {
        fd, err := l.Accept()
        if err != nil {
            log.Stderrf(err.String())
            break
        }
        go handleFcgiRequest(fd)

    }
}
<|endoftext|>


<commit_before>package contentaddressable

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"os"
	"path/filepath"
)

var (
	AlreadyClosed = errors.New("Already closed.")
	HasData       = errors.New("Destination file already has data.")
)

// File handles the atomic writing of a content addressable file.  It writes to
// a temp file, and then renames to the final location after Accept().
type File struct {
	Oid          string
	filename     string
	tempFilename string
	file         *os.File
	tempFile     *os.File
	hasher       hash.Hash
}

// NewFile initializes a content addressable file for writing.  It opens both
// the given filename, and a temp filename in exclusive mode.  The *File OID
// is taken from the base name of the given filename.
func NewFile(filename string) (*File, error) {
	oid := filepath.Base(filename)
	dir := filepath.Dir(filename)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, err
	}

	tempFilename := filename + "-temp"
	tempFile, err := os.OpenFile(tempFilename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0644)
	if err != nil {
		return nil, err
	}

	file, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0644)
	if err != nil {
		cleanupFile(tempFile)
		return nil, err
	}

	caw := &File{
		Oid:          oid,
		filename:     filename,
		tempFilename: tempFilename,
		file:         file,
		tempFile:     tempFile,
		hasher:       sha256.New(),
	}

	return caw, nil
}

// Write sends data to the temporary file.
func (w *File) Write(p []byte) (int, error) {
	if w.Closed() {
		return 0, AlreadyClosed
	}

	w.hasher.Write(p)
	return w.tempFile.Write(p)
}

// Accept verifies the written content SHA-256 signature matches the given OID.
// If it matches, the temp file is renamed to the original filename.  If not,
// an error is returned.
func (w *File) Accept() error {
	if w.Closed() {
		return AlreadyClosed
	}

	sig := hex.EncodeToString(w.hasher.Sum(nil))
	if sig != w.Oid {
		return fmt.Errorf("Content mismatch.  Expected OID %s, got %s", w.Oid, sig)
	}

	if err := cleanupFile(w.file); err != nil {
		return err
	}
	w.file = nil

	w.tempFile.Close()
	err := os.Rename(w.tempFilename, w.filename)
	w.Close()
	return err
}

// Close cleans up the internal file objects.
func (w *File) Close() error {
	if w.tempFile != nil {
		if err := cleanupFile(w.tempFile); err != nil {
			return err
		}
		w.tempFile = nil
	}

	if w.file != nil {
		if err := cleanupFile(w.file); err != nil {
			return err
		}
		w.file = nil
	}

	return nil
}

// Closed reports whether this file object has been closed.
func (w *File) Closed() bool {
	if w.tempFile == nil || w.file == nil {
		return true
	}
	return false
}

func cleanupFile(f *os.File) error {
	err := f.Close()
	if err := os.RemoveAll(f.Name()); err != nil {
		return err
	}

	return err
}
<commit_msg>don't try to close the temp file twice, otherwise Close() returns an error<commit_after>package contentaddressable

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"os"
	"path/filepath"
)

var (
	AlreadyClosed = errors.New("Already closed.")
	HasData       = errors.New("Destination file already has data.")
)

// File handles the atomic writing of a content addressable file.  It writes to
// a temp file, and then renames to the final location after Accept().
type File struct {
	Oid          string
	filename     string
	tempFilename string
	file         *os.File
	tempFile     *os.File
	hasher       hash.Hash
}

// NewFile initializes a content addressable file for writing.  It opens both
// the given filename, and a temp filename in exclusive mode.  The *File OID
// is taken from the base name of the given filename.
func NewFile(filename string) (*File, error) {
	oid := filepath.Base(filename)
	dir := filepath.Dir(filename)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, err
	}

	tempFilename := filename + "-temp"
	tempFile, err := os.OpenFile(tempFilename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0644)
	if err != nil {
		return nil, err
	}

	file, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0644)
	if err != nil {
		cleanupFile(tempFile)
		return nil, err
	}

	caw := &File{
		Oid:          oid,
		filename:     filename,
		tempFilename: tempFilename,
		file:         file,
		tempFile:     tempFile,
		hasher:       sha256.New(),
	}

	return caw, nil
}

// Write sends data to the temporary file.
func (w *File) Write(p []byte) (int, error) {
	if w.Closed() {
		return 0, AlreadyClosed
	}

	w.hasher.Write(p)
	return w.tempFile.Write(p)
}

// Accept verifies the written content SHA-256 signature matches the given OID.
// If it matches, the temp file is renamed to the original filename.  If not,
// an error is returned.
func (w *File) Accept() error {
	if w.Closed() {
		return AlreadyClosed
	}

	sig := hex.EncodeToString(w.hasher.Sum(nil))
	if sig != w.Oid {
		return fmt.Errorf("Content mismatch.  Expected OID %s, got %s", w.Oid, sig)
	}

	if err := cleanupFile(w.file); err != nil {
		return err
	}
	w.file = nil

	// flush any data to disk
	w.tempFile.Close()

	// rename the temp file to the real file
	if err := os.Rename(w.tempFilename, w.filename); err != nil {
		return w.Close()
	} else {
		// if successful, set the closed *os.File to nil
		w.tempFile = nil
		return nil
	}
}

// Close cleans up the internal file objects.
func (w *File) Close() error {
	if w.tempFile != nil {
		if err := cleanupFile(w.tempFile); err != nil {
			return err
		}
		w.tempFile = nil
	}

	if w.file != nil {
		if err := cleanupFile(w.file); err != nil {
			return err
		}
		w.file = nil
	}

	return nil
}

// Closed reports whether this file object has been closed.
func (w *File) Closed() bool {
	if w.tempFile == nil || w.file == nil {
		return true
	}
	return false
}

func cleanupFile(f *os.File) error {
	err := f.Close()
	if err := os.RemoveAll(f.Name()); err != nil {
		return err
	}

	return err
}
<|endoftext|>


<commit_before>package multiconfig

import (
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/BurntSushi/toml"
	yaml "gopkg.in/yaml.v2"
)

var (
	// ErrSourceNotSet states that neither the path or the reader is set on the loader
	ErrSourceNotSet = errors.New("config path or reader is not set")

	// ErrFileNotFound states that given file is not exists
	ErrFileNotFound = errors.New("config file not found")
)

// TOMLLoader satisifies the loader interface. It loads the configuration from
// the given toml file or Reader
type TOMLLoader struct {
	Path   string
	Reader io.Reader
}

// Load loads the source into the config defined by struct s
// Defaults to using the Reader if provided, otherwise tries to read from the
// file
func (t *TOMLLoader) Load(s interface{}) error {
	var r io.Reader

	if t.Reader != nil {
		r = t.Reader
	} else if t.Path != "" {
		file, err := getConfig(t.Path)
		if err != nil {
			return err
		}
		defer file.Close()
		r = file
	} else {
		return ErrSourceNotSet
	}

	if _, err := toml.DecodeReader(r, s); err != nil {
		return err
	}

	return nil
}

// JSONLoader satisifies the loader interface. It loads the configuration from
// the given json file or Reader
type JSONLoader struct {
	Path   string
	Reader io.Reader
}

// Load loads the source into the config defined by struct s
// Defaults to using the Reader if provided, otherwise tries to read from the
// file
func (j *JSONLoader) Load(s interface{}) error {
	var r io.Reader
	if j.Reader != nil {
		r = j.Reader
	} else if j.Path != "" {
		file, err := getConfig(j.Path)
		if err != nil {
			return err
		}
		defer file.Close()
		r = file
	} else {
		return ErrSourceNotSet
	}

	return json.NewDecoder(r).Decode(s)
}

// YAMLLoader satisifies the loader interface. It loads the configuration from
// the given yaml file
type YAMLLoader struct {
	Path   string
	Reader io.Reader
}

// Load loads the source into the config defined by struct s
// Defaults to using the Reader if provided, otherwise tries to read from the
// file
func (y *YAMLLoader) Load(s interface{}) error {
	var data []byte
	var err error

	if y.Reader != nil {
		data, err = ioutil.ReadAll(y.Reader)

	} else if y.Path != "" {

		var file *os.File

		file, err = getConfig(y.Path)
		if err != nil {
			return err
		}

		data, err = ioutil.ReadAll(file)

		file.Close()

	} else {
		return ErrSourceNotSet
	}

	if err != nil {
		return err
	}

	if err = yaml.Unmarshal(data, s); err != nil {
		return err
	}

	return nil
}

func getConfig(path string) (*os.File, error) {
	pwd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	configPath := path
	if !filepath.IsAbs(path) {
		configPath = filepath.Join(pwd, path)
	}

	// check if file with combined path is exists(relative path)
	if _, err := os.Stat(configPath); !os.IsNotExist(err) {
		return os.Open(configPath)
	}

	f, err := os.Open(path)
	if os.IsNotExist(err) {
		return nil, ErrFileNotFound
	}
	return f, err
}
<commit_msg>unified the code style for YAMLLoader.Load<commit_after>package multiconfig

import (
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/BurntSushi/toml"
	yaml "gopkg.in/yaml.v2"
)

var (
	// ErrSourceNotSet states that neither the path or the reader is set on the loader
	ErrSourceNotSet = errors.New("config path or reader is not set")

	// ErrFileNotFound states that given file is not exists
	ErrFileNotFound = errors.New("config file not found")
)

// TOMLLoader satisifies the loader interface. It loads the configuration from
// the given toml file or Reader
type TOMLLoader struct {
	Path   string
	Reader io.Reader
}

// Load loads the source into the config defined by struct s
// Defaults to using the Reader if provided, otherwise tries to read from the
// file
func (t *TOMLLoader) Load(s interface{}) error {
	var r io.Reader

	if t.Reader != nil {
		r = t.Reader
	} else if t.Path != "" {
		file, err := getConfig(t.Path)
		if err != nil {
			return err
		}
		defer file.Close()
		r = file
	} else {
		return ErrSourceNotSet
	}

	if _, err := toml.DecodeReader(r, s); err != nil {
		return err
	}

	return nil
}

// JSONLoader satisifies the loader interface. It loads the configuration from
// the given json file or Reader
type JSONLoader struct {
	Path   string
	Reader io.Reader
}

// Load loads the source into the config defined by struct s
// Defaults to using the Reader if provided, otherwise tries to read from the
// file
func (j *JSONLoader) Load(s interface{}) error {
	var r io.Reader
	if j.Reader != nil {
		r = j.Reader
	} else if j.Path != "" {
		file, err := getConfig(j.Path)
		if err != nil {
			return err
		}
		defer file.Close()
		r = file
	} else {
		return ErrSourceNotSet
	}

	return json.NewDecoder(r).Decode(s)
}

// YAMLLoader satisifies the loader interface. It loads the configuration from
// the given yaml file
type YAMLLoader struct {
	Path   string
	Reader io.Reader
}

// Load loads the source into the config defined by struct s
// Defaults to using the Reader if provided, otherwise tries to read from the
// file
func (y *YAMLLoader) Load(s interface{}) error {
	var r io.Reader

	if y.Reader != nil {
		r = y.Reader
	} else if y.Path != "" {
		file, err := getConfig(y.Path)
		if err != nil {
			return err
		}
		defer file.Close()
		r = file
	} else {
		return ErrSourceNotSet
	}

	data, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}

	return yaml.Unmarshal(data, s)
}

func getConfig(path string) (*os.File, error) {
	pwd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	configPath := path
	if !filepath.IsAbs(path) {
		configPath = filepath.Join(pwd, path)
	}

	// check if file with combined path is exists(relative path)
	if _, err := os.Stat(configPath); !os.IsNotExist(err) {
		return os.Open(configPath)
	}

	f, err := os.Open(path)
	if os.IsNotExist(err) {
		return nil, ErrFileNotFound
	}
	return f, err
}
<|endoftext|>