<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

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

   This program 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 Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package config

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"time"

	"github.com/fatih/color"
	"github.com/webx-top/codec"
	"github.com/webx-top/com"
	"github.com/webx-top/db"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/middleware/bytes"
	"github.com/webx-top/echo/middleware/language"

	"github.com/admpub/confl"
	"github.com/admpub/log"
	"github.com/admpub/nging/v4/application/library/config/extend"
	"github.com/admpub/nging/v4/application/library/config/subconfig/scookie"
	"github.com/admpub/nging/v4/application/library/config/subconfig/scron"
	"github.com/admpub/nging/v4/application/library/config/subconfig/sdb"
	"github.com/admpub/nging/v4/application/library/config/subconfig/ssystem"
	"github.com/admpub/securecookie"
)

func NewConfig() *Config {
	c := &Config{}
	c.InitExtend()
	c.Settings = NewSettings(c)
	c.DB.MaxIdleConns = db.DefaultSettings.MaxIdleConns()
	c.DB.MaxOpenConns = db.DefaultSettings.MaxOpenConns()
	return c
}

type Config struct {
	DB       sdb.DB          `json:"db"`
	Sys      ssystem.System  `json:"sys"`
	Cron     scron.Cron      `json:"cron"`
	Cookie   scookie.Config  `json:"cookie"`
	Language language.Config `json:"language"`
	Download struct {
		SavePath string `json:"savePath"`
	} `json:"download"`
	//License lib.LicenseData `json:"license,omitempty"`
	Extend    echo.H `json:"extend,omitempty"`
	*Settings `json:"-"`

	connectedDB bool
}

func (c *Config) IsEnv(env string) bool {
	return c.Sys.IsEnv(env)
}

func (c *Config) IsEnvProd() bool {
	return c.Sys.IsEnv(`prod`)
}

func (c *Config) IsEnvDev() bool {
	return c.Sys.IsEnv(`dev`)
}

func (c *Config) GetMaxRequestBodySize() int {
	if c.MaxRequestBodySize > 0 {
		return c.MaxRequestBodySize
	}
	if c.Sys.MaxRequestBodySize > 0 {
		return c.Sys.MaxRequestBodySize
	}
	return defaultMaxRequestBodyBytes
}

// ConnectedDB 数据库是否已连接，如果没有连接则自动连接
func (c *Config) ConnectedDB(autoConn ...bool) bool {
	if c.connectedDB {
		return c.connectedDB
	}
	n := len(autoConn)
	if n == 0 || (n > 0 && autoConn[0]) {
		err := c.connectDB()
		if err != nil {
			log.Error(err)
		}
	}
	return c.connectedDB
}

func (c *Config) connectDB() error {
	err := ConnectDB(c)
	if err != nil {
		return err
	}
	c.connectedDB = true
	return nil
}

func (c *Config) APIKey() string {
	return c.Settings.APIKey
}

func (c *Config) CookieConfig() scookie.Config {
	return c.Cookie
}

func (c *Config) InitExtend() *Config {
	c.Extend = echo.H{}
	extend.Range(c.registerExtend)
	return c
}

func (c *Config) registerExtend(key string, recv interface{}) {
	fmt.Printf(color.YellowString(`[Register Extend Config]`)+` `+color.MagentaString(`P%d`, DefaultCLIConfig.Pid())+` %s: %T`+"\n", key, recv)
	c.Extend[key] = recv
}

func (c *Config) UnregisterExtend(key string) {
	if recv, ok := c.Extend[key]; ok {
		fmt.Printf(color.YellowString(`[Unregister Extend Config]`)+` `+color.MagentaString(`P%d`, DefaultCLIConfig.Pid())+` %s: %T`+"\n", key, recv)
		delete(c.Extend, key)
	}
}

func (c *Config) ConfigFromDB() echo.H {
	return c.Settings.GetConfig()
}

func (c *Config) SetDebug(on bool) *Config {
	c.Settings.SetDebug(on)
	return c
}

func (c *Config) Codec(lengths ...int) codec.Codec {
	length := 128
	if len(lengths) > 0 {
		length = lengths[0]
	}
	if length == 256 {
		return default256Codec
	}
	return defaultCodec
}

var (
	defaultCodec    = codec.NewAesCrypto(`AES-128-CBC`)
	default256Codec = codec.NewAesCrypto(`AES-256-CBC`)
)

func (c *Config) Encode(raw string, keys ...string) string {
	var key string
	if len(keys) > 0 && len(keys[0]) > 0 {
		key = com.Md5(keys[0])
	} else {
		key = c.Cookie.HashKey
	}
	return c.Codec().Encode(raw, key)
}

func (c *Config) Decode(encrypted string, keys ...string) string {
	if len(encrypted) == 0 {
		return ``
	}
	var key string
	if len(keys) > 0 && len(keys[0]) > 0 {
		key = com.Md5(keys[0])
	} else {
		key = c.Cookie.HashKey
	}
	return c.Codec().Decode(encrypted, key)
}

func (c *Config) InitSecretKey() *Config {
	c.Cookie.BlockKey = c.GenerateRandomKey()
	c.Cookie.HashKey = c.GenerateRandomKey()
	return c
}

func (c *Config) GenerateRandomKey() string {
	return com.Md5(string(securecookie.GenerateRandomKey(32)))
}

func (c *Config) Reload(newConfig *Config) error {
	engines := []string{}
	for name, value := range newConfig.Extend {
		extConfig := c.Extend.Get(name)
		if !reflect.DeepEqual(value, extConfig) {
			if rd, ok := extConfig.(ReloadByConfig); ok {
				if err := rd.ReloadByConfig(newConfig); err != nil {
					log.Error(err)
				}
				continue
			}
			if rd, ok := extConfig.(extend.Reloader); ok {
				if err := rd.Reload(); err != nil {
					log.Error(err)
				}
				continue
			}
			engines = append(engines, name)
		}
	}
	return DefaultCLIConfig.Reload(newConfig, engines...)
}

func (c *Config) AsDefault() {
	echo.Set(`DefaultConfig`, c)
	DefaultConfig = c
	err := c.Settings.Init(nil)
	if err != nil {
		log.Error(err)
	}
}

func (c *Config) SaveToFile() error {
	b, err := confl.Marshal(c)
	if err != nil {
		return err
	}
	dir := filepath.Dir(DefaultCLIConfig.Conf)
	err = com.MkdirAll(dir, os.ModePerm)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(DefaultCLIConfig.Conf, b, os.ModePerm)
	return err
}

func (c *Config) GenerateSample() error {
	_, err := os.Stat(DefaultCLIConfig.Conf + `.sample`)
	if err == nil || !os.IsNotExist(err) {
		return err
	}
	var old []byte
	old, err = ioutil.ReadFile(DefaultCLIConfig.Conf)
	if err == nil {
		err = ioutil.WriteFile(DefaultCLIConfig.Conf+`.sample`, old, os.ModePerm)
	}
	return err
}

func (c *Config) SetDefaults(configFile string) {
	confDir := filepath.Dir(configFile)
	if len(c.Sys.VhostsfileDir) == 0 {
		c.Sys.VhostsfileDir = filepath.Join(confDir, `vhosts`)
	}
	if c.Sys.MaxRequestBodySize <= 0 {
		c.Sys.MaxRequestBodySize = defaultMaxRequestBodyBytes
	}
	if c.Sys.EditableFileMaxBytes < 1 && len(c.Sys.EditableFileMaxSize) > 0 {
		var err error
		c.Sys.EditableFileMaxBytes, err = bytes.Parse(c.Sys.EditableFileMaxSize)
		if err != nil {
			log.Error(err.Error())
		}
	}
	c.Sys.CmdTimeoutDuration = ParseTimeDuration(c.Sys.CmdTimeout)
	if c.Sys.CmdTimeoutDuration <= 0 {
		c.Sys.CmdTimeoutDuration = time.Second * 30
	}
	if len(c.Cookie.Path) == 0 {
		c.Cookie.Path = `/`
	}
	if len(c.Sys.SSLCacheDir) == 0 {
		c.Sys.SSLCacheDir = filepath.Join(echo.Wd(), `data`, `cache`, `autocert`)
	}
	if len(c.Cookie.Path) == 0 {
		c.Cookie.Path = `/`
	}
	for _, value := range c.Extend {
		if sd, ok := value.(extend.SetDefaults); ok {
			sd.SetDefaults()
		}
	}
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

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

   This program 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 Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package config

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"time"

	"github.com/admpub/confl"
	"github.com/admpub/log"
	"github.com/admpub/securecookie"
	"github.com/fatih/color"
	"github.com/webx-top/codec"
	"github.com/webx-top/com"
	"github.com/webx-top/db"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/middleware/bytes"
	"github.com/webx-top/echo/middleware/language"

	"github.com/admpub/nging/v4/application/library/config/extend"
	"github.com/admpub/nging/v4/application/library/config/subconfig/scookie"
	"github.com/admpub/nging/v4/application/library/config/subconfig/scron"
	"github.com/admpub/nging/v4/application/library/config/subconfig/sdb"
	"github.com/admpub/nging/v4/application/library/config/subconfig/ssystem"
)

func NewConfig() *Config {
	c := &Config{}
	c.InitExtend()
	c.Settings = NewSettings(c)
	c.DB.MaxIdleConns = db.DefaultSettings.MaxIdleConns()
	c.DB.MaxOpenConns = db.DefaultSettings.MaxOpenConns()
	return c
}

type Config struct {
	DB        sdb.DB          `json:"db"`
	Sys       ssystem.System  `json:"sys"`
	Cron      scron.Cron      `json:"cron"`
	Cookie    scookie.Config  `json:"cookie"`
	Language  language.Config `json:"language"`
	Extend    echo.H          `json:"extend,omitempty"`
	*Settings `json:"-"`

	connectedDB bool
}

func (c *Config) IsEnv(env string) bool {
	return c.Sys.IsEnv(env)
}

func (c *Config) IsEnvProd() bool {
	return c.Sys.IsEnv(`prod`)
}

func (c *Config) IsEnvDev() bool {
	return c.Sys.IsEnv(`dev`)
}

func (c *Config) GetMaxRequestBodySize() int {
	if c.MaxRequestBodySize > 0 {
		return c.MaxRequestBodySize
	}
	if c.Sys.MaxRequestBodySize > 0 {
		return c.Sys.MaxRequestBodySize
	}
	return defaultMaxRequestBodyBytes
}

// ConnectedDB 数据库是否已连接，如果没有连接则自动连接
func (c *Config) ConnectedDB(autoConn ...bool) bool {
	if c.connectedDB {
		return c.connectedDB
	}
	n := len(autoConn)
	if n == 0 || (n > 0 && autoConn[0]) {
		err := c.connectDB()
		if err != nil {
			log.Error(err)
		}
	}
	return c.connectedDB
}

func (c *Config) connectDB() error {
	err := ConnectDB(c)
	if err != nil {
		return err
	}
	c.connectedDB = true
	return nil
}

func (c *Config) APIKey() string {
	return c.Settings.APIKey
}

func (c *Config) CookieConfig() scookie.Config {
	return c.Cookie
}

func (c *Config) InitExtend() *Config {
	c.Extend = echo.H{}
	extend.Range(c.registerExtend)
	return c
}

func (c *Config) registerExtend(key string, recv interface{}) {
	fmt.Printf(color.YellowString(`[Register Extend Config]`)+` `+color.MagentaString(`P%d`, DefaultCLIConfig.Pid())+` %s: %T`+"\n", key, recv)
	c.Extend[key] = recv
}

func (c *Config) UnregisterExtend(key string) {
	if recv, ok := c.Extend[key]; ok {
		fmt.Printf(color.YellowString(`[Unregister Extend Config]`)+` `+color.MagentaString(`P%d`, DefaultCLIConfig.Pid())+` %s: %T`+"\n", key, recv)
		delete(c.Extend, key)
	}
}

func (c *Config) ConfigFromDB() echo.H {
	return c.Settings.GetConfig()
}

func (c *Config) SetDebug(on bool) *Config {
	c.Settings.SetDebug(on)
	return c
}

func (c *Config) Codec(lengths ...int) codec.Codec {
	length := 128
	if len(lengths) > 0 {
		length = lengths[0]
	}
	if length == 256 {
		return default256Codec
	}
	return defaultCodec
}

var (
	defaultCodec    = codec.NewAesCrypto(`AES-128-CBC`)
	default256Codec = codec.NewAesCrypto(`AES-256-CBC`)
)

func (c *Config) Encode(raw string, keys ...string) string {
	var key string
	if len(keys) > 0 && len(keys[0]) > 0 {
		key = com.Md5(keys[0])
	} else {
		key = c.Cookie.HashKey
	}
	return c.Codec().Encode(raw, key)
}

func (c *Config) Decode(encrypted string, keys ...string) string {
	if len(encrypted) == 0 {
		return ``
	}
	var key string
	if len(keys) > 0 && len(keys[0]) > 0 {
		key = com.Md5(keys[0])
	} else {
		key = c.Cookie.HashKey
	}
	return c.Codec().Decode(encrypted, key)
}

func (c *Config) InitSecretKey() *Config {
	c.Cookie.BlockKey = c.GenerateRandomKey()
	c.Cookie.HashKey = c.GenerateRandomKey()
	return c
}

func (c *Config) GenerateRandomKey() string {
	return com.Md5(string(securecookie.GenerateRandomKey(32)))
}

func (c *Config) Reload(newConfig *Config) error {
	engines := []string{}
	for name, value := range newConfig.Extend {
		extConfig := c.Extend.Get(name)
		if !reflect.DeepEqual(value, extConfig) {
			if rd, ok := extConfig.(ReloadByConfig); ok {
				if err := rd.ReloadByConfig(newConfig); err != nil {
					log.Error(err)
				}
				continue
			}
			if rd, ok := extConfig.(extend.Reloader); ok {
				if err := rd.Reload(); err != nil {
					log.Error(err)
				}
				continue
			}
			engines = append(engines, name)
		}
	}
	return DefaultCLIConfig.Reload(newConfig, engines...)
}

func (c *Config) AsDefault() {
	echo.Set(`DefaultConfig`, c)
	DefaultConfig = c
	err := c.Settings.Init(nil)
	if err != nil {
		log.Error(err)
	}
}

func (c *Config) SaveToFile() error {
	b, err := confl.Marshal(c)
	if err != nil {
		return err
	}
	dir := filepath.Dir(DefaultCLIConfig.Conf)
	err = com.MkdirAll(dir, os.ModePerm)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(DefaultCLIConfig.Conf, b, os.ModePerm)
	return err
}

func (c *Config) GenerateSample() error {
	_, err := os.Stat(DefaultCLIConfig.Conf + `.sample`)
	if err == nil || !os.IsNotExist(err) {
		return err
	}
	var old []byte
	old, err = ioutil.ReadFile(DefaultCLIConfig.Conf)
	if err == nil {
		err = ioutil.WriteFile(DefaultCLIConfig.Conf+`.sample`, old, os.ModePerm)
	}
	return err
}

func (c *Config) SetDefaults(configFile string) {
	confDir := filepath.Dir(configFile)
	if len(c.Sys.VhostsfileDir) == 0 {
		c.Sys.VhostsfileDir = filepath.Join(confDir, `vhosts`)
	}
	if c.Sys.MaxRequestBodySize <= 0 {
		c.Sys.MaxRequestBodySize = defaultMaxRequestBodyBytes
	}
	if c.Sys.EditableFileMaxBytes < 1 && len(c.Sys.EditableFileMaxSize) > 0 {
		var err error
		c.Sys.EditableFileMaxBytes, err = bytes.Parse(c.Sys.EditableFileMaxSize)
		if err != nil {
			log.Error(err.Error())
		}
	}
	c.Sys.CmdTimeoutDuration = ParseTimeDuration(c.Sys.CmdTimeout)
	if c.Sys.CmdTimeoutDuration <= 0 {
		c.Sys.CmdTimeoutDuration = time.Second * 30
	}
	if len(c.Cookie.Path) == 0 {
		c.Cookie.Path = `/`
	}
	if len(c.Sys.SSLCacheDir) == 0 {
		c.Sys.SSLCacheDir = filepath.Join(echo.Wd(), `data`, `cache`, `autocert`)
	}
	if len(c.Cookie.Path) == 0 {
		c.Cookie.Path = `/`
	}
	for _, value := range c.Extend {
		if sd, ok := value.(extend.SetDefaults); ok {
			sd.SetDefaults()
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
  ".."
  "strconv"
  "runtime"
  "fmt"
  "bytes"
)

// quick test for concurrent log messages
func main() {
  fmt.Printf("%d -> %d\n", runtime.NumCPU, runtime.GOMAXPROCS(1000))

  buf := bytes.NewBuffer([]byte(""))
  logger := scrolls.NewLogger(buf)
  chans := make([]chan bool, 1000)

  for i := range chans {
    chans[i] = make(chan bool)
    go LogLikeCrazy(strconv.Itoa(i), logger, chans[i])
  }

  for _, ch := range chans {
    <- ch
  }
  fmt.Println(buf.String())
}

func LogLikeCrazy(key string, logger scrolls.Logger, stopped chan bool) {
  for i := 0; i < 10; i += 1 {
    logger.Log(scrolls.LogData{"key": key, "i": strconv.Itoa(i+1)})
  }
  stopped <- true
}
<commit_msg>the logger breaks harder stdout<commit_after>package main

import (
  ".."
  "strconv"
  "runtime"
  "fmt"
)

// quick test for concurrent log messages
func main() {
  fmt.Printf("%d -> %d\n", runtime.NumCPU, runtime.GOMAXPROCS(1000))

  logger := scrolls.NewLogger(nil)
  chans := make([]chan bool, 1000)

  for i := range chans {
    chans[i] = make(chan bool)
    go LogLikeCrazy(strconv.Itoa(i), logger, chans[i])
  }

  for _, ch := range chans {
    <- ch
  }
}

func LogLikeCrazy(key string, logger scrolls.Logger, stopped chan bool) {
  for i := 0; i < 10; i += 1 {
    logger.Log(scrolls.LogData{"key": key, "i": strconv.Itoa(i+1)})
  }
  stopped <- true
}
<|endoftext|>


<commit_before>package gist8065433

import "net"

// GetPublicIps returns a string slice of non-loopback IPs.
func GetPublicIps() (publicIps []string, err error) {
	ifis, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	for _, ifi := range ifis {
		addrs, err := ifi.Addrs()
		if err != nil {
			return nil, err
		}

		for _, addr := range addrs {
			ipNet, ok := addr.(*net.IPNet)
			if !ok {
				continue
			}

			ip4 := ipNet.IP.To4()
			if ip4 == nil || ip4.IsLoopback() {
				continue
			}

			publicIps = append(publicIps, ipNet.IP.String())
		}
	}

	return publicIps, nil
}
<commit_msg>Add GetAllIps to get a list of all IPv4 interfaces, including loopback.<commit_after>package gist8065433

import "net"

// GetPublicIps returns a string slice of non-loopback IPs.
func GetPublicIps() (publicIps []string, err error) {
	ifis, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	for _, ifi := range ifis {
		addrs, err := ifi.Addrs()
		if err != nil {
			return nil, err
		}

		for _, addr := range addrs {
			ipNet, ok := addr.(*net.IPNet)
			if !ok {
				continue
			}

			ip4 := ipNet.IP.To4()
			if ip4 == nil || ip4.IsLoopback() {
				continue
			}

			publicIps = append(publicIps, ipNet.IP.String())
		}
	}

	return publicIps, nil
}

// GetAllIps returns a string slice of all IPs.
func GetAllIps() (ips []string, err error) {
	ifis, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	for _, ifi := range ifis {
		addrs, err := ifi.Addrs()
		if err != nil {
			return nil, err
		}

		for _, addr := range addrs {
			ipNet, ok := addr.(*net.IPNet)
			if !ok {
				continue
			}

			ip4 := ipNet.IP.To4()
			if ip4 == nil {
				continue
			}

			ips = append(ips, ipNet.IP.String())
		}
	}

	return ips, nil
}
<|endoftext|>


<commit_before>package download

import (
	"os"
	"path/filepath"
	"testing"
)

var testDownloadDir = "test_downloads"

func TestRepoRootForImportPath(t *testing.T) {
	cases := []struct {
		giveURL  string
		wantPath string
		wantVCS  string
	}{
		{"github.com/gojp/goreportcard", "github.com/gojp/goreportcard", "git"},
		{"https://github.com/boltdb/bolt", "github.com/boltdb/bolt", "git"},
		{"https://bitbucket.org/rickb777/go-talk", "bitbucket.org/rickb777/go-talk", "hg"},
		{"ssh://hg@bitbucket.org/rickb777/go-talk", "bitbucket.org/rickb777/go-talk", "hg"},
	}

	for _, tt := range cases {
		root, err := Download(tt.giveURL, testDownloadDir)
		if err != nil {
			t.Fatalf("Error calling Download(%q): %v", tt.giveURL, err)
		}

		if root.Root != tt.wantPath {
			t.Errorf("Download(%q): root.Repo = %q, want %q", tt.giveURL, root.Repo, tt.wantPath)
		}

		wantPath := filepath.Join(testDownloadDir, tt.wantPath)
		ex, _ := exists(wantPath)
		if !ex {
			t.Errorf("Download(%q): %q was not created", tt.giveURL, wantPath)
		}
	}

	// clean up the test
	os.RemoveAll(testDownloadDir)
}
<commit_msg>skip download test<commit_after>package download

import (
	"os"
	"path/filepath"
	"testing"
)

var testDownloadDir = "test_downloads"

func TestRepoRootForImportPath(t *testing.T) {
	t.Skip()
	cases := []struct {
		giveURL  string
		wantPath string
		wantVCS  string
	}{
		{"github.com/gojp/goreportcard", "github.com/gojp/goreportcard", "git"},
		{"https://github.com/boltdb/bolt", "github.com/boltdb/bolt", "git"},
		{"https://bitbucket.org/rickb777/go-talk", "bitbucket.org/rickb777/go-talk", "hg"},
		{"ssh://hg@bitbucket.org/rickb777/go-talk", "bitbucket.org/rickb777/go-talk", "hg"},
	}

	for _, tt := range cases {
		root, err := Download(tt.giveURL, testDownloadDir)
		if err != nil {
			t.Fatalf("Error calling Download(%q): %v", tt.giveURL, err)
		}

		if root.Root != tt.wantPath {
			t.Errorf("Download(%q): root.Repo = %q, want %q", tt.giveURL, root.Repo, tt.wantPath)
		}

		wantPath := filepath.Join(testDownloadDir, tt.wantPath)
		ex, _ := exists(wantPath)
		if !ex {
			t.Errorf("Download(%q): %q was not created", tt.giveURL, wantPath)
		}
	}

	// clean up the test
	os.RemoveAll(testDownloadDir)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"path"
	"runtime"
	"testing"
)

type ParseExample struct {
	raw           string
	expectedTerms []string
	expectedId    string
	expectedError error
}

var parseExamples = []ParseExample{
	{
		"{}",
		make([]string, 0),
		"",
		nil,
	},
	{
		`{"terms":["Government digital service","GDS"],"id":"1"}`,
		[]string{"Government digital service", "GDS"},
		"1",
		nil,
	},
}

func fixturePath(fixtureFile string) string {
	_, filename, _, _ := runtime.Caller(1)
	return path.Join(path.Dir(filename), "data", fixtureFile)
}

func exampleConfig() *Config {
	return &Config{
		extractAddress:     ":3096",
		dbConnectionString: "host=/var/run/postgresql dbname=entity-extractor_test sslmode=disable",
		logPath:            "STDERR",
	}
}

func TestLoadEntities(t *testing.T) {
	config := exampleConfig()
	extractor := NewExtractor(config)
	err := extractor.LoadEntities()
	if err != nil {
		fmt.Println(err)
		t.FailNow()
	}

	expectedTerms := [...]string{
		"Government digital service",
		"GDS",
		"Ministry of Justice",
		"MoJ",
	}
	require.Equal(t, len(expectedTerms), len(extractor.entities.terms))
	for i, expectedTerm := range expectedTerms {
		assert.Equal(t, expectedTerm, extractor.entities.terms[i])
	}
	expectedIds := [...]string{
		"1",
		"1",
		"2",
		"2",
	}
	for i, expectedId := range expectedIds {
		assert.Equal(t, expectedId, extractor.entities.termOffsetToEntity[i].id)
	}
}

type ExtractionExample struct {
	comment         string
	document        string
	expectedTermIds []string
}

var extractionExamples = []ExtractionExample{
	{
		"a document matching a single word term",
		"This document mentions GDS but it doesn't mention the Ministry of J...",
		[]string{"1"},
	},
	{
		"a document matching a multi-word term",
		"Government digital service",
		[]string{"1"},
	},
	{
		"terms are matched case sensitively",
		"gds",
		[]string{},
	},
}

func TestExtract(t *testing.T) {
	config := exampleConfig()
	extractor := NewExtractor(config)
	err := extractor.LoadEntities()
	require.Nil(t, err)

	for _, example := range extractionExamples {
		matchedTermIds := extractor.Extract(example.document)
		assert.Equal(t, example.expectedTermIds, matchedTermIds, example.comment)
	}
}
<commit_msg>More helpful error reporting on connection failures in test<commit_after>package main

import (
	"fmt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"path"
	"runtime"
	"testing"
)

type ParseExample struct {
	raw           string
	expectedTerms []string
	expectedId    string
	expectedError error
}

var parseExamples = []ParseExample{
	{
		"{}",
		make([]string, 0),
		"",
		nil,
	},
	{
		`{"terms":["Government digital service","GDS"],"id":"1"}`,
		[]string{"Government digital service", "GDS"},
		"1",
		nil,
	},
}

func fixturePath(fixtureFile string) string {
	_, filename, _, _ := runtime.Caller(1)
	return path.Join(path.Dir(filename), "data", fixtureFile)
}

func exampleConfig() *Config {
	return &Config{
		extractAddress:     ":3096",
		dbConnectionString: "host=/var/run/postgresql dbname=entity-extractor_test sslmode=disable",
		logPath:            "STDERR",
	}
}

func TestLoadEntities(t *testing.T) {
	config := exampleConfig()
	extractor := NewExtractor(config)
	err := extractor.LoadEntities()
	if err != nil {
		fmt.Println(err)
		t.FailNow()
	}

	expectedTerms := [...]string{
		"Government digital service",
		"GDS",
		"Ministry of Justice",
		"MoJ",
	}
	require.Equal(t, len(expectedTerms), len(extractor.entities.terms))
	for i, expectedTerm := range expectedTerms {
		assert.Equal(t, expectedTerm, extractor.entities.terms[i])
	}
	expectedIds := [...]string{
		"1",
		"1",
		"2",
		"2",
	}
	for i, expectedId := range expectedIds {
		assert.Equal(t, expectedId, extractor.entities.termOffsetToEntity[i].id)
	}
}

type ExtractionExample struct {
	comment         string
	document        string
	expectedTermIds []string
}

var extractionExamples = []ExtractionExample{
	{
		"a document matching a single word term",
		"This document mentions GDS but it doesn't mention the Ministry of J...",
		[]string{"1"},
	},
	{
		"a document matching a multi-word term",
		"Government digital service",
		[]string{"1"},
	},
	{
		"terms are matched case sensitively",
		"gds",
		[]string{},
	},
}

func TestExtract(t *testing.T) {
	config := exampleConfig()
	extractor := NewExtractor(config)
	err := extractor.LoadEntities()
	if err != nil {
		fmt.Println("ERROR: %v", err)
		t.FailNow()
	}

	for _, example := range extractionExamples {
		matchedTermIds := extractor.Extract(example.document)
		assert.Equal(t, example.expectedTermIds, matchedTermIds, example.comment)
	}
}
<|endoftext|>


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

package prod

import (
	"fmt"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"

	"github.com/luci/gae/service/info"
	"github.com/luci/gae/service/urlfetch"
	"golang.org/x/net/context"
	gOAuth "golang.org/x/oauth2/google"
	"google.golang.org/appengine"
	"google.golang.org/appengine/remote_api"
)

type key int

var (
	prodContextKey      key
	prodContextNoTxnKey key = 1
	probeCacheKey       key = 2
)

// AEContext retrieves the raw "google.golang.org/appengine" compatible Context.
//
// It also transfers deadline of `c` to AE context, since deadline is used for
// RPCs. Doesn't transfer cancelation ability though (since it's ignored by GAE
// anyway).
func AEContext(c context.Context) context.Context {
	aeCtx, _ := c.Value(prodContextKey).(context.Context)
	if aeCtx == nil {
		return nil
	}
	if deadline, ok := c.Deadline(); ok {
		aeCtx, _ = context.WithDeadline(aeCtx, deadline)
	}
	return aeCtx
}

// AEContextNoTxn retrieves the raw "google.golang.org/appengine" compatible
// Context that's not part of a transaction.
func AEContextNoTxn(c context.Context) context.Context {
	aeCtx, _ := c.Value(prodContextNoTxnKey).(context.Context)
	if aeCtx == nil {
		return nil
	}
	aeCtx, err := appengine.Namespace(aeCtx, info.Get(c).GetNamespace())
	if err != nil {
		panic(err)
	}
	if deadline, ok := c.Deadline(); ok {
		aeCtx, _ = context.WithDeadline(aeCtx, deadline)
	}
	return aeCtx
}

func setupAECtx(c, aeCtx context.Context) context.Context {
	c = context.WithValue(c, prodContextKey, aeCtx)
	c = context.WithValue(c, prodContextNoTxnKey, aeCtx)
	return useMail(useUser(useURLFetch(useRDS(useMC(useTQ(useGI(useLogging(c))))))))
}

// Use adds production implementations for all the gae services to the
// context.
//
// The services added are:
//   - github.com/luci-go/common/logging
//   - github.com/luci/gae/service/datastore
//   - github.com/luci/gae/service/info
//   - github.com/luci/gae/service/mail
//   - github.com/luci/gae/service/memcache
//   - github.com/luci/gae/service/taskqueue
//   - github.com/luci/gae/service/urlfetch
//   - github.com/luci/gae/service/user
//
// These can be retrieved with the <service>.Get functions.
//
// The implementations are all backed by the real appengine SDK functionality,
func Use(c context.Context, r *http.Request) context.Context {
	return setupAECtx(c, appengine.NewContext(r))
}

// UseRemote is the same as Use, except that it lets you attach a context to
// a remote host using the Remote API feature. See the docs for the
// prerequisites.
//
// docs: https://cloud.google.com/appengine/docs/go/tools/remoteapi
//
// inOutCtx will be replaced with the new, derived context, if err is nil,
// otherwise it's unchanged and continues to be safe-to-use.
//
// If client is nil, this will use create a new client, and will try to be
// clever about it:
//   * If you're creating a remote context FROM AppEngine, this will use
//     urlfetch.Transport. This can be used to allow app-to-app remote_api
//     control.
//
//   * If host starts with "localhost", this will create a regular http.Client
//     with a cookiejar, and call the _ah/login API to log in as an admin with
//     the user "admin@example.com".
//
//   * Otherwise, it will create a Google OAuth2 client with the following scopes:
//       - "https://www.googleapis.com/auth/appengine.apis"
//       - "https://www.googleapis.com/auth/userinfo.email"
//       - "https://www.googleapis.com/auth/cloud.platform"
func UseRemote(inOutCtx *context.Context, host string, client *http.Client) (err error) {
	if client == nil {
		if strings.HasPrefix(host, "localhost") {
			transp := http.DefaultTransport
			if aeCtx := AEContextNoTxn(*inOutCtx); aeCtx != nil {
				transp = urlfetch.Get(*inOutCtx)
			}

			client = &http.Client{Transport: transp}
			client.Jar, err = cookiejar.New(nil)
			if err != nil {
				return
			}
			u := fmt.Sprintf("http://%s/_ah/login?%s", host, url.Values{
				"email":  {"admin@example.com"},
				"admin":  {"True"},
				"action": {"Login"},
			}.Encode())

			var rsp *http.Response
			rsp, err = client.Get(u)
			if err != nil {
				return
			}
			defer rsp.Body.Close()
		} else {
			aeCtx := AEContextNoTxn(*inOutCtx)
			if aeCtx == nil {
				aeCtx = context.Background()
			}
			client, err = gOAuth.DefaultClient(aeCtx,
				"https://www.googleapis.com/auth/appengine.apis",
				"https://www.googleapis.com/auth/userinfo.email",
				"https://www.googleapis.com/auth/cloud.platform",
			)
			if err != nil {
				return
			}
		}
	}

	aeCtx, err := remote_api.NewRemoteContext(host, client)
	if err != nil {
		return
	}
	*inOutCtx = setupAECtx(*inOutCtx, aeCtx)
	return nil
}
<commit_msg>impl/prod: Expose Remote API scopes.<commit_after>// Copyright 2015 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package prod

import (
	"fmt"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"

	"github.com/luci/gae/service/info"
	"github.com/luci/gae/service/urlfetch"
	"golang.org/x/net/context"
	gOAuth "golang.org/x/oauth2/google"
	"google.golang.org/appengine"
	"google.golang.org/appengine/remote_api"
)

// RemoteAPIScopes is the set of OAuth2 scopes needed for Remote API access.
var RemoteAPIScopes = []string{
	"https://www.googleapis.com/auth/appengine.apis",
	"https://www.googleapis.com/auth/userinfo.email",
	"https://www.googleapis.com/auth/cloud.platform",
}

type key int

var (
	prodContextKey      key
	prodContextNoTxnKey key = 1
	probeCacheKey       key = 2
)

// AEContext retrieves the raw "google.golang.org/appengine" compatible Context.
//
// It also transfers deadline of `c` to AE context, since deadline is used for
// RPCs. Doesn't transfer cancelation ability though (since it's ignored by GAE
// anyway).
func AEContext(c context.Context) context.Context {
	aeCtx, _ := c.Value(prodContextKey).(context.Context)
	if aeCtx == nil {
		return nil
	}
	if deadline, ok := c.Deadline(); ok {
		aeCtx, _ = context.WithDeadline(aeCtx, deadline)
	}
	return aeCtx
}

// AEContextNoTxn retrieves the raw "google.golang.org/appengine" compatible
// Context that's not part of a transaction.
func AEContextNoTxn(c context.Context) context.Context {
	aeCtx, _ := c.Value(prodContextNoTxnKey).(context.Context)
	if aeCtx == nil {
		return nil
	}
	aeCtx, err := appengine.Namespace(aeCtx, info.Get(c).GetNamespace())
	if err != nil {
		panic(err)
	}
	if deadline, ok := c.Deadline(); ok {
		aeCtx, _ = context.WithDeadline(aeCtx, deadline)
	}
	return aeCtx
}

func setupAECtx(c, aeCtx context.Context) context.Context {
	c = context.WithValue(c, prodContextKey, aeCtx)
	c = context.WithValue(c, prodContextNoTxnKey, aeCtx)
	return useMail(useUser(useURLFetch(useRDS(useMC(useTQ(useGI(useLogging(c))))))))
}

// Use adds production implementations for all the gae services to the
// context.
//
// The services added are:
//   - github.com/luci-go/common/logging
//   - github.com/luci/gae/service/datastore
//   - github.com/luci/gae/service/info
//   - github.com/luci/gae/service/mail
//   - github.com/luci/gae/service/memcache
//   - github.com/luci/gae/service/taskqueue
//   - github.com/luci/gae/service/urlfetch
//   - github.com/luci/gae/service/user
//
// These can be retrieved with the <service>.Get functions.
//
// The implementations are all backed by the real appengine SDK functionality,
func Use(c context.Context, r *http.Request) context.Context {
	return setupAECtx(c, appengine.NewContext(r))
}

// UseRemote is the same as Use, except that it lets you attach a context to
// a remote host using the Remote API feature. See the docs for the
// prerequisites.
//
// docs: https://cloud.google.com/appengine/docs/go/tools/remoteapi
//
// inOutCtx will be replaced with the new, derived context, if err is nil,
// otherwise it's unchanged and continues to be safe-to-use.
//
// If client is nil, this will use create a new client, and will try to be
// clever about it:
//   * If you're creating a remote context FROM AppEngine, this will use
//     urlfetch.Transport. This can be used to allow app-to-app remote_api
//     control.
//
//   * If host starts with "localhost", this will create a regular http.Client
//     with a cookiejar, and call the _ah/login API to log in as an admin with
//     the user "admin@example.com".
//
//   * Otherwise, it will create a Google OAuth2 client with the following scopes:
//       - "https://www.googleapis.com/auth/appengine.apis"
//       - "https://www.googleapis.com/auth/userinfo.email"
//       - "https://www.googleapis.com/auth/cloud.platform"
func UseRemote(inOutCtx *context.Context, host string, client *http.Client) (err error) {
	if client == nil {
		if strings.HasPrefix(host, "localhost") {
			transp := http.DefaultTransport
			if aeCtx := AEContextNoTxn(*inOutCtx); aeCtx != nil {
				transp = urlfetch.Get(*inOutCtx)
			}

			client = &http.Client{Transport: transp}
			client.Jar, err = cookiejar.New(nil)
			if err != nil {
				return
			}
			u := fmt.Sprintf("http://%s/_ah/login?%s", host, url.Values{
				"email":  {"admin@example.com"},
				"admin":  {"True"},
				"action": {"Login"},
			}.Encode())

			var rsp *http.Response
			rsp, err = client.Get(u)
			if err != nil {
				return
			}
			defer rsp.Body.Close()
		} else {
			aeCtx := AEContextNoTxn(*inOutCtx)
			if aeCtx == nil {
				aeCtx = context.Background()
			}
			client, err = gOAuth.DefaultClient(aeCtx, RemoteAPIScopes...)
			if err != nil {
				return
			}
		}
	}

	aeCtx, err := remote_api.NewRemoteContext(host, client)
	if err != nil {
		return
	}
	*inOutCtx = setupAECtx(*inOutCtx, aeCtx)
	return nil
}
<|endoftext|>


<commit_before><commit_msg>smaller size<commit_after><|endoftext|>


<commit_before>package gin

import (
	"fmt"
	"os/exec"
	"runtime"
	"strings"
)

type Builder interface {
	Build() error
	Binary() string
	Errors() string
}

type builder struct {
	dir      string
	binary   string
	errors   string
	useGodep bool
}

func NewBuilder(dir string, bin string, useGodep bool) Builder {
	if len(bin) == 0 {
		bin = "bin"
	}

	// does not work on Windows without the ".exe" extension
	if runtime.GOOS == "windows" {
		if !strings.HasSuffix(bin, ".exe") { // check if it already has the .exe extension
			bin += ".exe"
		}
	}

	return &builder{dir: dir, binary: bin, useGodep: useGodep}
}

func (b *builder) Binary() string {
	return b.binary
}

func (b *builder) Errors() string {
	return b.errors
}

func (b *builder) Build() error {
	var command *exec.Cmd
	if b.useGodep {
		command = exec.Command("godep", "go", "build", "-o", b.binary)
	} else {
		command = exec.Command("go", "build", "-o", b.binary)
	}
	command.Dir = b.dir

	output, err := command.CombinedOutput()

	if command.ProcessState.Success() {
		b.errors = ""
	} else {
		b.errors = string(output)
	}

	if len(b.errors) > 0 {
		return fmt.Errorf(b.errors)
	}

	return err
}
<commit_msg>Add package name<commit_after>package gin

import (
	"fmt"
	"os/exec"
	"runtime"
	"strings"
)

type Builder interface {
	Build() error
	Binary() string
	Errors() string
}

type builder struct {
	dir      string
	binary   string
	errors   string
	useGodep bool
}

func NewBuilder(dir string, bin string, useGodep bool) Builder {
	if len(bin) == 0 {
		bin = "bin"
	}

	// does not work on Windows without the ".exe" extension
	if runtime.GOOS == "windows" {
		if !strings.HasSuffix(bin, ".exe") { // check if it already has the .exe extension
			bin += ".exe"
		}
	}

	return &builder{dir: dir, binary: bin, useGodep: useGodep}
}

func (b *builder) Binary() string {
	return b.binary
}

func (b *builder) Errors() string {
	return b.errors
}

func (b *builder) Build() error {
	var command *exec.Cmd
	if b.useGodep {
		command = exec.Command("godep", "go", "build", "-o", b.binary, "github.com/eave/eave-go")
	} else {
		command = exec.Command("go", "build", "-o", b.binary, "github.com/eave/eave-go")
	}
	command.Dir = b.dir

	output, err := command.CombinedOutput()

	if command.ProcessState.Success() {
		b.errors = ""
	} else {
		b.errors = string(output)
	}

	if len(b.errors) > 0 {
		return fmt.Errorf(b.errors)
	}

	return err
}
<|endoftext|>


<commit_before>/*
Library for parsing versions.yaml file.
*/
package versions

import (
	"fmt"
	"io/ioutil"
	"log"

	yaml "gopkg.in/yaml.v2"
)

type Package struct {
	Version string
	Major   string
	Gpg     string
	Sha1    string
	Sha256  string
	Md5     string
}

type Version struct {
	Dir          string
	Repo         string
	Tags         []string
	From         string
	Cmd          string
	Packages     map[string]Package
	ExcludeTests []string `yaml:"excludeTests"`
}

type Spec struct {
	Versions []Version
}

func LoadVersions(path string) Spec {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		log.Fatalf("error: %v", err)
	}

	spec := Spec{}
	err = yaml.Unmarshal([]byte(data), &spec)
	if err != nil {
		log.Fatalf("error: %v", err)
	}

	validateUniqueTags(spec)

	return spec
}

func validateUniqueTags(spec Spec) {
	repoTags := make(map[string]bool)
	for _, version := range spec.Versions {
		for _, tag := range version.Tags {
			repoTag := fmt.Sprintf("%s:%s", version.Repo, tag)
			if repoTags[repoTag] {
				log.Fatalf("error: duplicate repo tag %v", repoTag)
			}
			repoTags[repoTag] = true
		}
	}
}
<commit_msg>support for many custom commands addded to dockerfiles<commit_after>/*
Library for parsing versions.yaml file.
*/
package versions

import (
	"fmt"
	"io/ioutil"
	"log"

	yaml "gopkg.in/yaml.v2"
)

type Package struct {
	Version string
	Major   string
	Gpg     string
	Sha1    string
	Sha256  string
	Md5     string
}

// this will allow for passing many custom commands to Dockerfile.template
// usage in versions.yaml
//    custom:
//       your_command:
//                  cmd: "<command to be added to Dockerfile>"
// usage in Dockerfile.template
// {{- $cmd1 := index .Custom "your_command" }}
// {{ $cmd1.Cmd }} \
type Command struct {
	Cmd string
}

type Version struct {
	Dir          string
	Repo         string
	Tags         []string
	From         string
	Cmd          string
	Custom       map[string]Command
	Packages     map[string]Package
	ExcludeTests []string `yaml:"excludeTests"`
}

type Spec struct {
	Versions []Version
}

func LoadVersions(path string) Spec {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		log.Fatalf("error: %v", err)
	}

	spec := Spec{}
	err = yaml.Unmarshal([]byte(data), &spec)
	if err != nil {
		log.Fatalf("error: %v", err)
	}

	validateUniqueTags(spec)

	return spec
}

func validateUniqueTags(spec Spec) {
	repoTags := make(map[string]bool)
	for _, version := range spec.Versions {
		for _, tag := range version.Tags {
			repoTag := fmt.Sprintf("%s:%s", version.Repo, tag)
			if repoTags[repoTag] {
				log.Fatalf("error: duplicate repo tag %v", repoTag)
			}
			repoTags[repoTag] = true
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package fetch

import (
	"testing"

	"github.com/issue9/assert"
)

func TestMap(t *testing.T) {
	a := assert.New(t)
	db := initDB(a)
	defer closeDB(db, a)

	// 正常匹配数据，读取多行
	sql := `SELECT id,Email FROM user WHERE id<2 ORDER BY id`
	rows, err := db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err := Map(false, rows)
	a.NotError(err).NotNil(mapped)

	a.Equal([]map[string]interface{}{
		map[string]interface{}{"id": 0, "Email": "email-0"},
		map[string]interface{}{"id": 1, "Email": "email-1"},
	}, mapped)
	a.NotError(rows.Close())

	// 正常匹配数据，读取一行
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = Map(true, rows)
	a.NotError(err).NotNil(mapped)

	a.Equal([]map[string]interface{}{
		map[string]interface{}{"id": 0, "Email": "email-0"},
	}, mapped)
	a.NotError(rows.Close())

	// 没有匹配的数据，读取多行
	sql = `SELECT id,Email FROM user WHERE id<0 ORDER BY id`
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = Map(false, rows)
	a.NotError(err)

	a.Equal([]map[string]interface{}{}, mapped)
	a.NotError(rows.Close())

	// 没有匹配的数据，读取一行
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = Map(true, rows)
	a.NotError(err)

	a.Equal([]map[string]interface{}{}, mapped)
	a.NotError(rows.Close())
}

func TestMapString(t *testing.T) {
	a := assert.New(t)
	db := initDB(a)
	defer closeDB(db, a)

	// 正常数据匹配，读取多行
	sql := `SELECT id,Email FROM user WHERE id<2 ORDER BY id`
	rows, err := db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err := MapString(false, rows)
	a.NotError(err).NotNil(mapped)

	a.Equal(mapped, []map[string]string{
		map[string]string{"id": "0", "Email": "email-0"},
		map[string]string{"id": "1", "Email": "email-1"},
	})
	a.NotError(rows.Close())

	// 正常数据匹配，读取一行
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = MapString(true, rows)
	a.NotError(err).NotNil(mapped)

	a.Equal(mapped, []map[string]string{
		map[string]string{"id": "0", "Email": "email-0"},
	})
	a.NotError(rows.Close())

	// 没有数据匹配，读取多行
	sql = `SELECT id,Email FROM user WHERE id<0 ORDER BY id`
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = MapString(false, rows)
	a.NotError(err)

	a.Equal(mapped, []map[string]string{})
	a.NotError(rows.Close())

	// 没有数据匹配，读取一行
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = MapString(true, rows)
	a.NotError(err)

	a.Equal(mapped, []map[string]string{})
	a.NotError(rows.Close())
}
<commit_msg>添加性能测试<commit_after>// Copyright 2014 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package fetch

import (
	"testing"

	"github.com/issue9/assert"
)

func BenchmarkMap(b *testing.B) {
	a := assert.New(b)
	db := initDB(a)
	defer closeDB(db, a)

	// 正常匹配数据，读取多行
	sql := `SELECT id,Email FROM user WHERE id<2 ORDER BY id`

	for i := 0; i < b.N; i++ {
		rows, err := db.Query(sql)
		a.NotError(err)

		mapped, err := Map(false, rows)
		a.NotError(err).NotNil(mapped)
		rows.Close()
	}
}

func TestMap(t *testing.T) {
	a := assert.New(t)
	db := initDB(a)
	defer closeDB(db, a)

	// 正常匹配数据，读取多行
	sql := `SELECT id,Email FROM user WHERE id<2 ORDER BY id`
	rows, err := db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err := Map(false, rows)
	a.NotError(err).NotNil(mapped)

	a.Equal([]map[string]interface{}{
		map[string]interface{}{"id": 0, "Email": "email-0"},
		map[string]interface{}{"id": 1, "Email": "email-1"},
	}, mapped)
	a.NotError(rows.Close())

	// 正常匹配数据，读取一行
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = Map(true, rows)
	a.NotError(err).NotNil(mapped)

	a.Equal([]map[string]interface{}{
		map[string]interface{}{"id": 0, "Email": "email-0"},
	}, mapped)
	a.NotError(rows.Close())

	// 没有匹配的数据，读取多行
	sql = `SELECT id,Email FROM user WHERE id<0 ORDER BY id`
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = Map(false, rows)
	a.NotError(err)

	a.Equal([]map[string]interface{}{}, mapped)
	a.NotError(rows.Close())

	// 没有匹配的数据，读取一行
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = Map(true, rows)
	a.NotError(err)

	a.Equal([]map[string]interface{}{}, mapped)
	a.NotError(rows.Close())
}

func TestMapString(t *testing.T) {
	a := assert.New(t)
	db := initDB(a)
	defer closeDB(db, a)

	// 正常数据匹配，读取多行
	sql := `SELECT id,Email FROM user WHERE id<2 ORDER BY id`
	rows, err := db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err := MapString(false, rows)
	a.NotError(err).NotNil(mapped)

	a.Equal(mapped, []map[string]string{
		map[string]string{"id": "0", "Email": "email-0"},
		map[string]string{"id": "1", "Email": "email-1"},
	})
	a.NotError(rows.Close())

	// 正常数据匹配，读取一行
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = MapString(true, rows)
	a.NotError(err).NotNil(mapped)

	a.Equal(mapped, []map[string]string{
		map[string]string{"id": "0", "Email": "email-0"},
	})
	a.NotError(rows.Close())

	// 没有数据匹配，读取多行
	sql = `SELECT id,Email FROM user WHERE id<0 ORDER BY id`
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = MapString(false, rows)
	a.NotError(err)

	a.Equal(mapped, []map[string]string{})
	a.NotError(rows.Close())

	// 没有数据匹配，读取一行
	rows, err = db.Query(sql)
	a.NotError(err).NotNil(rows)

	mapped, err = MapString(true, rows)
	a.NotError(err)

	a.Equal(mapped, []map[string]string{})
	a.NotError(rows.Close())
}
<|endoftext|>


<commit_before>package file

import (
	"errors"
	"go/format"
	"regexp"
	"strings"
)

func (file *File) replaceBuffer(newBuffer Buffer) {
	n := len(newBuffer)
	file.buffer = file.buffer[:n]
	for k, line := range newBuffer {
		if k > len(file.buffer) {
			file.buffer = append(file.buffer, line)
		} else {
			if file.buffer[k].ToString() != line.ToString() {
				file.buffer[k] = line
			}
		}
	}
}

func (file *File) GoFmt() error {
	filetype := file.SyntaxRules.GetFileType(file.Name)
	if filetype != "go" {
		return errors.New("Will not gofmt a non-go file.")
	}
	contents := file.toString()
	bytes, err := format.Source([]byte(contents))
	if err == nil {
		stringBuf := strings.Split(string(bytes), file.newline)
		newBuffer := MakeBuffer(stringBuf)
		file.replaceBuffer(newBuffer)
	}
	file.Snapshot()
	return nil
}

func (file *File) InsertChar(ch rune) {
	maxCol := 0
	maxLineLen := 0
	for _, cursor := range file.MultiCursor {
		if cursor.col > maxCol {
			maxCol = cursor.col
		}
		if len(file.buffer[cursor.row]) > maxLineLen {
			maxLineLen = len(file.buffer[cursor.row])
		}
	}
	for idx, cursor := range file.MultiCursor {
		col, row := cursor.col, cursor.row
		if maxCol > 0 && col == 0 {
			continue
		}
		line := file.buffer[row]
		if (ch == ' ' || ch == '\t') && col == 0 && len(line) == 0 && maxLineLen > 0 {
			continue
		}
		insertStr := string(ch)
		if ch == '\t' && file.autoTab && file.tabString != "\t" {
			insertStr = file.tabString
		}
		file.buffer[row] = Line(string(line[0:col]) + insertStr + string(line[col:]))
		file.MultiCursor[idx].col += len(insertStr)
		file.MultiCursor[idx].colwant = file.MultiCursor[idx].col
	}
	file.Snapshot()
}

func (file *File) Backspace() {
	for idx, cursor := range file.MultiCursor {
		col, row := cursor.col, cursor.row
		if col == 0 {
			if len(file.MultiCursor) > 1 {
				continue
			}
			if row == 0 {
				return
			}
			row -= 1
			if row+1 >= len(file.buffer) {
				return
			}
			col = len(file.buffer[row])
			file.buffer[row] = append(file.buffer[row], file.buffer[row+1]...)
			file.buffer = append(file.buffer[0:row+1], file.buffer[row+2:]...)
			file.MultiCursor[idx].col = col
			file.MultiCursor[idx].row = row
		} else {
			line := file.buffer[row]
			if col > len(line) {
				continue
			}

			// Handle multi-char indents.
			nDel := 1
			if file.autoTab && len(file.tabString) > 0 {
				if string(line[0:col]) == strings.Repeat(" ", col) {
					n := len(file.tabString)
					if n*(col/n) == col {
						nDel = n
					}
				}
			}

			file.buffer[row] = Line(string(line[0:col-nDel]) + string(line[col:]))
			file.MultiCursor[idx].col = col - nDel
			file.MultiCursor[idx].row = row
		}
		file.MultiCursor[idx].colwant = file.MultiCursor[idx].col
	}
	file.EnforceRowBounds()
	file.EnforceColBounds()
	file.Snapshot()
}

func (file *File) Delete() {
	file.CursorRight()
	file.Backspace()
}

func (file *File) Newline() {

	rate := file.timer.Tick()

	for idx, cursor := range file.MultiCursor {

		col, row := cursor.col, cursor.row
		lineStart := file.buffer[row][0:col]
		lineEnd := file.buffer[row][col:]

		file.buffer = append(file.buffer, Line(""))
		copy(file.buffer[row+2:], file.buffer[row+1:])
		file.buffer[row+1] = lineEnd

		file.MultiCursor[idx].row = row + 1
		file.MultiCursor[idx].col = 0

		if file.autoIndent && rate < file.maxRate && len(lineEnd) == 0 {
			file.DoAutoIndent(idx)
		}

		file.buffer[row] = lineStart.RemoveTrailingWhitespace()

	}

	file.Snapshot()
}

func (file *File) DoAutoIndent(cursorIdx int) {

	row := file.MultiCursor[cursorIdx].row
	if row == 0 {
		return
	}

	origLine := file.buffer[row].Dup()

	// Whitespace-only indent.
	re, _ := regexp.Compile("^[ \t]+")
	ws := Line(re.FindString(file.buffer[row-1].ToString()))
	if len(ws) > 0 {
		file.buffer[row] = append(ws, file.buffer[row]...)
		file.MultiCursor[cursorIdx].col += len(ws)
		if len(file.buffer[row-1]) == len(ws) {
			file.buffer[row-1] = Line("")
		}
	}

	if row < 2 {
		return
	}

	// Non-whitespace indent.
	indent := file.buffer[row-1].CommonStart(file.buffer[row-2])
	if len(indent) > len(ws) {
		file.ForceSnapshot()
		file.buffer[row] = append(indent, origLine...)
		file.MultiCursor[cursorIdx].col += len(indent) - len(ws)
	}

}

func (file *File) Justify(lineLen int) {
	minRow, maxRow := file.MultiCursor.MinMaxRow()
	lines := file.buffer[minRow : maxRow+1]
	bigString := lines.ToString(" ")
	lines = MakeSplitBuffer(bigString, lineLen)
	file.buffer = file.buffer.ReplaceLines(lines, minRow, maxRow)
	file.MultiCursor = file.MultiCursor.Clear()
	file.Snapshot()
}

func (file *File) Cut() Buffer {
	row := file.MultiCursor[0].row
	cutBuffer := file.buffer[row : row+1].Dup()
	if len(file.buffer) == 1 {
		file.buffer = MakeBuffer([]string{""})
	} else if row == 0 {
		file.buffer = file.buffer[1:]
	} else if row < len(file.buffer)-1 {
		file.buffer = append(file.buffer[:row], file.buffer[row+1:]...)
	} else {
		file.buffer = file.buffer[:row]
	}
	file.EnforceRowBounds()
	file.EnforceColBounds()
	file.Snapshot()
	return cutBuffer
}

func (file *File) Paste(buffer Buffer) {
	row := file.MultiCursor[0].row
	newBuffer := file.buffer[:row].Dup()
	for _, line := range buffer {
		newBuffer = append(newBuffer, line.Dup())
	}
	file.buffer = append(newBuffer, file.buffer[row:].Dup()...)
	file.CursorDown(len(buffer))
	file.EnforceRowBounds()
	file.EnforceColBounds()
	file.Snapshot()
}

func (file *File) CutToStartOfLine() {
	allAtZero := true
	for _, cursor := range file.MultiCursor {
		if cursor.col != 0 {
			allAtZero = false
			break
		}
	}
	if allAtZero {
		re := regexp.MustCompile("^[ \t]*")
		for idx, cursor := range file.MultiCursor {
			row := cursor.row
			line := file.buffer[row]
			match := re.FindStringIndex(line.ToString())
			file.MultiCursor[idx].col = match[1]
			file.MultiCursor[idx].colwant = file.MultiCursor[idx].col
		}
	} else {
		for idx, _ := range file.MultiCursor {
			file.MultiCursor[idx].col = 0
			file.MultiCursor[idx].colwant = file.MultiCursor[idx].col
		}
	}
}

func (file *File) CutToEndOfLine() {
	for idx, _ := range file.MultiCursor {
		row := file.MultiCursor[idx].row
		file.MultiCursor[idx].col = len(file.buffer[row])
		file.MultiCursor[idx].colwant = file.MultiCursor[idx].col
	}
}
<commit_msg>cut-to-end and cut-to-start of line<commit_after>package file

import (
	"errors"
	"go/format"
	"regexp"
	"strings"
)

func (file *File) replaceBuffer(newBuffer Buffer) {
	n := len(newBuffer)
	file.buffer = file.buffer[:n]
	for k, line := range newBuffer {
		if k > len(file.buffer) {
			file.buffer = append(file.buffer, line)
		} else {
			if file.buffer[k].ToString() != line.ToString() {
				file.buffer[k] = line
			}
		}
	}
}

func (file *File) GoFmt() error {
	filetype := file.SyntaxRules.GetFileType(file.Name)
	if filetype != "go" {
		return errors.New("Will not gofmt a non-go file.")
	}
	contents := file.toString()
	bytes, err := format.Source([]byte(contents))
	if err == nil {
		stringBuf := strings.Split(string(bytes), file.newline)
		newBuffer := MakeBuffer(stringBuf)
		file.replaceBuffer(newBuffer)
	}
	file.Snapshot()
	return nil
}

func (file *File) InsertChar(ch rune) {
	maxCol := 0
	maxLineLen := 0
	for _, cursor := range file.MultiCursor {
		if cursor.col > maxCol {
			maxCol = cursor.col
		}
		if len(file.buffer[cursor.row]) > maxLineLen {
			maxLineLen = len(file.buffer[cursor.row])
		}
	}
	for idx, cursor := range file.MultiCursor {
		col, row := cursor.col, cursor.row
		if maxCol > 0 && col == 0 {
			continue
		}
		line := file.buffer[row]
		if (ch == ' ' || ch == '\t') && col == 0 && len(line) == 0 && maxLineLen > 0 {
			continue
		}
		insertStr := string(ch)
		if ch == '\t' && file.autoTab && file.tabString != "\t" {
			insertStr = file.tabString
		}
		file.buffer[row] = Line(string(line[0:col]) + insertStr + string(line[col:]))
		file.MultiCursor[idx].col += len(insertStr)
		file.MultiCursor[idx].colwant = file.MultiCursor[idx].col
	}
	file.Snapshot()
}

func (file *File) Backspace() {
	for idx, cursor := range file.MultiCursor {
		col, row := cursor.col, cursor.row
		if col == 0 {
			if len(file.MultiCursor) > 1 {
				continue
			}
			if row == 0 {
				return
			}
			row -= 1
			if row+1 >= len(file.buffer) {
				return
			}
			col = len(file.buffer[row])
			file.buffer[row] = append(file.buffer[row], file.buffer[row+1]...)
			file.buffer = append(file.buffer[0:row+1], file.buffer[row+2:]...)
			file.MultiCursor[idx].col = col
			file.MultiCursor[idx].row = row
		} else {
			line := file.buffer[row]
			if col > len(line) {
				continue
			}

			// Handle multi-char indents.
			nDel := 1
			if file.autoTab && len(file.tabString) > 0 {
				if string(line[0:col]) == strings.Repeat(" ", col) {
					n := len(file.tabString)
					if n*(col/n) == col {
						nDel = n
					}
				}
			}

			file.buffer[row] = Line(string(line[0:col-nDel]) + string(line[col:]))
			file.MultiCursor[idx].col = col - nDel
			file.MultiCursor[idx].row = row
		}
		file.MultiCursor[idx].colwant = file.MultiCursor[idx].col
	}
	file.EnforceRowBounds()
	file.EnforceColBounds()
	file.Snapshot()
}

func (file *File) Delete() {
	file.CursorRight()
	file.Backspace()
}

func (file *File) Newline() {

	rate := file.timer.Tick()

	for idx, cursor := range file.MultiCursor {

		col, row := cursor.col, cursor.row
		lineStart := file.buffer[row][0:col]
		lineEnd := file.buffer[row][col:]

		file.buffer = append(file.buffer, Line(""))
		copy(file.buffer[row+2:], file.buffer[row+1:])
		file.buffer[row+1] = lineEnd

		file.MultiCursor[idx].row = row + 1
		file.MultiCursor[idx].col = 0

		if file.autoIndent && rate < file.maxRate && len(lineEnd) == 0 {
			file.DoAutoIndent(idx)
		}

		file.buffer[row] = lineStart.RemoveTrailingWhitespace()

	}

	file.Snapshot()
}

func (file *File) DoAutoIndent(cursorIdx int) {

	row := file.MultiCursor[cursorIdx].row
	if row == 0 {
		return
	}

	origLine := file.buffer[row].Dup()

	// Whitespace-only indent.
	re, _ := regexp.Compile("^[ \t]+")
	ws := Line(re.FindString(file.buffer[row-1].ToString()))
	if len(ws) > 0 {
		file.buffer[row] = append(ws, file.buffer[row]...)
		file.MultiCursor[cursorIdx].col += len(ws)
		if len(file.buffer[row-1]) == len(ws) {
			file.buffer[row-1] = Line("")
		}
	}

	if row < 2 {
		return
	}

	// Non-whitespace indent.
	indent := file.buffer[row-1].CommonStart(file.buffer[row-2])
	if len(indent) > len(ws) {
		file.ForceSnapshot()
		file.buffer[row] = append(indent, origLine...)
		file.MultiCursor[cursorIdx].col += len(indent) - len(ws)
	}

}

func (file *File) Justify(lineLen int) {
	minRow, maxRow := file.MultiCursor.MinMaxRow()
	lines := file.buffer[minRow : maxRow+1]
	bigString := lines.ToString(" ")
	lines = MakeSplitBuffer(bigString, lineLen)
	file.buffer = file.buffer.ReplaceLines(lines, minRow, maxRow)
	file.MultiCursor = file.MultiCursor.Clear()
	file.Snapshot()
}

func (file *File) Cut() Buffer {
	row := file.MultiCursor[0].row
	cutBuffer := file.buffer[row : row+1].Dup()
	if len(file.buffer) == 1 {
		file.buffer = MakeBuffer([]string{""})
	} else if row == 0 {
		file.buffer = file.buffer[1:]
	} else if row < len(file.buffer)-1 {
		file.buffer = append(file.buffer[:row], file.buffer[row+1:]...)
	} else {
		file.buffer = file.buffer[:row]
	}
	file.EnforceRowBounds()
	file.EnforceColBounds()
	file.Snapshot()
	return cutBuffer
}

func (file *File) Paste(buffer Buffer) {
	row := file.MultiCursor[0].row
	newBuffer := file.buffer[:row].Dup()
	for _, line := range buffer {
		newBuffer = append(newBuffer, line.Dup())
	}
	file.buffer = append(newBuffer, file.buffer[row:].Dup()...)
	file.CursorDown(len(buffer))
	file.EnforceRowBounds()
	file.EnforceColBounds()
	file.Snapshot()
}

func (file *File) CutToStartOfLine() {
	for idx, _ := range file.MultiCursor {
		row := file.MultiCursor[idx].row
		col := file.MultiCursor[idx].col
		file.buffer[row] = file.buffer[row][col:]
		file.MultiCursor[idx].col = 0
	}
	file.Snapshot()
}

func (file *File) CutToEndOfLine() {
	for idx, _ := range file.MultiCursor {
		row := file.MultiCursor[idx].row
		col := file.MultiCursor[idx].col
		file.buffer[row] = file.buffer[row][:col]
	}
	file.Snapshot()
}
<|endoftext|>


<commit_before>package filecache_test

import (
	"fmt"
	"github.com/DannyBen/filecache"
)

func Example() {
	// Get a handler and set a directory + 1 hour cache life
	handler := filecache.Handler{"./cache", 1}

	// Data to store in cache
	data := []byte("Joey doesn't share food")

	// Store the data in the cache
	handler.Set("testkey", data)

	// Retrieve the data from the cache
	r := handler.Get("testkey")

	// Show the result
	fmt.Println(string(r))

	// Show the filename
	fmt.Println(handler.Filename("testkey"))

	// Wait for some seconds
	// time.Sleep(7 * time.Second)

	// By now the cache is invalid
	// r = handler.Get("testkey")
	// if r == nil {
	// 	fmt.Println("Cache expired")
	// }

	// Output:
	// Joey doesn't share food
	// ./cache/221b368d7f5f597867f525971f28ff75
}
<commit_msg>Changed example to use temp folder, to work with Travis<commit_after>package filecache_test

import (
	"fmt"
	"github.com/DannyBen/filecache"
)

func Example() {
	// Get a cache handler with default (temp) directory as
	// the cache directory, and a one hour cache life
	handler := filecache.Handler{Life: 1}

	// Data to store in cache
	data := []byte("Joey doesn't share food")

	// Store the data in the cache
	handler.Set("testkey", data)

	// Retrieve the data from the cache
	r := handler.Get("testkey")

	// Show the result
	fmt.Println(string(r))

	// Output:
	// Joey doesn't share food
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"runtime"
	d "runtime/debug"
	"strconv"
	"strings"
	"time"

	bimg "gopkg.in/h2non/bimg.v1"

	. "github.com/tj/go-debug"
)

var debug = Debug("imaginary")

var (
	aAddr              = flag.String("a", "", "Bind address")
	aPort              = flag.Int("p", 8088, "Port to listen")
	aVers              = flag.Bool("v", false, "Show version")
	aVersl             = flag.Bool("version", false, "Show version")
	aHelp              = flag.Bool("h", false, "Show help")
	aHelpl             = flag.Bool("help", false, "Show help")
	aPathPrefix        = flag.String("path-prefix", "/", "Url path prefix to listen to")
	aCors              = flag.Bool("cors", false, "Enable CORS support")
	aGzip              = flag.Bool("gzip", false, "Enable gzip compression (deprecated)")
	aAuthForwarding    = flag.Bool("enable-auth-forwarding", false, "Forwards X-Forward-Authorization or Authorization header to the image source server. -enable-url-source flag must be defined. Tip: secure your server from public access to prevent attack vectors")
	aEnableURLSource   = flag.Bool("enable-url-source", false, "Enable remote HTTP URL image source processing")
	aEnablePlaceholder = flag.Bool("enable-placeholder", false, "Enable image response placeholder to be used in case of error")
	aAlloweOrigins     = flag.String("allowed-origins", "", "Restrict remote image source processing to certain origins (separated by commas)")
	aMaxAllowedSize    = flag.Int("max-allowed-size", 0, "Restrict maximum size of http image source (in bytes)")
	aKey               = flag.String("key", "", "Define API key for authorization")
	aMount             = flag.String("mount", "", "Mount server local directory")
	aCertFile          = flag.String("certfile", "", "TLS certificate file path")
	aKeyFile           = flag.String("keyfile", "", "TLS private key file path")
	aAuthorization     = flag.String("authorization", "", "Defines a constant Authorization header value passed to all the image source servers. -enable-url-source flag must be defined. This overwrites authorization headers forwarding behavior via X-Forward-Authorization")
	aPlaceholder       = flag.String("placeholder", "", "Image path to image custom placeholder to be used in case of error. Recommended minimum image size is: 1200x1200")
	aDisableEndpoints  = flag.String("disable-endpoints", "", "Comma separated endpoints to disable. E.g: form,crop,rotate,health")
	aHTTPCacheTTL      = flag.Int("http-cache-ttl", -1, "The TTL in seconds")
	aReadTimeout       = flag.Int("http-read-timeout", 60, "HTTP read timeout in seconds")
	aWriteTimeout      = flag.Int("http-write-timeout", 60, "HTTP write timeout in seconds")
	aConcurrency       = flag.Int("concurrency", 0, "Throttle concurrency limit per second")
	aBurst             = flag.Int("burst", 100, "Throttle burst max cache size")
	aMRelease          = flag.Int("mrelease", 30, "OS memory release interval in seconds")
	aCpus              = flag.Int("cpus", runtime.GOMAXPROCS(-1), "Number of cpu cores to use")
)

const usage = `imaginary %s

Usage:
  imaginary -p 80
  imaginary -cors
  imaginary -concurrency 10
  imaginary -path-prefix /api/v1
  imaginary -enable-url-source
  imaginary -disable-endpoints form,health,crop,rotate
  imaginary -enable-url-source -allowed-origins http://localhost,http://server.com
  imaginary -enable-url-source -enable-auth-forwarding
  imaginary -enable-url-source -authorization "Basic AwDJdL2DbwrD=="
  imaginary -enable-placeholder
  imaginary -enable-url-source -placeholder ./placeholder.jpg
  imaginary -h | -help
  imaginary -v | -version

Options:
  -a <addr>                 Bind address [default: *]
  -p <port>                 Bind port [default: 8088]
  -h, -help                 Show help
  -v, -version              Show version
  -path-prefix <value>      Url path prefix to listen to [default: "/"]
  -cors                     Enable CORS support [default: false]
  -gzip                     Enable gzip compression (deprecated) [default: false]
  -disable-endpoints        Comma separated endpoints to disable. E.g: form,crop,rotate,health [default: ""]
  -key <key>                Define API key for authorization
  -mount <path>             Mount server local directory
  -http-cache-ttl <num>     The TTL in seconds. Adds caching headers to locally served files.
  -http-read-timeout <num>  HTTP read timeout in seconds [default: 30]
  -http-write-timeout <num> HTTP write timeout in seconds [default: 30]
  -enable-url-source        Restrict remote image source processing to certain origins (separated by commas)
  -enable-placeholder       Enable image response placeholder to be used in case of error [default: false]
  -enable-auth-forwarding   Forwards X-Forward-Authorization or Authorization header to the image source server. -enable-url-source flag must be defined. Tip: secure your server from public access to prevent attack vectors
  -allowed-origins <urls>   Restrict remote image source processing to certain origins (separated by commas)
  -max-allowed-size <bytes> Restrict maximum size of http image source (in bytes)
  -certfile <path>          TLS certificate file path
  -keyfile <path>           TLS private key file path
  -authorization <value>    Defines a constant Authorization header value passed to all the image source servers. -enable-url-source flag must be defined. This overwrites authorization headers forwarding behavior via X-Forward-Authorization
  -placeholder <path>       Image path to image custom placeholder to be used in case of error. Recommended minimum image size is: 1200x1200
  -concurrency <num>        Throttle concurrency limit per second [default: disabled]
  -burst <num>              Throttle burst max cache size [default: 100]
  -mrelease <num>           OS memory release interval in seconds [default: 30]
  -cpus <num>               Number of used cpu cores.
                            (default for current machine is %d cores)
`

func main() {
	flag.Usage = func() {
		fmt.Fprint(os.Stderr, fmt.Sprintf(usage, Version, runtime.NumCPU()))
	}
	flag.Parse()

	if *aHelp || *aHelpl {
		showUsage()
	}
	if *aVers || *aVersl {
		showVersion()
	}

	// Only required in Go < 1.5
	runtime.GOMAXPROCS(*aCpus)

	port := getPort(*aPort)
	opts := ServerOptions{
		Port:              port,
		Address:           *aAddr,
		CORS:              *aCors,
		AuthForwarding:    *aAuthForwarding,
		EnableURLSource:   *aEnableURLSource,
		EnablePlaceholder: *aEnablePlaceholder,
		PathPrefix:        *aPathPrefix,
		APIKey:            *aKey,
		Concurrency:       *aConcurrency,
		Burst:             *aBurst,
		Mount:             *aMount,
		CertFile:          *aCertFile,
		KeyFile:           *aKeyFile,
		Placeholder:       *aPlaceholder,
		HTTPCacheTTL:      *aHTTPCacheTTL,
		HTTPReadTimeout:   *aReadTimeout,
		HTTPWriteTimeout:  *aWriteTimeout,
		Authorization:     *aAuthorization,
		AlloweOrigins:     parseOrigins(*aAlloweOrigins),
		MaxAllowedSize:    *aMaxAllowedSize,
	}

	// Show warning if gzip flag is passed
	if *aGzip {
		fmt.Println("warning: -gzip flag is deprecated and will not have effect")
	}

	// Create a memory release goroutine
	if *aMRelease > 0 {
		memoryRelease(*aMRelease)
	}

	// Check if the mount directory exists, if present
	if *aMount != "" {
		checkMountDirectory(*aMount)
	}

	// Validate HTTP cache param, if present
	if *aHTTPCacheTTL != -1 {
		checkHttpCacheTtl(*aHTTPCacheTTL)
	}

	// Parse endpoint names to disabled, if present
	if *aDisableEndpoints != "" {
		opts.Endpoints = parseEndpoints(*aDisableEndpoints)
	}

	// Read placeholder image, if required
	if *aPlaceholder != "" {
		buf, err := ioutil.ReadFile(*aPlaceholder)
		if err != nil {
			exitWithError("cannot start the server: %s", err)
		}

		imageType := bimg.DetermineImageType(buf)
		if !bimg.IsImageTypeSupportedByVips(imageType).Load {
			exitWithError("Placeholder image type is not supported. Only JPEG, PNG or WEBP are supported")
		}

		opts.PlaceholderImage = buf
	} else if *aEnablePlaceholder {
		// Expose default placeholder
		opts.PlaceholderImage = placeholder
	}

	debug("imaginary server listening on port :%d/%s", opts.Port, strings.TrimPrefix(opts.PathPrefix, "/"))

	// Load image source providers
	LoadSources(opts)

	// Start the server
	err := Server(opts)
	if err != nil {
		exitWithError("cannot start the server: %s", err)
	}
}

func getPort(port int) int {
	if portEnv := os.Getenv("PORT"); portEnv != "" {
		newPort, _ := strconv.Atoi(portEnv)
		if newPort > 0 {
			port = newPort
		}
	}
	return port
}

func showUsage() {
	flag.Usage()
	os.Exit(1)
}

func showVersion() {
	fmt.Println(Version)
	os.Exit(1)
}

func checkMountDirectory(path string) {
	src, err := os.Stat(path)
	if err != nil {
		exitWithError("error while mounting directory: %s", err)
	}
	if src.IsDir() == false {
		exitWithError("mount path is not a directory: %s", path)
	}
	if path == "/" {
		exitWithError("cannot mount root directory for security reasons")
	}
}

func checkHttpCacheTtl(ttl int) {
	if ttl < -1 || ttl > 31556926 {
		exitWithError("The -http-cache-ttl flag only accepts a value from 0 to 31556926")
	}

	if ttl == 0 {
		debug("Adding HTTP cache control headers set to prevent caching.")
	}
}

func parseOrigins(origins string) []*url.URL {
	urls := []*url.URL{}
	if origins == "" {
		return urls
	}
	for _, origin := range strings.Split(origins, ",") {
		u, err := url.Parse(origin)
		if err != nil {
			continue
		}
		urls = append(urls, u)
	}
	return urls
}

func parseEndpoints(input string) Endpoints {
	endpoints := Endpoints{}
	for _, endpoint := range strings.Split(input, ",") {
		endpoint = strings.ToLower(strings.TrimSpace(endpoint))
		if endpoint != "" {
			endpoints = append(endpoints, endpoint)
		}
	}
	return endpoints
}

func memoryRelease(interval int) {
	ticker := time.NewTicker(time.Duration(interval) * time.Second)
	go func() {
		for range ticker.C {
			debug("FreeOSMemory()")
			d.FreeOSMemory()
		}
	}()
}

func exitWithError(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args)
	os.Exit(1)
}
<commit_msg>fix(debug): implement custom debug function<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"runtime"
	d "runtime/debug"
	"strconv"
	"strings"
	"time"

	bimg "gopkg.in/h2non/bimg.v1"
)

var (
	aAddr              = flag.String("a", "", "Bind address")
	aPort              = flag.Int("p", 8088, "Port to listen")
	aVers              = flag.Bool("v", false, "Show version")
	aVersl             = flag.Bool("version", false, "Show version")
	aHelp              = flag.Bool("h", false, "Show help")
	aHelpl             = flag.Bool("help", false, "Show help")
	aPathPrefix        = flag.String("path-prefix", "/", "Url path prefix to listen to")
	aCors              = flag.Bool("cors", false, "Enable CORS support")
	aGzip              = flag.Bool("gzip", false, "Enable gzip compression (deprecated)")
	aAuthForwarding    = flag.Bool("enable-auth-forwarding", false, "Forwards X-Forward-Authorization or Authorization header to the image source server. -enable-url-source flag must be defined. Tip: secure your server from public access to prevent attack vectors")
	aEnableURLSource   = flag.Bool("enable-url-source", false, "Enable remote HTTP URL image source processing")
	aEnablePlaceholder = flag.Bool("enable-placeholder", false, "Enable image response placeholder to be used in case of error")
	aAlloweOrigins     = flag.String("allowed-origins", "", "Restrict remote image source processing to certain origins (separated by commas)")
	aMaxAllowedSize    = flag.Int("max-allowed-size", 0, "Restrict maximum size of http image source (in bytes)")
	aKey               = flag.String("key", "", "Define API key for authorization")
	aMount             = flag.String("mount", "", "Mount server local directory")
	aCertFile          = flag.String("certfile", "", "TLS certificate file path")
	aKeyFile           = flag.String("keyfile", "", "TLS private key file path")
	aAuthorization     = flag.String("authorization", "", "Defines a constant Authorization header value passed to all the image source servers. -enable-url-source flag must be defined. This overwrites authorization headers forwarding behavior via X-Forward-Authorization")
	aPlaceholder       = flag.String("placeholder", "", "Image path to image custom placeholder to be used in case of error. Recommended minimum image size is: 1200x1200")
	aDisableEndpoints  = flag.String("disable-endpoints", "", "Comma separated endpoints to disable. E.g: form,crop,rotate,health")
	aHTTPCacheTTL      = flag.Int("http-cache-ttl", -1, "The TTL in seconds")
	aReadTimeout       = flag.Int("http-read-timeout", 60, "HTTP read timeout in seconds")
	aWriteTimeout      = flag.Int("http-write-timeout", 60, "HTTP write timeout in seconds")
	aConcurrency       = flag.Int("concurrency", 0, "Throttle concurrency limit per second")
	aBurst             = flag.Int("burst", 100, "Throttle burst max cache size")
	aMRelease          = flag.Int("mrelease", 30, "OS memory release interval in seconds")
	aCpus              = flag.Int("cpus", runtime.GOMAXPROCS(-1), "Number of cpu cores to use")
)

const usage = `imaginary %s

Usage:
  imaginary -p 80
  imaginary -cors
  imaginary -concurrency 10
  imaginary -path-prefix /api/v1
  imaginary -enable-url-source
  imaginary -disable-endpoints form,health,crop,rotate
  imaginary -enable-url-source -allowed-origins http://localhost,http://server.com
  imaginary -enable-url-source -enable-auth-forwarding
  imaginary -enable-url-source -authorization "Basic AwDJdL2DbwrD=="
  imaginary -enable-placeholder
  imaginary -enable-url-source -placeholder ./placeholder.jpg
  imaginary -h | -help
  imaginary -v | -version

Options:
  -a <addr>                 Bind address [default: *]
  -p <port>                 Bind port [default: 8088]
  -h, -help                 Show help
  -v, -version              Show version
  -path-prefix <value>      Url path prefix to listen to [default: "/"]
  -cors                     Enable CORS support [default: false]
  -gzip                     Enable gzip compression (deprecated) [default: false]
  -disable-endpoints        Comma separated endpoints to disable. E.g: form,crop,rotate,health [default: ""]
  -key <key>                Define API key for authorization
  -mount <path>             Mount server local directory
  -http-cache-ttl <num>     The TTL in seconds. Adds caching headers to locally served files.
  -http-read-timeout <num>  HTTP read timeout in seconds [default: 30]
  -http-write-timeout <num> HTTP write timeout in seconds [default: 30]
  -enable-url-source        Restrict remote image source processing to certain origins (separated by commas)
  -enable-placeholder       Enable image response placeholder to be used in case of error [default: false]
  -enable-auth-forwarding   Forwards X-Forward-Authorization or Authorization header to the image source server. -enable-url-source flag must be defined. Tip: secure your server from public access to prevent attack vectors
  -allowed-origins <urls>   Restrict remote image source processing to certain origins (separated by commas)
  -max-allowed-size <bytes> Restrict maximum size of http image source (in bytes)
  -certfile <path>          TLS certificate file path
  -keyfile <path>           TLS private key file path
  -authorization <value>    Defines a constant Authorization header value passed to all the image source servers. -enable-url-source flag must be defined. This overwrites authorization headers forwarding behavior via X-Forward-Authorization
  -placeholder <path>       Image path to image custom placeholder to be used in case of error. Recommended minimum image size is: 1200x1200
  -concurrency <num>        Throttle concurrency limit per second [default: disabled]
  -burst <num>              Throttle burst max cache size [default: 100]
  -mrelease <num>           OS memory release interval in seconds [default: 30]
  -cpus <num>               Number of used cpu cores.
                            (default for current machine is %d cores)
`

func main() {
	flag.Usage = func() {
		fmt.Fprint(os.Stderr, fmt.Sprintf(usage, Version, runtime.NumCPU()))
	}
	flag.Parse()

	if *aHelp || *aHelpl {
		showUsage()
	}
	if *aVers || *aVersl {
		showVersion()
	}

	// Only required in Go < 1.5
	runtime.GOMAXPROCS(*aCpus)

	port := getPort(*aPort)
	opts := ServerOptions{
		Port:              port,
		Address:           *aAddr,
		CORS:              *aCors,
		AuthForwarding:    *aAuthForwarding,
		EnableURLSource:   *aEnableURLSource,
		EnablePlaceholder: *aEnablePlaceholder,
		PathPrefix:        *aPathPrefix,
		APIKey:            *aKey,
		Concurrency:       *aConcurrency,
		Burst:             *aBurst,
		Mount:             *aMount,
		CertFile:          *aCertFile,
		KeyFile:           *aKeyFile,
		Placeholder:       *aPlaceholder,
		HTTPCacheTTL:      *aHTTPCacheTTL,
		HTTPReadTimeout:   *aReadTimeout,
		HTTPWriteTimeout:  *aWriteTimeout,
		Authorization:     *aAuthorization,
		AlloweOrigins:     parseOrigins(*aAlloweOrigins),
		MaxAllowedSize:    *aMaxAllowedSize,
	}

	// Show warning if gzip flag is passed
	if *aGzip {
		fmt.Println("warning: -gzip flag is deprecated and will not have effect")
	}

	// Create a memory release goroutine
	if *aMRelease > 0 {
		memoryRelease(*aMRelease)
	}

	// Check if the mount directory exists, if present
	if *aMount != "" {
		checkMountDirectory(*aMount)
	}

	// Validate HTTP cache param, if present
	if *aHTTPCacheTTL != -1 {
		checkHttpCacheTtl(*aHTTPCacheTTL)
	}

	// Parse endpoint names to disabled, if present
	if *aDisableEndpoints != "" {
		opts.Endpoints = parseEndpoints(*aDisableEndpoints)
	}

	// Read placeholder image, if required
	if *aPlaceholder != "" {
		buf, err := ioutil.ReadFile(*aPlaceholder)
		if err != nil {
			exitWithError("cannot start the server: %s", err)
		}

		imageType := bimg.DetermineImageType(buf)
		if !bimg.IsImageTypeSupportedByVips(imageType).Load {
			exitWithError("Placeholder image type is not supported. Only JPEG, PNG or WEBP are supported")
		}

		opts.PlaceholderImage = buf
	} else if *aEnablePlaceholder {
		// Expose default placeholder
		opts.PlaceholderImage = placeholder
	}

	debug("imaginary server listening on port :%d/%s", opts.Port, strings.TrimPrefix(opts.PathPrefix, "/"))

	// Load image source providers
	LoadSources(opts)

	// Start the server
	err := Server(opts)
	if err != nil {
		exitWithError("cannot start the server: %s", err)
	}
}

func getPort(port int) int {
	if portEnv := os.Getenv("PORT"); portEnv != "" {
		newPort, _ := strconv.Atoi(portEnv)
		if newPort > 0 {
			port = newPort
		}
	}
	return port
}

func showUsage() {
	flag.Usage()
	os.Exit(1)
}

func showVersion() {
	fmt.Println(Version)
	os.Exit(1)
}

func checkMountDirectory(path string) {
	src, err := os.Stat(path)
	if err != nil {
		exitWithError("error while mounting directory: %s", err)
	}
	if src.IsDir() == false {
		exitWithError("mount path is not a directory: %s", path)
	}
	if path == "/" {
		exitWithError("cannot mount root directory for security reasons")
	}
}

func checkHttpCacheTtl(ttl int) {
	if ttl < -1 || ttl > 31556926 {
		exitWithError("The -http-cache-ttl flag only accepts a value from 0 to 31556926")
	}

	if ttl == 0 {
		debug("Adding HTTP cache control headers set to prevent caching.")
	}
}

func parseOrigins(origins string) []*url.URL {
	urls := []*url.URL{}
	if origins == "" {
		return urls
	}
	for _, origin := range strings.Split(origins, ",") {
		u, err := url.Parse(origin)
		if err != nil {
			continue
		}
		urls = append(urls, u)
	}
	return urls
}

func parseEndpoints(input string) Endpoints {
	endpoints := Endpoints{}
	for _, endpoint := range strings.Split(input, ",") {
		endpoint = strings.ToLower(strings.TrimSpace(endpoint))
		if endpoint != "" {
			endpoints = append(endpoints, endpoint)
		}
	}
	return endpoints
}

func memoryRelease(interval int) {
	ticker := time.NewTicker(time.Duration(interval) * time.Second)
	go func() {
		for range ticker.C {
			debug("FreeOSMemory()")
			d.FreeOSMemory()
		}
	}()
}

func exitWithError(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args)
	os.Exit(1)
}

func debug(msg string, values ...interface{}) {
	debug := os.Getenv("DEBUG")
	if debug == "imaginary" || debug == "*" {
		log.Printf(msg, values...)
	}
}
<|endoftext|>


<commit_before>package manta

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"path"

	uuid "github.com/hashicorp/go-uuid"
	"github.com/hashicorp/terraform/state"
	"github.com/hashicorp/terraform/state/remote"
	tritonErrors "github.com/joyent/triton-go/errors"
	"github.com/joyent/triton-go/storage"
)

const (
	mantaDefaultRootStore = "/stor"
	lockFileName          = "tflock"
)

type RemoteClient struct {
	storageClient *storage.StorageClient
	directoryName string
	keyName       string
	statePath     string
}

func (c *RemoteClient) Get() (*remote.Payload, error) {
	output, err := c.storageClient.Objects().Get(context.Background(), &storage.GetObjectInput{
		ObjectPath: path.Join(mantaDefaultRootStore, c.directoryName, c.keyName),
	})
	if err != nil {
		if tritonErrors.IsResourceNotFound(err) {
			return nil, nil
		}
		return nil, err
	}
	defer output.ObjectReader.Close()

	buf := bytes.NewBuffer(nil)
	if _, err := io.Copy(buf, output.ObjectReader); err != nil {
		return nil, fmt.Errorf("Failed to read remote state: %s", err)
	}

	payload := &remote.Payload{
		Data: buf.Bytes(),
	}

	// If there was no data, then return nil
	if len(payload.Data) == 0 {
		return nil, nil
	}

	return payload, nil

}

func (c *RemoteClient) Put(data []byte) error {
	contentType := "application/json"
	contentLength := int64(len(data))

	params := &storage.PutObjectInput{
		ContentType:   contentType,
		ContentLength: uint64(contentLength),
		ObjectPath:    path.Join(mantaDefaultRootStore, c.directoryName, c.keyName),
		ObjectReader:  bytes.NewReader(data),
	}

	log.Printf("[DEBUG] Uploading remote state to Manta: %#v", params)
	err := c.storageClient.Objects().Put(context.Background(), params)
	if err != nil {
		return err
	}

	return nil
}

func (c *RemoteClient) Delete() error {
	err := c.storageClient.Objects().Delete(context.Background(), &storage.DeleteObjectInput{
		ObjectPath: path.Join(mantaDefaultRootStore, c.directoryName, c.keyName),
	})

	return err
}

func (c *RemoteClient) Lock(info *state.LockInfo) (string, error) {
	//At Joyent, we want to make sure that the State directory exists before we interact with it
	//We don't expect users to have to create it in advance
	//The order of operations of Backend State as follows:
	// * Get - if this doesn't exist then we continue as though it's new
	// * Lock - we make sure that the state directory exists as it's the entrance to writing to Manta
	// * Put - put the state up there
	// * Unlock - unlock the directory
	//We can always guarantee that the user can put their state in the specified location because of this
	err := c.storageClient.Dir().Put(context.Background(), &storage.PutDirectoryInput{
		DirectoryName: path.Join(mantaDefaultRootStore, c.directoryName),
	})
	if err != nil {
		return "", err
	}

	//firstly we want to check that a lock doesn't already exist
	lockErr := &state.LockError{}
	lockInfo, err := c.getLockInfo()
	if err != nil {
		if tritonErrors.IsResourceNotFound(err) {
			lockErr.Err = fmt.Errorf("failed to retrieve lock info: %s", err)
			return "", lockErr
		}
	}

	if lockInfo != nil {
		lockErr := &state.LockError{
			Err:  fmt.Errorf("A lock is already acquired"),
			Info: lockInfo,
		}
		return "", lockErr
	}

	info.Path = path.Join(c.directoryName, lockFileName)

	if info.ID == "" {
		lockID, err := uuid.GenerateUUID()
		if err != nil {
			return "", err
		}

		info.ID = lockID
	}

	data := info.Marshal()

	contentType := "application/json"
	contentLength := int64(len(data))

	params := &storage.PutObjectInput{
		ContentType:   contentType,
		ContentLength: uint64(contentLength),
		ObjectPath:    path.Join(mantaDefaultRootStore, c.directoryName, lockFileName),
		ObjectReader:  bytes.NewReader(data),
	}

	log.Printf("[DEBUG] Creating manta state lock: %#v", params)
	err = c.storageClient.Objects().Put(context.Background(), params)
	if err != nil {
		return "", err
	}

	return info.ID, nil
}

func (c *RemoteClient) Unlock(id string) error {
	lockErr := &state.LockError{}

	lockInfo, err := c.getLockInfo()
	if err != nil {
		lockErr.Err = fmt.Errorf("failed to retrieve lock info: %s", err)
		return lockErr
	}
	lockErr.Info = lockInfo

	if lockInfo.ID != id {
		lockErr.Err = fmt.Errorf("lock id %q does not match existing lock", id)
		return lockErr
	}

	err = c.storageClient.Objects().Delete(context.Background(), &storage.DeleteObjectInput{
		ObjectPath: path.Join(mantaDefaultRootStore, c.directoryName, lockFileName),
	})

	return err
}

func (c *RemoteClient) getLockInfo() (*state.LockInfo, error) {
	output, err := c.storageClient.Objects().Get(context.Background(), &storage.GetObjectInput{
		ObjectPath: path.Join(mantaDefaultRootStore, c.directoryName, lockFileName),
	})
	if err != nil {
		return nil, err
	}

	defer output.ObjectReader.Close()

	buf := bytes.NewBuffer(nil)
	if _, err := io.Copy(buf, output.ObjectReader); err != nil {
		return nil, fmt.Errorf("Failed to read lock info: %s", err)
	}

	lockInfo := &state.LockInfo{}
	err = json.Unmarshal(buf.Bytes(), lockInfo)
	if err != nil {
		return nil, err
	}

	return lockInfo, nil
}
<commit_msg>backend/manta: Manta Backend was not dealing with a ResourceNotFound<commit_after>package manta

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"path"

	uuid "github.com/hashicorp/go-uuid"
	"github.com/hashicorp/terraform/state"
	"github.com/hashicorp/terraform/state/remote"
	tritonErrors "github.com/joyent/triton-go/errors"
	"github.com/joyent/triton-go/storage"
)

const (
	mantaDefaultRootStore = "/stor"
	lockFileName          = "tflock"
)

type RemoteClient struct {
	storageClient *storage.StorageClient
	directoryName string
	keyName       string
	statePath     string
}

func (c *RemoteClient) Get() (*remote.Payload, error) {
	output, err := c.storageClient.Objects().Get(context.Background(), &storage.GetObjectInput{
		ObjectPath: path.Join(mantaDefaultRootStore, c.directoryName, c.keyName),
	})
	if err != nil {
		if tritonErrors.IsResourceNotFound(err) {
			return nil, nil
		}
		return nil, err
	}
	defer output.ObjectReader.Close()

	buf := bytes.NewBuffer(nil)
	if _, err := io.Copy(buf, output.ObjectReader); err != nil {
		return nil, fmt.Errorf("Failed to read remote state: %s", err)
	}

	payload := &remote.Payload{
		Data: buf.Bytes(),
	}

	// If there was no data, then return nil
	if len(payload.Data) == 0 {
		return nil, nil
	}

	return payload, nil

}

func (c *RemoteClient) Put(data []byte) error {
	contentType := "application/json"
	contentLength := int64(len(data))

	params := &storage.PutObjectInput{
		ContentType:   contentType,
		ContentLength: uint64(contentLength),
		ObjectPath:    path.Join(mantaDefaultRootStore, c.directoryName, c.keyName),
		ObjectReader:  bytes.NewReader(data),
	}

	log.Printf("[DEBUG] Uploading remote state to Manta: %#v", params)
	err := c.storageClient.Objects().Put(context.Background(), params)
	if err != nil {
		return err
	}

	return nil
}

func (c *RemoteClient) Delete() error {
	err := c.storageClient.Objects().Delete(context.Background(), &storage.DeleteObjectInput{
		ObjectPath: path.Join(mantaDefaultRootStore, c.directoryName, c.keyName),
	})

	return err
}

func (c *RemoteClient) Lock(info *state.LockInfo) (string, error) {
	//At Joyent, we want to make sure that the State directory exists before we interact with it
	//We don't expect users to have to create it in advance
	//The order of operations of Backend State as follows:
	// * Get - if this doesn't exist then we continue as though it's new
	// * Lock - we make sure that the state directory exists as it's the entrance to writing to Manta
	// * Put - put the state up there
	// * Unlock - unlock the directory
	//We can always guarantee that the user can put their state in the specified location because of this
	err := c.storageClient.Dir().Put(context.Background(), &storage.PutDirectoryInput{
		DirectoryName: path.Join(mantaDefaultRootStore, c.directoryName),
	})
	if err != nil {
		return "", err
	}

	//firstly we want to check that a lock doesn't already exist
	lockErr := &state.LockError{}
	lockInfo, err := c.getLockInfo()
	if err != nil {
		if !tritonErrors.IsResourceNotFound(err) {
			lockErr.Err = fmt.Errorf("failed to retrieve lock info: %s", err)
			return "", lockErr
		}
	}

	if lockInfo != nil {
		lockErr := &state.LockError{
			Err:  fmt.Errorf("A lock is already acquired"),
			Info: lockInfo,
		}
		return "", lockErr
	}

	info.Path = path.Join(c.directoryName, lockFileName)

	if info.ID == "" {
		lockID, err := uuid.GenerateUUID()
		if err != nil {
			return "", err
		}

		info.ID = lockID
	}

	data := info.Marshal()

	contentType := "application/json"
	contentLength := int64(len(data))

	params := &storage.PutObjectInput{
		ContentType:   contentType,
		ContentLength: uint64(contentLength),
		ObjectPath:    path.Join(mantaDefaultRootStore, c.directoryName, lockFileName),
		ObjectReader:  bytes.NewReader(data),
	}

	log.Printf("[DEBUG] Creating manta state lock: %#v", params)
	err = c.storageClient.Objects().Put(context.Background(), params)
	if err != nil {
		return "", err
	}

	return info.ID, nil
}

func (c *RemoteClient) Unlock(id string) error {
	lockErr := &state.LockError{}

	lockInfo, err := c.getLockInfo()
	if err != nil {
		lockErr.Err = fmt.Errorf("failed to retrieve lock info: %s", err)
		return lockErr
	}
	lockErr.Info = lockInfo

	if lockInfo.ID != id {
		lockErr.Err = fmt.Errorf("lock id %q does not match existing lock", id)
		return lockErr
	}

	err = c.storageClient.Objects().Delete(context.Background(), &storage.DeleteObjectInput{
		ObjectPath: path.Join(mantaDefaultRootStore, c.directoryName, lockFileName),
	})

	return err
}

func (c *RemoteClient) getLockInfo() (*state.LockInfo, error) {
	output, err := c.storageClient.Objects().Get(context.Background(), &storage.GetObjectInput{
		ObjectPath: path.Join(mantaDefaultRootStore, c.directoryName, lockFileName),
	})
	if err != nil {
		return nil, err
	}

	defer output.ObjectReader.Close()

	buf := bytes.NewBuffer(nil)
	if _, err := io.Copy(buf, output.ObjectReader); err != nil {
		return nil, fmt.Errorf("Failed to read lock info: %s", err)
	}

	lockInfo := &state.LockInfo{}
	err = json.Unmarshal(buf.Bytes(), lockInfo)
	if err != nil {
		return nil, err
	}

	return lockInfo, nil
}
<|endoftext|>


<commit_before>/*
   Copyright 2017 Shlomi Noach, GitHub Inc.

   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 logic

import (
	"encoding/json"
	"io"

	"github.com/github/orchestrator/go/db"
	"github.com/github/orchestrator/go/inst"

	"github.com/openark/golib/log"
	"github.com/openark/golib/sqlutils"
)

type SnapshotData struct {
	Keys             []inst.InstanceKey
	RecoveryDisabled bool

	ClusterAlias,
	ClusterAliasOverride,
	ClusterDomainName,
	HostAttributes,
	AccessToken,
	PoolInstances,
	HostnameResolves,
	HostnameUnresolves,
	DowntimedInstances,
	Candidates,
	Detections,
	Recovery,
	RecoverySteps sqlutils.NamedResultData
}

func NewSnapshotData() *SnapshotData {
	return &SnapshotData{}
}

func readTableData(tableName string, data *sqlutils.NamedResultData) error {
	orcdb, err := db.OpenOrchestrator()
	if err != nil {
		return log.Errore(err)
	}
	*data, err = sqlutils.ScanTable(orcdb, tableName)
	return log.Errore(err)
}

func writeTableData(tableName string, data *sqlutils.NamedResultData) error {
	orcdb, err := db.OpenOrchestrator()
	if err != nil {
		return log.Errore(err)
	}
	err = sqlutils.WriteTable(orcdb, tableName, *data)
	return log.Errore(err)
}

func CreateSnapshotData() *SnapshotData {
	snapshotData := NewSnapshotData()

	// keys
	snapshotData.Keys, _ = inst.ReadAllInstanceKeys()
	snapshotData.RecoveryDisabled, _ = IsRecoveryDisabled()

	readTableData("cluster_alias", &snapshotData.ClusterAlias)
	readTableData("cluster_alias_override", &snapshotData.ClusterAliasOverride)
	readTableData("cluster_domain_name", &snapshotData.ClusterDomainName)
	readTableData("access_token", &snapshotData.AccessToken)
	readTableData("host_attributes", &snapshotData.HostAttributes)
	readTableData("database_instance_pool", &snapshotData.PoolInstances)
	readTableData("hostname_resolve", &snapshotData.HostnameResolves)
	readTableData("hostname_unresolve", &snapshotData.HostnameUnresolves)
	readTableData("database_instance_downtime", &snapshotData.DowntimedInstances)
	readTableData("candidate_database_instance", &snapshotData.Candidates)
	readTableData("topology_failure_detection", &snapshotData.Detections)
	readTableData("topology_recovery", &snapshotData.Recovery)
	readTableData("topology_recovery_steps", &snapshotData.RecoverySteps)

	log.Debugf("raft snapshot data created")
	return snapshotData
}

type SnapshotDataCreatorApplier struct {
}

func NewSnapshotDataCreatorApplier() *SnapshotDataCreatorApplier {
	generator := &SnapshotDataCreatorApplier{}
	return generator
}

func (this *SnapshotDataCreatorApplier) GetData() (data []byte, err error) {
	snapshotData := CreateSnapshotData()
	return json.Marshal(snapshotData)
}

func (this *SnapshotDataCreatorApplier) Restore(rc io.ReadCloser) error {
	snapshotData := NewSnapshotData()
	// if err := json.Unmarshal(data, &snapshotData); err != nil {
	if err := json.NewDecoder(rc).Decode(&snapshotData); err != nil {
		return err
	}

	// keys
	{
		instanceKeyMap := inst.NewInstanceKeyMap()
		instanceKeyMap.AddKeys(snapshotData.Keys)

		allKeys, _ := inst.ReadAllInstanceKeys()
		for _, key := range allKeys {
			if !instanceKeyMap.HasKey(key) {
				inst.ForgetInstance(&key)
			}
		}
		for _, key := range instanceKeyMap.GetInstanceKeys() {
			key := key
			go func() {
				snapshotDiscoveryKeys <- key
			}()
		}
	}
	writeTableData("cluster_alias", &snapshotData.ClusterAlias)
	writeTableData("cluster_alias_override", &snapshotData.ClusterAliasOverride)
	writeTableData("cluster_domain_name", &snapshotData.ClusterDomainName)
	writeTableData("access_token", &snapshotData.AccessToken)
	writeTableData("host_attributes", &snapshotData.HostAttributes)
	writeTableData("database_instance_pool", &snapshotData.PoolInstances)
	writeTableData("hostname_resolve", &snapshotData.HostnameResolves)
	writeTableData("hostname_unresolve", &snapshotData.HostnameUnresolves)
	writeTableData("database_instance_downtime", &snapshotData.DowntimedInstances)
	writeTableData("candidate_database_instance", &snapshotData.Candidates)
	writeTableData("topology_recovery", &snapshotData.Recovery)
	writeTableData("topology_failure_detection", &snapshotData.Detections)
	writeTableData("topology_recovery_steps", &snapshotData.RecoverySteps)

	// recovery disable
	{
		SetRecoveryDisabled(snapshotData.RecoveryDisabled)
	}
	log.Debugf("raft restore applied")
	return nil
}
<commit_msg>snapshots are gzipped<commit_after>/*
   Copyright 2017 Shlomi Noach, GitHub Inc.

   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 logic

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"io"

	"github.com/github/orchestrator/go/db"
	"github.com/github/orchestrator/go/inst"

	"github.com/openark/golib/log"
	"github.com/openark/golib/sqlutils"
)

type SnapshotData struct {
	Keys             []inst.InstanceKey
	RecoveryDisabled bool

	ClusterAlias,
	ClusterAliasOverride,
	ClusterDomainName,
	HostAttributes,
	AccessToken,
	PoolInstances,
	HostnameResolves,
	HostnameUnresolves,
	DowntimedInstances,
	Candidates,
	Detections,
	Recovery,
	RecoverySteps sqlutils.NamedResultData
}

func NewSnapshotData() *SnapshotData {
	return &SnapshotData{}
}

func readTableData(tableName string, data *sqlutils.NamedResultData) error {
	orcdb, err := db.OpenOrchestrator()
	if err != nil {
		return log.Errore(err)
	}
	*data, err = sqlutils.ScanTable(orcdb, tableName)
	return log.Errore(err)
}

func writeTableData(tableName string, data *sqlutils.NamedResultData) error {
	orcdb, err := db.OpenOrchestrator()
	if err != nil {
		return log.Errore(err)
	}
	err = sqlutils.WriteTable(orcdb, tableName, *data)
	return log.Errore(err)
}

func CreateSnapshotData() *SnapshotData {
	snapshotData := NewSnapshotData()

	// keys
	snapshotData.Keys, _ = inst.ReadAllInstanceKeys()
	snapshotData.RecoveryDisabled, _ = IsRecoveryDisabled()

	readTableData("cluster_alias", &snapshotData.ClusterAlias)
	readTableData("cluster_alias_override", &snapshotData.ClusterAliasOverride)
	readTableData("cluster_domain_name", &snapshotData.ClusterDomainName)
	readTableData("access_token", &snapshotData.AccessToken)
	readTableData("host_attributes", &snapshotData.HostAttributes)
	readTableData("database_instance_pool", &snapshotData.PoolInstances)
	readTableData("hostname_resolve", &snapshotData.HostnameResolves)
	readTableData("hostname_unresolve", &snapshotData.HostnameUnresolves)
	readTableData("database_instance_downtime", &snapshotData.DowntimedInstances)
	readTableData("candidate_database_instance", &snapshotData.Candidates)
	readTableData("topology_failure_detection", &snapshotData.Detections)
	readTableData("topology_recovery", &snapshotData.Recovery)
	readTableData("topology_recovery_steps", &snapshotData.RecoverySteps)

	log.Debugf("raft snapshot data created")
	return snapshotData
}

type SnapshotDataCreatorApplier struct {
}

func NewSnapshotDataCreatorApplier() *SnapshotDataCreatorApplier {
	generator := &SnapshotDataCreatorApplier{}
	return generator
}

func (this *SnapshotDataCreatorApplier) GetData() (data []byte, err error) {
	log.Debugf(".............. writing snap")
	snapshotData := CreateSnapshotData()
	b, err := json.Marshal(snapshotData)
	if err != nil {
		return b, err
	}
	var buf bytes.Buffer
	zw := gzip.NewWriter(&buf)
	if _, err := zw.Write(b); err != nil {
		return b, err
	}
	if err := zw.Close(); err != nil {
		return b, err
	}
	return buf.Bytes(), nil
}

func (this *SnapshotDataCreatorApplier) Restore(rc io.ReadCloser) error {
	log.Debugf(".............. reading snap")
	snapshotData := NewSnapshotData()
	zr, err := gzip.NewReader(rc)
	if err != nil {
		return err
	}
	if err := json.NewDecoder(zr).Decode(&snapshotData); err != nil {
		return err
	}

	// keys
	{
		instanceKeyMap := inst.NewInstanceKeyMap()
		instanceKeyMap.AddKeys(snapshotData.Keys)

		allKeys, _ := inst.ReadAllInstanceKeys()
		for _, key := range allKeys {
			if !instanceKeyMap.HasKey(key) {
				inst.ForgetInstance(&key)
			}
		}
		for _, key := range instanceKeyMap.GetInstanceKeys() {
			key := key
			go func() {
				snapshotDiscoveryKeys <- key
			}()
		}
	}
	writeTableData("cluster_alias", &snapshotData.ClusterAlias)
	writeTableData("cluster_alias_override", &snapshotData.ClusterAliasOverride)
	writeTableData("cluster_domain_name", &snapshotData.ClusterDomainName)
	writeTableData("access_token", &snapshotData.AccessToken)
	writeTableData("host_attributes", &snapshotData.HostAttributes)
	writeTableData("database_instance_pool", &snapshotData.PoolInstances)
	writeTableData("hostname_resolve", &snapshotData.HostnameResolves)
	writeTableData("hostname_unresolve", &snapshotData.HostnameUnresolves)
	writeTableData("database_instance_downtime", &snapshotData.DowntimedInstances)
	writeTableData("candidate_database_instance", &snapshotData.Candidates)
	writeTableData("topology_recovery", &snapshotData.Recovery)
	writeTableData("topology_failure_detection", &snapshotData.Detections)
	writeTableData("topology_recovery_steps", &snapshotData.RecoverySteps)

	// recovery disable
	{
		SetRecoveryDisabled(snapshotData.RecoveryDisabled)
	}
	log.Debugf("raft restore applied")
	return nil
}
<|endoftext|>