<commit_before>package data

import (
	"net/url"
	"testing"

	"github.com/spf13/afero"

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

func TestReadMerge(t *testing.T) {
	jsonContent := `{"hello": "world"}`
	yamlContent := "hello: earth\ngoodnight: moon\n"
	arrayContent := `["hello", "world"]`

	mergedContent := "goodnight: moon\nhello: world\n"

	fs := afero.NewMemMapFs()

	_ = fs.Mkdir("/tmp", 0777)
	f, _ := fs.Create("/tmp/jsonfile.json")
	_, _ = f.WriteString(jsonContent)
	f, _ = fs.Create("/tmp/array.json")
	_, _ = f.WriteString(arrayContent)
	f, _ = fs.Create("/tmp/yamlfile.yaml")
	_, _ = f.WriteString(yamlContent)
	f, _ = fs.Create("/tmp/textfile.txt")
	_, _ = f.WriteString(`plain text...`)

	source := &Source{Alias: "foo", URL: mustParseURL("merge:file:///tmp/jsonfile.json|file:///tmp/yamlfile.yaml")}
	source.fs = fs
	d := &Data{
		Sources: map[string]*Source{
			"foo":       source,
			"bar":       {Alias: "bar", URL: mustParseURL("file:///tmp/jsonfile.json")},
			"baz":       {Alias: "baz", URL: mustParseURL("file:///tmp/yamlfile.yaml")},
			"text":      {Alias: "text", URL: mustParseURL("file:///tmp/textfile.txt")},
			"badscheme": {Alias: "badscheme", URL: mustParseURL("bad:///scheme.json")},
			"badtype":   {Alias: "badtype", URL: mustParseURL("file:///tmp/textfile.txt?type=foo/bar")},
			"array":     {Alias: "array", URL: mustParseURL("file:///tmp/array.json?type=" + url.QueryEscape(jsonArrayMimetype))},
		},
	}

	actual, err := d.readMerge(source)
	assert.NoError(t, err)
	assert.Equal(t, mergedContent, string(actual))

	source.URL = mustParseURL("merge:bar|baz")
	actual, err = d.readMerge(source)
	assert.NoError(t, err)
	assert.Equal(t, mergedContent, string(actual))

	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json")
	_, err = d.readMerge(source)
	assert.Error(t, err)

	source.URL = mustParseURL("merge:bogusalias|file:///tmp/jsonfile.json")
	_, err = d.readMerge(source)
	assert.Error(t, err)

	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json|badscheme")
	_, err = d.readMerge(source)
	assert.Error(t, err)

	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json|badtype")
	_, err = d.readMerge(source)
	assert.Error(t, err)

	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json|array")
	_, err = d.readMerge(source)
	assert.Error(t, err)
}

func TestMergeData(t *testing.T) {
	def := map[string]interface{}{
		"f": true,
		"t": false,
		"z": "def",
	}
	out, err := mergeData([]map[string]interface{}{def})
	assert.NoError(t, err)
	assert.Equal(t, "f: true\nt: false\nz: def\n", string(out))

	over := map[string]interface{}{
		"f": false,
		"t": true,
		"z": "over",
	}
	out, err = mergeData([]map[string]interface{}{over, def})
	assert.NoError(t, err)
	assert.Equal(t, "f: false\nt: true\nz: over\n", string(out))

	over = map[string]interface{}{
		"f": false,
		"t": true,
		"z": "over",
		"m": map[string]interface{}{
			"a": "aaa",
		},
	}
	out, err = mergeData([]map[string]interface{}{over, def})
	assert.NoError(t, err)
	assert.Equal(t, "f: false\nm:\n  a: aaa\nt: true\nz: over\n", string(out))

	uber := map[string]interface{}{
		"z": "über",
	}
	out, err = mergeData([]map[string]interface{}{uber, over, def})
	assert.NoError(t, err)
	assert.Equal(t, "f: false\nm:\n  a: aaa\nt: true\nz: über\n", string(out))

	uber = map[string]interface{}{
		"m": "notamap",
		"z": map[string]interface{}{
			"b": "bbb",
		},
	}
	out, err = mergeData([]map[string]interface{}{uber, over, def})
	assert.NoError(t, err)
	assert.Equal(t, "f: false\nm: notamap\nt: true\nz:\n  b: bbb\n", string(out))

	uber = map[string]interface{}{
		"m": map[string]interface{}{
			"b": "bbb",
		},
	}
	out, err = mergeData([]map[string]interface{}{uber, over, def})
	assert.NoError(t, err)
	assert.Equal(t, "f: false\nm:\n  a: aaa\n  b: bbb\nt: true\nz: over\n", string(out))
}
<commit_msg>Add test for relative paths in merge datasources<commit_after>package data

import (
	"net/url"
	"os"
	"path/filepath"
	"testing"

	"github.com/spf13/afero"

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

func TestReadMerge(t *testing.T) {
	jsonContent := `{"hello": "world"}`
	yamlContent := "hello: earth\ngoodnight: moon\n"
	arrayContent := `["hello", "world"]`

	mergedContent := "goodnight: moon\nhello: world\n"

	fs := afero.NewMemMapFs()

	_ = fs.Mkdir("/tmp", 0777)
	f, _ := fs.Create("/tmp/jsonfile.json")
	_, _ = f.WriteString(jsonContent)
	f, _ = fs.Create("/tmp/array.json")
	_, _ = f.WriteString(arrayContent)
	f, _ = fs.Create("/tmp/yamlfile.yaml")
	_, _ = f.WriteString(yamlContent)
	f, _ = fs.Create("/tmp/textfile.txt")
	_, _ = f.WriteString(`plain text...`)

	wd, _ := os.Getwd()
	_ = fs.Mkdir(wd, 0777)
	f, _ = fs.Create(filepath.Join(wd, "jsonfile.json"))
	_, _ = f.WriteString(jsonContent)
	f, _ = fs.Create(filepath.Join(wd, "array.json"))
	_, _ = f.WriteString(arrayContent)
	f, _ = fs.Create(filepath.Join(wd, "yamlfile.yaml"))
	_, _ = f.WriteString(yamlContent)
	f, _ = fs.Create(filepath.Join(wd, "textfile.txt"))
	_, _ = f.WriteString(`plain text...`)

	source := &Source{Alias: "foo", URL: mustParseURL("merge:file:///tmp/jsonfile.json|file:///tmp/yamlfile.yaml")}
	source.fs = fs
	d := &Data{
		Sources: map[string]*Source{
			"foo":       source,
			"bar":       {Alias: "bar", URL: mustParseURL("file:///tmp/jsonfile.json")},
			"baz":       {Alias: "baz", URL: mustParseURL("file:///tmp/yamlfile.yaml")},
			"text":      {Alias: "text", URL: mustParseURL("file:///tmp/textfile.txt")},
			"badscheme": {Alias: "badscheme", URL: mustParseURL("bad:///scheme.json")},
			"badtype":   {Alias: "badtype", URL: mustParseURL("file:///tmp/textfile.txt?type=foo/bar")},
			"array":     {Alias: "array", URL: mustParseURL("file:///tmp/array.json?type=" + url.QueryEscape(jsonArrayMimetype))},
		},
	}

	actual, err := d.readMerge(source)
	assert.NoError(t, err)
	assert.Equal(t, mergedContent, string(actual))

	source.URL = mustParseURL("merge:bar|baz")
	actual, err = d.readMerge(source)
	assert.NoError(t, err)
	assert.Equal(t, mergedContent, string(actual))

	source.URL = mustParseURL("merge:./jsonfile.json|baz")
	actual, err = d.readMerge(source)
	assert.NoError(t, err)
	assert.Equal(t, mergedContent, string(actual))

	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json")
	_, err = d.readMerge(source)
	assert.Error(t, err)

	source.URL = mustParseURL("merge:bogusalias|file:///tmp/jsonfile.json")
	_, err = d.readMerge(source)
	assert.Error(t, err)

	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json|badscheme")
	_, err = d.readMerge(source)
	assert.Error(t, err)

	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json|badtype")
	_, err = d.readMerge(source)
	assert.Error(t, err)

	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json|array")
	_, err = d.readMerge(source)
	assert.Error(t, err)
}

func TestMergeData(t *testing.T) {
	def := map[string]interface{}{
		"f": true,
		"t": false,
		"z": "def",
	}
	out, err := mergeData([]map[string]interface{}{def})
	assert.NoError(t, err)
	assert.Equal(t, "f: true\nt: false\nz: def\n", string(out))

	over := map[string]interface{}{
		"f": false,
		"t": true,
		"z": "over",
	}
	out, err = mergeData([]map[string]interface{}{over, def})
	assert.NoError(t, err)
	assert.Equal(t, "f: false\nt: true\nz: over\n", string(out))

	over = map[string]interface{}{
		"f": false,
		"t": true,
		"z": "over",
		"m": map[string]interface{}{
			"a": "aaa",
		},
	}
	out, err = mergeData([]map[string]interface{}{over, def})
	assert.NoError(t, err)
	assert.Equal(t, "f: false\nm:\n  a: aaa\nt: true\nz: over\n", string(out))

	uber := map[string]interface{}{
		"z": "über",
	}
	out, err = mergeData([]map[string]interface{}{uber, over, def})
	assert.NoError(t, err)
	assert.Equal(t, "f: false\nm:\n  a: aaa\nt: true\nz: über\n", string(out))

	uber = map[string]interface{}{
		"m": "notamap",
		"z": map[string]interface{}{
			"b": "bbb",
		},
	}
	out, err = mergeData([]map[string]interface{}{uber, over, def})
	assert.NoError(t, err)
	assert.Equal(t, "f: false\nm: notamap\nt: true\nz:\n  b: bbb\n", string(out))

	uber = map[string]interface{}{
		"m": map[string]interface{}{
			"b": "bbb",
		},
	}
	out, err = mergeData([]map[string]interface{}{uber, over, def})
	assert.NoError(t, err)
	assert.Equal(t, "f: false\nm:\n  a: aaa\n  b: bbb\nt: true\nz: over\n", string(out))
}
<|endoftext|>


<commit_before>package manager

import (
	"fmt"
	"github.com/hidu/goutils"
	"log"
	//	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"
)

type ProxyPool struct {
	proxyListActive map[string]*Proxy
	proxyListAll    map[string]*Proxy
	mu              sync.RWMutex

	SessionProxys map[int64]map[string]*Proxy
	ProxyManager  *ProxyManager

	aliveCheckUrl      string
	aliveCheckResponse *http.Response

	checkChan     chan string
	testRunChan   chan bool
	timeout       int
	checkInterval int64

	proxyActiveUsed map[string]string

	Count *ProxyCount
}

func LoadProxyPool(manager *ProxyManager) *ProxyPool {
	log.Println("loading proxy pool...")
	pool := &ProxyPool{}
	pool.ProxyManager = manager
	pool.proxyListActive = make(map[string]*Proxy)
	pool.proxyListAll = make(map[string]*Proxy)
	pool.SessionProxys = make(map[int64]map[string]*Proxy)

	pool.proxyActiveUsed = make(map[string]string)

	pool.checkChan = make(chan string, 100)
	pool.testRunChan = make(chan bool, 1)
	pool.timeout = manager.config.timeout

	pool.aliveCheckUrl = manager.config.aliveCheckUrl
	pool.checkInterval = manager.config.checkInterval
	pool.Count = NewProxyCount()

	if pool.aliveCheckUrl != "" {
		var err error
		urlStr := strings.Replace(pool.aliveCheckUrl, "{%rand}", fmt.Sprintf("%d", time.Now().UnixNano()), -1)
		pool.aliveCheckResponse, err = doRequestGet(urlStr, nil, 3)
		if err != nil {
			log.Println("get origin alive response failed,url:", pool.aliveCheckUrl, "err:", err)
			return nil
		} else {
			log.Println("get alive info suc!url:", pool.aliveCheckUrl, "resp_header:", pool.aliveCheckResponse.Header)
		}
	}

	proxyAll, err := pool.loadConf("pool.conf")
	if err != nil {
		log.Println("pool.conf not exists")
	}
	proxyAllChecked, _ := pool.loadConf("pool_checked.conf")

	pool.proxyListAll = proxyAllChecked
	for _url, proxy := range proxyAll {
		if _, has := pool.proxyListAll[_url]; !has {
			pool.proxyListAll[_url] = proxy
		}
	}
	if len(pool.proxyListAll) == 0 {
		log.Println("proxy pool list is empty")
	}

	go pool.runTest()

	utils.SetInterval(func() {
		pool.runTest()
	}, pool.checkInterval)

	return pool
}

func (pool *ProxyPool) String() string {
	allProxy := []string{}
	for _, proxy := range pool.proxyListAll {
		allProxy = append(allProxy, proxy.String())
	}
	return strings.Join(allProxy, "\n")
}

func (pool *ProxyPool) loadConf(confName string) (map[string]*Proxy, error) {
	proxys := make(map[string]*Proxy)
	confPath := pool.ProxyManager.config.confDir + "/" + confName

	txtFile, err := utils.NewTxtFile(confPath)
	if err != nil {
		log.Println("load proxy pool failed[", confName, "]")
		return proxys, err
	}
	return pool.loadProxysFromTxtFile(txtFile)
}

func (pool *ProxyPool) loadProxysFromTxtFile(txtFile *utils.TxtFile) (map[string]*Proxy, error) {
	proxys := make(map[string]*Proxy)
	defaultValues := make(map[string]string)
	defaultValues["proxy"] = "required"
	defaultValues["weight"] = "1"
	defaultValues["status"] = "1"
	defaultValues["last_check"] = "0"
	defaultValues["check_used"] = "0"

	datas, err := txtFile.KvMapSlice("=", true, defaultValues)
	if err != nil {
		return proxys, err
	}
	for _, kv := range datas {
		proxy := pool.parseProxy(kv)
		if proxy != nil {
			proxys[proxy.proxy] = proxy
		}
	}
	return proxys, nil
}

func (pool *ProxyPool) parseProxy(info map[string]string) *Proxy {
	if info == nil {
		return nil
	}
	proxy := NewProxy(info["proxy"])
	if proxy == nil {
		return nil
	}
	intValues := make(map[string]int)
	intFields := []string{"weight", "status", "check_used", "last_check", "last_check_ok"}
	var err error
	for _, fieldName := range intFields {
		intValues[fieldName], err = strconv.Atoi(info[fieldName])
		if err != nil {
			log.Println("parse [", fieldName, "]failed,not int.err:", err)
			intValues[fieldName] = 0
		}
	}
	proxy.Weight = intValues["weight"]
	proxy.StatusCode = PROXY_STATUS(intValues["status"])
	proxy.CheckUsed = int64(intValues["check_used"])
	proxy.LastCheck = int64(intValues["last_check"])
	proxy.LastCheckOk = int64(intValues["last_check_ok"])
	return proxy
}

func (pool *ProxyPool) GetProxy(proxy_url string) *Proxy {
	pool.mu.RLock()
	defer pool.mu.RUnlock()
	if proxy, has := pool.proxyListAll[proxy_url]; has {
		return proxy
	}
	return nil
}

func (pool *ProxyPool) addProxyActive(proxy_url string) bool {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	if proxy, has := pool.proxyListAll[proxy_url]; has {
		if _, hasAct := pool.proxyListActive[proxy_url]; !hasAct {
			pool.proxyListActive[proxy_url] = proxy
			return true
		}
	}
	return false
}

func (pool *ProxyPool) addProxy(proxy *Proxy) bool {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	if _, has := pool.proxyListAll[proxy.proxy]; !has {
		pool.proxyListAll[proxy.proxy] = proxy
		return true
	}
	return false
}

func (pool *ProxyPool) removeProxyActive(proxy_url string) {
	pool.mu.Lock()
	defer pool.mu.Unlock()
	if _, hasAct := pool.proxyListActive[proxy_url]; hasAct {
		delete(pool.proxyListActive, proxy_url)
	}
}

var errorNoProxy error = fmt.Errorf("no active proxy")

func (pool *ProxyPool) GetOneProxy(logid int64) (*Proxy, error) {
	pool.mu.Lock()
	defer pool.mu.Unlock()
	l := len(pool.proxyListActive)
	if l == 0 {
		return nil, errorNoProxy
	}

	sessionProxys, has := pool.SessionProxys[logid]

	if !has {
		sessionProxys = make(map[string]*Proxy)
		pool.SessionProxys[logid] = sessionProxys
	}

	for _, proxy := range pool.proxyListActive {
		if _, has := pool.proxyActiveUsed[proxy.proxy]; has {
			continue
		}
		if _, has := sessionProxys[proxy.proxy]; !has {
			sessionProxys[proxy.proxy] = proxy
			proxy.Used++
			pool.proxyActiveUsed[proxy.proxy] = "1"
			if len(pool.proxyActiveUsed) >= len(pool.proxyListActive) {
				pool.proxyActiveUsed = make(map[string]string)
			}
			return proxy, nil
		}
	}
	return nil, errorNoProxy
}

func (pool *ProxyPool) CleanSessionProxy(logid int64) {
	pool.mu.Lock()
	defer pool.mu.Unlock()
	if _, has := pool.SessionProxys[logid]; has {
		delete(pool.SessionProxys, logid)
	}
}
func (pool *ProxyPool) runTest() {
	pool.testRunChan <- true
	defer (func() {
		<-pool.testRunChan
	})()
	start := time.Now()
	proxyTotal := len(pool.proxyListAll)
	log.Println("start test all proxy,total=", proxyTotal)

	var wg sync.WaitGroup
	for name := range pool.proxyListAll {
		wg.Add(1)
		go (func(proxyUrl string) {
			pool.TestProxyAddActive(proxyUrl)
			wg.Done()
		})(name)
	}
	wg.Wait()

	used := time.Now().Sub(start)
	log.Println("test all proxy finish,total:", proxyTotal, "used:", used, "activeTotal:", len(pool.proxyListActive))

	testResultFile := pool.ProxyManager.config.confDir + "/pool_checked.conf"
	utils.File_put_contents(testResultFile, []byte(pool.String()))
}

func (pool *ProxyPool) TestProxyAddActive(proxy_url string) bool {
	proxy := pool.GetProxy(proxy_url)
	if proxy == nil {
		return false
	}
	isOk := pool.TestProxy(proxy)
	if isOk {
		pool.addProxyActive(proxy.proxy)
	} else {
		pool.removeProxyActive(proxy.proxy)
	}
	return true
}

func (pool *ProxyPool) TestProxy(proxy *Proxy) bool {
	pool.checkChan <- proxy.proxy
	start := time.Now()
	defer (func() {
		<-pool.checkChan
	})()

	if start.Unix()-proxy.LastCheck < pool.checkInterval {
		return proxy.IsOk()
	}

	proxy.StatusCode = PROXY_STATUS_UNAVAILABLE

	testlog := func(msg ...interface{}) {
		used := time.Now().Sub(start)
		proxy.CheckUsed = used.Nanoseconds() / 1000000
		proxy.LastCheck = start.Unix()
		log.Println("test proxy", proxy.proxy, fmt.Sprint(msg...), "used:", proxy.CheckUsed, "ms")
	}

	if pool.aliveCheckUrl != "" {
		urlStr := strings.Replace(pool.aliveCheckUrl, "{%rand}", fmt.Sprintf("%d", start.UnixNano()), -1)
		resp, err := doRequestGet(urlStr, proxy, pool.timeout/2)
		if err != nil {
			testlog("failed,", err.Error())
			return false
		} else {
			cur_len := resp.Header.Get("Content-Length")
			check_len := pool.aliveCheckResponse.Header.Get("Content-Length")
			if cur_len != check_len {
				testlog("failed ,content-length wrong,[", check_len, "!=", cur_len, "]")
				return false
			}
		}
	} else {
		host, port, err := utils.Net_getHostPortFromUrl(proxy.proxy)
		if err != nil {
			testlog("failed,proxy url err:", err)
			return false
		}
		conn, netErr := net.Dial("tcp", fmt.Sprintf("%s:%d", host, port))
		if netErr != nil {
			testlog("failed", netErr)
			return false
		}
		conn.Close()
	}
	proxy.StatusCode = PROXY_STATUS_ACTIVE
	proxy.LastCheckOk = proxy.LastCheck
	testlog("pass")
	return true
}

func (pool *ProxyPool) MarkProxyStatus(proxy *Proxy, status PROXY_USED_STATUS) {
	proxy.Count.MarkStatus(status)
	pool.Count.MarkStatus(status)
}

func doRequestGet(urlStr string, proxy *Proxy, timeout_sec int) (resp *http.Response, err error) {
	client := &http.Client{}
	if timeout_sec > 0 {
		client.Timeout = time.Duration(timeout_sec) * time.Second
	}
	if proxy != nil {
		proxyGetFn := func(req *http.Request) (*url.URL, error) {
			return proxy.URL, nil
		}
		client.Transport = &http.Transport{Proxy: proxyGetFn}
	}
	req, _ := http.NewRequest("GET", urlStr, nil)
	return client.Do(req)
}
<commit_msg>clean bad proxy<commit_after>package manager

import (
	"fmt"
	"github.com/hidu/goutils"
	"log"
	//	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"
)

type ProxyPool struct {
	proxyListActive map[string]*Proxy
	proxyListAll    map[string]*Proxy
	mu              sync.RWMutex

	SessionProxys map[int64]map[string]*Proxy
	ProxyManager  *ProxyManager

	aliveCheckUrl      string
	aliveCheckResponse *http.Response

	checkChan     chan string
	testRunChan   chan bool
	timeout       int
	checkInterval int64

	proxyActiveUsed map[string]string

	Count *ProxyCount
}

func LoadProxyPool(manager *ProxyManager) *ProxyPool {
	log.Println("loading proxy pool...")
	pool := &ProxyPool{}
	pool.ProxyManager = manager
	pool.proxyListActive = make(map[string]*Proxy)
	pool.proxyListAll = make(map[string]*Proxy)
	pool.SessionProxys = make(map[int64]map[string]*Proxy)

	pool.proxyActiveUsed = make(map[string]string)

	pool.checkChan = make(chan string, 100)
	pool.testRunChan = make(chan bool, 1)
	pool.timeout = manager.config.timeout

	pool.aliveCheckUrl = manager.config.aliveCheckUrl
	pool.checkInterval = manager.config.checkInterval
	pool.Count = NewProxyCount()

	if pool.aliveCheckUrl != "" {
		var err error
		urlStr := strings.Replace(pool.aliveCheckUrl, "{%rand}", fmt.Sprintf("%d", time.Now().UnixNano()), -1)
		pool.aliveCheckResponse, err = doRequestGet(urlStr, nil, 3)
		if err != nil {
			log.Println("get origin alive response failed,url:", pool.aliveCheckUrl, "err:", err)
			return nil
		} else {
			log.Println("get alive info suc!url:", pool.aliveCheckUrl, "resp_header:", pool.aliveCheckResponse.Header)
		}
	}

	proxyAll, err := pool.loadConf("pool.conf")
	if err != nil {
		log.Println("pool.conf not exists")
	}
	proxyAllChecked, _ := pool.loadConf("pool_checked.conf")

	pool.proxyListAll = proxyAllChecked
	for _url, proxy := range proxyAll {
		if _, has := pool.proxyListAll[_url]; !has {
			pool.proxyListAll[_url] = proxy
		}
	}
	if len(pool.proxyListAll) == 0 {
		log.Println("proxy pool list is empty")
	}

	go pool.runTest()

	utils.SetInterval(func() {
		pool.runTest()
	}, pool.checkInterval)

	return pool
}

func (pool *ProxyPool) String() string {
	allProxy := []string{}
	for _, proxy := range pool.proxyListAll {
		allProxy = append(allProxy, proxy.String())
	}
	return strings.Join(allProxy, "\n")
}

func (pool *ProxyPool) loadConf(confName string) (map[string]*Proxy, error) {
	proxys := make(map[string]*Proxy)
	confPath := pool.ProxyManager.config.confDir + "/" + confName

	txtFile, err := utils.NewTxtFile(confPath)
	if err != nil {
		log.Println("load proxy pool failed[", confName, "]")
		return proxys, err
	}
	return pool.loadProxysFromTxtFile(txtFile)
}

func (pool *ProxyPool) loadProxysFromTxtFile(txtFile *utils.TxtFile) (map[string]*Proxy, error) {
	proxys := make(map[string]*Proxy)
	defaultValues := make(map[string]string)
	defaultValues["proxy"] = "required"
	defaultValues["weight"] = "1"
	defaultValues["status"] = "1"
	defaultValues["last_check"] = "0"
	defaultValues["check_used"] = "0"

	datas, err := txtFile.KvMapSlice("=", true, defaultValues)
	if err != nil {
		return proxys, err
	}
	for _, kv := range datas {
		proxy := pool.parseProxy(kv)
		if proxy != nil {
			proxys[proxy.proxy] = proxy
		}
	}
	return proxys, nil
}

func (pool *ProxyPool) parseProxy(info map[string]string) *Proxy {
	if info == nil {
		return nil
	}
	proxy := NewProxy(info["proxy"])
	if proxy == nil {
		return nil
	}
	intValues := make(map[string]int)
	intFields := []string{"weight", "status", "check_used", "last_check", "last_check_ok"}
	var err error
	for _, fieldName := range intFields {
		intValues[fieldName], err = strconv.Atoi(info[fieldName])
		if err != nil {
			log.Println("parse [", fieldName, "]failed,not int.err:", err)
			intValues[fieldName] = 0
		}
	}
	proxy.Weight = intValues["weight"]
	proxy.StatusCode = PROXY_STATUS(intValues["status"])
	proxy.CheckUsed = int64(intValues["check_used"])
	proxy.LastCheck = int64(intValues["last_check"])
	proxy.LastCheckOk = int64(intValues["last_check_ok"])
	return proxy
}

func (pool *ProxyPool) GetProxy(proxy_url string) *Proxy {
	pool.mu.RLock()
	defer pool.mu.RUnlock()
	if proxy, has := pool.proxyListAll[proxy_url]; has {
		return proxy
	}
	return nil
}

func (pool *ProxyPool) addProxyActive(proxy_url string) bool {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	if proxy, has := pool.proxyListAll[proxy_url]; has {
		if _, hasAct := pool.proxyListActive[proxy_url]; !hasAct {
			pool.proxyListActive[proxy_url] = proxy
			return true
		}
	}
	return false
}

func (pool *ProxyPool) addProxy(proxy *Proxy) bool {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	if _, has := pool.proxyListAll[proxy.proxy]; !has {
		pool.proxyListAll[proxy.proxy] = proxy
		return true
	}
	return false
}

func (pool *ProxyPool) removeProxyActive(proxy_url string) {
	pool.mu.Lock()
	defer pool.mu.Unlock()
	if _, hasAct := pool.proxyListActive[proxy_url]; hasAct {
		delete(pool.proxyListActive, proxy_url)
	}
}

func (pool *ProxyPool) removeProxy(proxy_url string) {
	pool.mu.Lock()
	defer pool.mu.Unlock()
	if _, hasAct := pool.proxyListAll[proxy_url]; hasAct {
		delete(pool.proxyListAll, proxy_url)
	}
	if _, hasAct := pool.proxyListActive[proxy_url]; hasAct {
		delete(pool.proxyListActive, proxy_url)
	}
}

var errorNoProxy error = fmt.Errorf("no active proxy")

func (pool *ProxyPool) GetOneProxy(logid int64) (*Proxy, error) {
	pool.mu.Lock()
	defer pool.mu.Unlock()
	l := len(pool.proxyListActive)
	if l == 0 {
		return nil, errorNoProxy
	}

	sessionProxys, has := pool.SessionProxys[logid]

	if !has {
		sessionProxys = make(map[string]*Proxy)
		pool.SessionProxys[logid] = sessionProxys
	}

	for _, proxy := range pool.proxyListActive {
		if _, has := pool.proxyActiveUsed[proxy.proxy]; has {
			continue
		}
		if _, has := sessionProxys[proxy.proxy]; !has {
			sessionProxys[proxy.proxy] = proxy
			proxy.Used++
			pool.proxyActiveUsed[proxy.proxy] = "1"
			if len(pool.proxyActiveUsed) >= len(pool.proxyListActive) {
				pool.proxyActiveUsed = make(map[string]string)
			}
			return proxy, nil
		}
	}
	return nil, errorNoProxy
}

func (pool *ProxyPool) CleanSessionProxy(logid int64) {
	pool.mu.Lock()
	defer pool.mu.Unlock()
	if _, has := pool.SessionProxys[logid]; has {
		delete(pool.SessionProxys, logid)
	}
}
func (pool *ProxyPool) runTest() {
	pool.testRunChan <- true
	defer (func() {
		<-pool.testRunChan
	})()
	start := time.Now()
	proxyTotal := len(pool.proxyListAll)
	log.Println("start test all proxy,total=", proxyTotal)

	var wg sync.WaitGroup
	for name := range pool.proxyListAll {
		wg.Add(1)
		go (func(proxyUrl string) {
			pool.TestProxyAddActive(proxyUrl)
			wg.Done()
		})(name)
	}
	wg.Wait()

	used := time.Now().Sub(start)
	log.Println("test all proxy finish,total:", proxyTotal, "used:", used, "activeTotal:", len(pool.proxyListActive))

	pool.cleanBadProxy(86400)

	testResultFile := pool.ProxyManager.config.confDir + "/pool_checked.conf"
	utils.File_put_contents(testResultFile, []byte(pool.String()))
}

func (pool *ProxyPool) TestProxyAddActive(proxy_url string) bool {
	proxy := pool.GetProxy(proxy_url)
	if proxy == nil {
		return false
	}
	isOk := pool.TestProxy(proxy)
	if isOk {
		pool.addProxyActive(proxy.proxy)
	} else {
		pool.removeProxyActive(proxy.proxy)
	}
	return true
}

func (pool *ProxyPool) TestProxy(proxy *Proxy) bool {
	pool.checkChan <- proxy.proxy
	start := time.Now()
	defer (func() {
		<-pool.checkChan
	})()

	if start.Unix()-proxy.LastCheck < pool.checkInterval {
		return proxy.IsOk()
	}

	proxy.StatusCode = PROXY_STATUS_UNAVAILABLE

	testlog := func(msg ...interface{}) {
		used := time.Now().Sub(start)
		proxy.CheckUsed = used.Nanoseconds() / 1000000
		proxy.LastCheck = start.Unix()
		log.Println("test proxy", proxy.proxy, fmt.Sprint(msg...), "used:", proxy.CheckUsed, "ms")
	}

	if pool.aliveCheckUrl != "" {
		urlStr := strings.Replace(pool.aliveCheckUrl, "{%rand}", fmt.Sprintf("%d", start.UnixNano()), -1)
		resp, err := doRequestGet(urlStr, proxy, pool.timeout/2)
		if err != nil {
			testlog("failed,", err.Error())
			return false
		} else {
			cur_len := resp.Header.Get("Content-Length")
			check_len := pool.aliveCheckResponse.Header.Get("Content-Length")
			if cur_len != check_len {
				testlog("failed ,content-length wrong,[", check_len, "!=", cur_len, "]")
				return false
			}
		}
	} else {
		host, port, err := utils.Net_getHostPortFromUrl(proxy.proxy)
		if err != nil {
			testlog("failed,proxy url err:", err)
			return false
		}
		conn, netErr := net.Dial("tcp", fmt.Sprintf("%s:%d", host, port))
		if netErr != nil {
			testlog("failed", netErr)
			return false
		}
		conn.Close()
	}
	proxy.StatusCode = PROXY_STATUS_ACTIVE
	proxy.LastCheckOk = proxy.LastCheck
	testlog("pass")
	return true
}

func (pool *ProxyPool) MarkProxyStatus(proxy *Proxy, status PROXY_USED_STATUS) {
	proxy.Count.MarkStatus(status)
	pool.Count.MarkStatus(status)
}

func doRequestGet(urlStr string, proxy *Proxy, timeout_sec int) (resp *http.Response, err error) {
	client := &http.Client{}
	if timeout_sec > 0 {
		client.Timeout = time.Duration(timeout_sec) * time.Second
	}
	if proxy != nil {
		proxyGetFn := func(req *http.Request) (*url.URL, error) {
			return proxy.URL, nil
		}
		client.Transport = &http.Transport{Proxy: proxyGetFn}
	}
	req, _ := http.NewRequest("GET", urlStr, nil)
	return client.Do(req)
}

func (pool *ProxyPool)cleanBadProxy(sec int64){
	last:=time.Now().Unix()-sec
	proxyBad:=make([]*Proxy,0)
	for _,proxy:= range pool.proxyListAll {
		if proxy.LastCheckOk<last{
			proxyBad=append(proxyBad,proxy)
		}
	}

	for _,proxy:=range proxyBad{
		pool.removeProxy(proxy.proxy)
		utils.File_put_contents(pool.ProxyManager.config.confDir+"/pool_bad.list", []byte(proxy.String()+"\n"),utils.FILE_APPEND)
	}
}
<|endoftext|>


<commit_before>package mat64

import (
	"math/rand"

	"github.com/gonum/blas/blas64"
	"github.com/gonum/floats"
	"gopkg.in/check.v1"
)

func (s *S) TestNewVector(c *check.C) {
	for i, test := range []struct {
		n      int
		data   []float64
		vector *Vector
	}{
		{
			n:    3,
			data: []float64{4, 5, 6},
			vector: &Vector{
				mat: blas64.Vector{
					Data: []float64{4, 5, 6},
					Inc:  1,
				},
				n: 3,
			},
		},
	} {
		v := NewVector(test.n, test.data)
		rows, cols := v.Dims()
		c.Check(rows, check.Equals, test.n, check.Commentf("Test %d", i))
		c.Check(cols, check.Equals, 1, check.Commentf("Test %d", i))
		c.Check(v, check.DeepEquals, test.vector, check.Commentf("Test %d", i))
		v2 := NewVector(test.n, nil)
		c.Check(v2.mat.Data, check.DeepEquals, []float64{0, 0, 0}, check.Commentf("Test %d", i))
	}
}

func (s *S) TestVectorAtSet(c *check.C) {
	for i, test := range []struct {
		vector *Vector
	}{
		{
			vector: &Vector{
				mat: blas64.Vector{
					Data: []float64{0, 1, 2},
					Inc:  1,
				},
				n: 3,
			},
		},
		{
			vector: &Vector{
				mat: blas64.Vector{
					Data: []float64{0, 10, 10, 1, 10, 10, 2},
					Inc:  3,
				},
				n: 3,
			},
		},
	} {
		v := test.vector
		n := test.vector.n
		c.Check(func() { v.At(n, 0) }, check.PanicMatches, ErrRowAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.At(-1, 0) }, check.PanicMatches, ErrRowAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.At(0, 1) }, check.PanicMatches, ErrColAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.At(0, -1) }, check.PanicMatches, ErrColAccess.Error(), check.Commentf("Test %d", i))

		c.Check(v.At(0, 0), check.Equals, 0.0, check.Commentf("Test %d", i))
		c.Check(v.At(1, 0), check.Equals, 1.0, check.Commentf("Test %d", i))
		c.Check(v.At(n-1, 0), check.Equals, float64(n-1), check.Commentf("Test %d", i))

		c.Check(func() { v.Set(n, 0, 100) }, check.PanicMatches, ErrRowAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.Set(-1, 0, 100) }, check.PanicMatches, ErrRowAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.Set(0, 1, 100) }, check.PanicMatches, ErrColAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.Set(0, -1, 100) }, check.PanicMatches, ErrColAccess.Error(), check.Commentf("Test %d", i))

		v.Set(0, 0, 100)
		c.Check(v.At(0, 0), check.Equals, 100.0, check.Commentf("Test %d", i))
		v.Set(2, 0, 101)
		c.Check(v.At(2, 0), check.Equals, 101.0, check.Commentf("Test %d", i))
	}
}

func (s *S) TestVectorMul(c *check.C) {

	for i, test := range []struct {
		m int
		n int
	}{
		{
			m: 10,
			n: 5,
		},
		/*
			{
				m: 5,
				n: 5,
			},
			{
				m: 5,
				n: 10,
			},
		*/
	} {
		vData := make([]float64, test.n)
		for i := range vData {
			vData[i] = rand.Float64()
		}
		vDataCopy := make([]float64, test.n)
		copy(vDataCopy, vData)
		v := NewVector(test.n, vData)
		aData := make([]float64, test.n*test.m)
		for i := range aData {
			aData[i] = rand.Float64()
		}
		a := NewDense(test.m, test.n, aData)
		var v2 Vector
		v2.MulVec(a, false, v)
		var v2M Dense
		v2M.Mul(a, v)
		same := floats.EqualApprox(v2.mat.Data, v2M.mat.Data, 1e-14)
		c.Check(same, check.Equals, true, check.Commentf("Test %d", i))

		var aT Dense
		aT.TCopy(a)
		v2.MulVec(&aT, true, v)
		same = floats.EqualApprox(v2.mat.Data, v2M.mat.Data, 1e-14)
		c.Check(same, check.Equals, true, check.Commentf("Test %d", i))

		v.MulVec(&aT, true, v)
		same = floats.EqualApprox(v.mat.Data, v2M.mat.Data, 1e-14)
		c.Check(same, check.Equals, true, check.Commentf("Test %d", i))
	}
}
<commit_msg>Uncommented valid tests and commented invalid tests<commit_after>package mat64

import (
	"math/rand"

	"github.com/gonum/blas/blas64"
	"github.com/gonum/floats"
	"gopkg.in/check.v1"
)

func (s *S) TestNewVector(c *check.C) {
	for i, test := range []struct {
		n      int
		data   []float64
		vector *Vector
	}{
		{
			n:    3,
			data: []float64{4, 5, 6},
			vector: &Vector{
				mat: blas64.Vector{
					Data: []float64{4, 5, 6},
					Inc:  1,
				},
				n: 3,
			},
		},
	} {
		v := NewVector(test.n, test.data)
		rows, cols := v.Dims()
		c.Check(rows, check.Equals, test.n, check.Commentf("Test %d", i))
		c.Check(cols, check.Equals, 1, check.Commentf("Test %d", i))
		c.Check(v, check.DeepEquals, test.vector, check.Commentf("Test %d", i))
		v2 := NewVector(test.n, nil)
		c.Check(v2.mat.Data, check.DeepEquals, []float64{0, 0, 0}, check.Commentf("Test %d", i))
	}
}

func (s *S) TestVectorAtSet(c *check.C) {
	for i, test := range []struct {
		vector *Vector
	}{
		{
			vector: &Vector{
				mat: blas64.Vector{
					Data: []float64{0, 1, 2},
					Inc:  1,
				},
				n: 3,
			},
		},
		{
			vector: &Vector{
				mat: blas64.Vector{
					Data: []float64{0, 10, 10, 1, 10, 10, 2},
					Inc:  3,
				},
				n: 3,
			},
		},
	} {
		v := test.vector
		n := test.vector.n
		c.Check(func() { v.At(n, 0) }, check.PanicMatches, ErrRowAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.At(-1, 0) }, check.PanicMatches, ErrRowAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.At(0, 1) }, check.PanicMatches, ErrColAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.At(0, -1) }, check.PanicMatches, ErrColAccess.Error(), check.Commentf("Test %d", i))

		c.Check(v.At(0, 0), check.Equals, 0.0, check.Commentf("Test %d", i))
		c.Check(v.At(1, 0), check.Equals, 1.0, check.Commentf("Test %d", i))
		c.Check(v.At(n-1, 0), check.Equals, float64(n-1), check.Commentf("Test %d", i))

		c.Check(func() { v.Set(n, 0, 100) }, check.PanicMatches, ErrRowAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.Set(-1, 0, 100) }, check.PanicMatches, ErrRowAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.Set(0, 1, 100) }, check.PanicMatches, ErrColAccess.Error(), check.Commentf("Test %d", i))
		c.Check(func() { v.Set(0, -1, 100) }, check.PanicMatches, ErrColAccess.Error(), check.Commentf("Test %d", i))

		v.Set(0, 0, 100)
		c.Check(v.At(0, 0), check.Equals, 100.0, check.Commentf("Test %d", i))
		v.Set(2, 0, 101)
		c.Check(v.At(2, 0), check.Equals, 101.0, check.Commentf("Test %d", i))
	}
}

func (s *S) TestVectorMul(c *check.C) {

	for i, test := range []struct {
		m int
		n int
	}{
		{
			m: 10,
			n: 5,
		},
		{
			m: 5,
			n: 5,
		},
		{
			m: 5,
			n: 10,
		},
	} {
		vData := make([]float64, test.n)
		for i := range vData {
			vData[i] = rand.Float64()
		}
		vDataCopy := make([]float64, test.n)
		copy(vDataCopy, vData)
		v := NewVector(test.n, vData)
		aData := make([]float64, test.n*test.m)
		for i := range aData {
			aData[i] = rand.Float64()
		}
		a := NewDense(test.m, test.n, aData)
		var v2 Vector
		v2.MulVec(a, false, v)
		var v2M Dense
		v2M.Mul(a, v)
		same := floats.EqualApprox(v2.mat.Data, v2M.mat.Data, 1e-14)
		c.Check(same, check.Equals, true, check.Commentf("Test %d", i))

		var aT Dense
		aT.TCopy(a)
		v2.MulVec(&aT, true, v)
		same = floats.EqualApprox(v2.mat.Data, v2M.mat.Data, 1e-14)
		c.Check(same, check.Equals, true, check.Commentf("Test %d", i))

		/*
			v.MulVec(&aT, true, v)
			same = floats.EqualApprox(v.mat.Data, v2M.mat.Data, 1e-14)
			c.Check(same, check.Equals, true, check.Commentf("Test %d", i))
		*/
	}
}
<|endoftext|>


<commit_before>package goose

import (
	"fmt"
	log "github.com/getwe/goose/log"
	flags "github.com/jessevdk/go-flags"
	"os"
	"time"
)

// goose的入口程序.
type Goose struct {
	// 建库策略
	indexSty IndexStrategy
	// 检索策略
	searchSty SearchStrategy

	// 配置文件
	confPath string

	// 日志配置
	logConfPath string

	// 建库模式数据文件
	dataPath string
}

func (this *Goose) SetIndexStrategy(sty IndexStrategy) {
	this.indexSty = sty
}

func (this *Goose) SetSearchStrategy(sty SearchStrategy) {
	this.searchSty = sty
}

// 程序入口,解析程序参数,启动[建库|检索]模式
func (this *Goose) Run() {
	defer func() {
		if r := recover(); r != nil {
			os.Exit(1)
		}
	}()

	// 解析命令行参数
	var opts struct {
		// build mode
		BuildMode bool `short:"b" long:"build" description:"run in build mode"`

		// configure file
		Configure string `short:"c" long:"conf" description:"congfigure file" default:"conf/goose.toml"`

		// log configure file
		LogConf string `short:"l" long:"logconf" description:"log congfigure file" default:"conf/log.toml"`

		// build mode data file
		DataFile string `short:"d" long:"datafile" description:"build mode data file"`
	}
	parser := flags.NewParser(&opts, flags.HelpFlag)
	_, err := parser.ParseArgs(os.Args)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	if opts.BuildMode && len(opts.DataFile) == 0 {
		parser.WriteHelp(os.Stderr)
		os.Exit(1)
	}

	this.confPath = opts.Configure
	this.dataPath = opts.DataFile
	this.logConfPath = opts.LogConf

	// init log
	err = log.LoadConfiguration(this.logConfPath)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	log.Debug("Load log conf finish")

	// run
	if opts.BuildMode {
		this.buildModeRun()
	} else {
		this.searchModeRun()
	}

	// BUG(log4go) log4go need time to sync ...(wtf)
	// see http://stackoverflow.com/questions/14252766/abnormal-behavior-of-log4go
	time.Sleep(100 * time.Millisecond)
}

// 建库模式运行
func (this *Goose) buildModeRun() {

	if this.indexSty == nil {
		log.Error("Please set index strategy,see Goose.SetIndexStrategy()")
		return
	}

	gooseBuild := NewGooseBuild()
	err := gooseBuild.Init(this.confPath, this.indexSty, this.dataPath)
	if err != nil {
		fmt.Println(err)
		log.Error(err)
		return
	}

	err = gooseBuild.Run()
	if err != nil {
		log.Error(err)
		return
	}

}

// 检索模式运行
func (this *Goose) searchModeRun() {

	log.Debug("run in search mode")

	if this.searchSty == nil {
		log.Error("Please set search strategy,see Goose.SetSearchStrategy()")
		return
	}

	if this.indexSty == nil {
		log.Warn("can't build index real time witout Index Strategy")
	}

	gooseSearch := NewGooseSearch()
	err := gooseSearch.Init(this.confPath, this.indexSty, this.searchSty)
	if err != nil {
		log.Error(err)
		return
	}

	log.Debug("goose search init succ")

	err = gooseSearch.Run()
	if err != nil {
		log.Error(err)
		return
	}
}

func NewGoose() *Goose {
	g := Goose{}
	g.indexSty = nil
	g.searchSty = nil
	return &g
}
<commit_msg>add logo<commit_after>package goose

import (
	"fmt"
	"github.com/fatih/color"
	"github.com/getwe/figlet4go"
	log "github.com/getwe/goose/log"
	flags "github.com/jessevdk/go-flags"
	"os"
	"time"
)

// goose的入口程序.
type Goose struct {
	// 建库策略
	indexSty IndexStrategy
	// 检索策略
	searchSty SearchStrategy

	// 配置文件
	confPath string

	// 日志配置
	logConfPath string

	// 建库模式数据文件
	dataPath string
}

func (this *Goose) SetIndexStrategy(sty IndexStrategy) {
	this.indexSty = sty
}

func (this *Goose) SetSearchStrategy(sty SearchStrategy) {
	this.searchSty = sty
}

// 程序入口,解析程序参数,启动[建库|检索]模式
func (this *Goose) Run() {
	defer func() {
		if r := recover(); r != nil {
			os.Exit(1)
		}
	}()

	// 解析命令行参数
	var opts struct {
		// build mode
		BuildMode bool `short:"b" long:"build" description:"run in build mode"`

		// configure file
		Configure string `short:"c" long:"conf" description:"congfigure file" default:"conf/goose.toml"`

		// log configure file
		LogConf string `short:"l" long:"logconf" description:"log congfigure file" default:"conf/log.toml"`

		// build mode data file
		DataFile string `short:"d" long:"datafile" description:"build mode data file"`
	}
	parser := flags.NewParser(&opts, flags.HelpFlag)
	_, err := parser.ParseArgs(os.Args)
	if err != nil {
		fmt.Println(this.showLogo())
		fmt.Println(err)
		os.Exit(1)
	}
	if opts.BuildMode && len(opts.DataFile) == 0 {
		fmt.Println(this.showLogo())
		parser.WriteHelp(os.Stderr)
		os.Exit(1)
	}

	this.confPath = opts.Configure
	this.dataPath = opts.DataFile
	this.logConfPath = opts.LogConf

	// init log
	err = log.LoadConfiguration(this.logConfPath)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	log.Debug("Load log conf finish")

	// run
	if opts.BuildMode {
		this.buildModeRun()
	} else {
		this.searchModeRun()
	}

	// BUG(log4go) log4go need time to sync ...(wtf)
	// see http://stackoverflow.com/questions/14252766/abnormal-behavior-of-log4go
	time.Sleep(100 * time.Millisecond)
}

func (this *Goose) showLogo() string {
	str := "goose"
	ascii := figlet4go.NewAsciiRender()

	// change the font color
	options := figlet4go.NewRenderOptions()
	options.FontColor = make([]color.Attribute, len(str))
	options.FontColor[0] = color.FgMagenta
	options.FontColor[1] = color.FgYellow
	options.FontColor[2] = color.FgBlue
	options.FontColor[3] = color.FgCyan
	options.FontColor[4] = color.FgRed
	renderStr, _ := ascii.RenderOpts(str, options)
	return renderStr
}

// 建库模式运行
func (this *Goose) buildModeRun() {

	if this.indexSty == nil {
		log.Error("Please set index strategy,see Goose.SetIndexStrategy()")
		return
	}

	gooseBuild := NewGooseBuild()
	err := gooseBuild.Init(this.confPath, this.indexSty, this.dataPath)
	if err != nil {
		fmt.Println(err)
		log.Error(err)
		return
	}

	err = gooseBuild.Run()
	if err != nil {
		log.Error(err)
		return
	}

}

// 检索模式运行
func (this *Goose) searchModeRun() {

	log.Debug("run in search mode")

	if this.searchSty == nil {
		log.Error("Please set search strategy,see Goose.SetSearchStrategy()")
		return
	}

	if this.indexSty == nil {
		log.Warn("can't build index real time witout Index Strategy")
	}

	gooseSearch := NewGooseSearch()
	err := gooseSearch.Init(this.confPath, this.indexSty, this.searchSty)
	if err != nil {
		log.Error(err)
		return
	}

	log.Debug("goose search init succ")

	err = gooseSearch.Run()
	if err != nil {
		log.Error(err)
		return
	}
}

func NewGoose() *Goose {
	g := Goose{}
	g.indexSty = nil
	g.searchSty = nil
	return &g
}
<|endoftext|>


<commit_before>// +build linux

/*
 * Copyright (C) 2016 Red Hat, 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 ofthe 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 specificlanguage governing permissions and
 * limitations under the License.
 *
 */

package docker

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/events"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/client"
	"github.com/vishvananda/netns"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/graffiti/graph"
	"github.com/skydive-project/skydive/probe"
	"github.com/skydive-project/skydive/topology"
	tp "github.com/skydive-project/skydive/topology/probes"
	ns "github.com/skydive-project/skydive/topology/probes/netns"
	sversion "github.com/skydive-project/skydive/version"
)

// ClientAPIVersion Client API version used
const ClientAPIVersion = "1.18"

type containerInfo struct {
	Pid  int
	Node *graph.Node
}

// ProbeHandler describes a Docker topology graph that enhance the graph
type ProbeHandler struct {
	common.RWMutex
	*ns.ProbeHandler
	url          string
	client       *client.Client
	cancel       context.CancelFunc
	state        int64
	connected    atomic.Value
	wg           sync.WaitGroup
	hostNs       netns.NsHandle
	containerMap map[string]containerInfo
}

func (p *ProbeHandler) containerNamespace(pid int) string {
	return fmt.Sprintf("/proc/%d/ns/net", pid)
}

func (p *ProbeHandler) registerContainer(id string) {
	p.Lock()
	defer p.Unlock()

	if _, ok := p.containerMap[id]; ok {
		return
	}
	info, err := p.client.ContainerInspect(context.Background(), id)
	if err != nil {
		p.Ctx.Logger.Errorf("Failed to inspect Docker container %s: %s", id, err)
		return
	}

	nsHandle, err := netns.GetFromPid(info.State.Pid)
	if err != nil {
		return
	}
	defer nsHandle.Close()

	namespace := p.containerNamespace(info.State.Pid)
	p.Ctx.Logger.Debugf("Register docker container %s and PID %d", info.ID, info.State.Pid)

	var n *graph.Node
	if p.hostNs.Equal(nsHandle) {
		// The container is in net=host mode
		n = p.Ctx.RootNode
	} else {
		if n, err = p.Register(namespace, info.Name[1:]); err != nil {
			p.Ctx.Logger.Debugf("Failed to register probe for namespace %s: %s", namespace, err)
			return
		}

		p.Ctx.Graph.Lock()
		if err := p.Ctx.Graph.AddMetadata(n, "Manager", "docker"); err != nil {
			p.Ctx.Logger.Error(err)
		}
		p.Ctx.Graph.Unlock()
	}

	pid := int64(info.State.Pid)

	dockerMetadata := Metadata{
		ContainerID:   info.ID,
		ContainerName: info.Name[1:],
	}

	if len(info.Config.Labels) != 0 {
		dockerMetadata.Labels = graph.Metadata(common.NormalizeValue(info.Config.Labels).(map[string]interface{}))
	}

	p.Ctx.Graph.Lock()
	defer p.Ctx.Graph.Unlock()

	containerNode := p.Ctx.Graph.LookupFirstNode(graph.Metadata{"InitProcessPID": pid})
	if containerNode != nil {
		if err := p.Ctx.Graph.AddMetadata(containerNode, "Docker", dockerMetadata); err != nil {
			p.Ctx.Logger.Error(err)
		}
	} else {
		metadata := graph.Metadata{
			"Type":           "container",
			"Name":           info.Name[1:],
			"Manager":        "docker",
			"InitProcessPID": pid,
			"Docker":         dockerMetadata,
		}

		if containerNode, err = p.Ctx.Graph.NewNode(graph.GenID(), metadata); err != nil {
			p.Ctx.Logger.Error(err)
			return
		}
	}
	topology.AddOwnershipLink(p.Ctx.Graph, n, containerNode, nil)

	p.containerMap[info.ID] = containerInfo{
		Pid:  info.State.Pid,
		Node: containerNode,
	}
}

func (p *ProbeHandler) unregisterContainer(id string) {
	p.Lock()
	defer p.Unlock()

	infos, ok := p.containerMap[id]
	if !ok {
		return
	}

	p.Ctx.Graph.Lock()
	if err := p.Ctx.Graph.DelNode(infos.Node); err != nil {
		p.Ctx.Graph.Unlock()
		p.Ctx.Logger.Error(err)
		return
	}
	p.Ctx.Graph.Unlock()

	namespace := p.containerNamespace(infos.Pid)
	p.Ctx.Logger.Debugf("Stop listening for namespace %s with PID %d", namespace, infos.Pid)
	p.Unregister(namespace)

	delete(p.containerMap, id)
}

func (p *ProbeHandler) handleDockerEvent(event *events.Message) {
	if event.Status == "start" {
		p.registerContainer(event.ID)
	} else if event.Status == "die" {
		p.unregisterContainer(event.ID)
	}
}

func (p *ProbeHandler) connect() error {
	var err error

	p.Ctx.Logger.Debugf("Connecting to Docker daemon: %s", p.url)
	defaultHeaders := map[string]string{"User-Agent": fmt.Sprintf("skydive-agent-%s", sversion.Version)}
	p.client, err = client.NewClient(p.url, ClientAPIVersion, nil, defaultHeaders)
	if err != nil {
		p.Ctx.Logger.Errorf("Failed to create client to Docker daemon: %s", err)
		return err
	}
	defer p.client.Close()

	if _, err := p.client.ServerVersion(context.Background()); err != nil {
		p.Ctx.Logger.Errorf("Failed to connect to Docker daemon: %s", err)
		return err
	}

	if p.hostNs, err = netns.Get(); err != nil {
		return err
	}
	defer p.hostNs.Close()

	for id := range p.containerMap {
		p.unregisterContainer(id)
	}

	eventsFilter := filters.NewArgs()
	eventsFilter.Add("event", "start")
	eventsFilter.Add("event", "die")

	ctx, cancel := context.WithCancel(context.Background())
	eventChan, errChan := p.client.Events(ctx, types.EventsOptions{Filters: eventsFilter})

	p.cancel = cancel
	p.wg.Add(2)

	p.connected.Store(true)
	defer p.connected.Store(false)

	go func() {
		defer p.wg.Done()

		containers, err := p.client.ContainerList(ctx, types.ContainerListOptions{})
		if err != nil {
			p.Ctx.Logger.Errorf("Failed to list containers: %s", err)
			return
		}

		for _, c := range containers {
			if atomic.LoadInt64(&p.state) != common.RunningState {
				break
			}
			p.registerContainer(c.ID)
		}
	}()

	defer p.wg.Done()

	for {
		select {
		case err := <-errChan:
			if atomic.LoadInt64(&p.state) != common.StoppingState {
				err = fmt.Errorf("Got error while waiting for Docker event: %s", err)
			}
			return err
		case event := <-eventChan:
			p.handleDockerEvent(&event)
		}
	}
}

// Start the probe
func (p *ProbeHandler) Start() {
	if !atomic.CompareAndSwapInt64(&p.state, common.StoppedState, common.RunningState) {
		return
	}

	go func() {
		for {
			state := atomic.LoadInt64(&p.state)
			if state == common.StoppingState || state == common.StoppedState {
				break
			}

			if p.connect() != nil {
				time.Sleep(1 * time.Second)
			}

			p.wg.Wait()
		}
	}()
}

// Stop the probe
func (p *ProbeHandler) Stop() {
	if !atomic.CompareAndSwapInt64(&p.state, common.RunningState, common.StoppingState) {
		return
	}

	if p.connected.Load() == true {
		p.cancel()
		p.wg.Wait()
	}

	atomic.StoreInt64(&p.state, common.StoppedState)
}

// Init initializes a new topology Docker probe
func (p *ProbeHandler) Init(ctx tp.Context, bundle *probe.Bundle) (probe.Handler, error) {
	nsHandler := bundle.GetHandler("netns")
	if nsHandler == nil {
		return nil, errors.New("unable to find the netns handler")
	}

	dockerURL := ctx.Config.GetString("agent.topology.docker.url")
	netnsRunPath := ctx.Config.GetString("agent.topology.docker.netns.run_path")

	p.ProbeHandler = nsHandler.(*ns.ProbeHandler)
	p.url = dockerURL
	p.containerMap = make(map[string]containerInfo)
	p.state = common.StoppedState

	if netnsRunPath != "" {
		p.Exclude(netnsRunPath + "/default")
		p.Watch(netnsRunPath)
	}

	return p, nil
}
<commit_msg>docker: make use of retry probe wrapper<commit_after>// +build linux

/*
 * Copyright (C) 2016 Red Hat, 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 ofthe 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 specificlanguage governing permissions and
 * limitations under the License.
 *
 */

package docker

import (
	"context"
	"errors"
	"fmt"
	"sync"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/events"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/client"
	"github.com/vishvananda/netns"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/graffiti/graph"
	"github.com/skydive-project/skydive/probe"
	"github.com/skydive-project/skydive/topology"
	"github.com/skydive-project/skydive/topology/probes"
	tp "github.com/skydive-project/skydive/topology/probes"
	ns "github.com/skydive-project/skydive/topology/probes/netns"
	sversion "github.com/skydive-project/skydive/version"
)

// ClientAPIVersion Client API version used
const ClientAPIVersion = "1.18"

type containerInfo struct {
	Pid  int
	Node *graph.Node
}

// ProbeHandler describes a Docker topology graph that enhance the graph
type ProbeHandler struct {
	common.RWMutex
	Ctx          tp.Context
	nsProbe      *ns.ProbeHandler
	url          string
	client       *client.Client
	hostNs       netns.NsHandle
	containerMap map[string]containerInfo
}

func (p *ProbeHandler) containerNamespace(pid int) string {
	return fmt.Sprintf("/proc/%d/ns/net", pid)
}

func (p *ProbeHandler) registerContainer(id string) {
	p.Lock()
	defer p.Unlock()

	if _, ok := p.containerMap[id]; ok {
		return
	}
	info, err := p.client.ContainerInspect(context.Background(), id)
	if err != nil {
		p.Ctx.Logger.Errorf("Failed to inspect Docker container %s: %s", id, err)
		return
	}

	nsHandle, err := netns.GetFromPid(info.State.Pid)
	if err != nil {
		return
	}
	defer nsHandle.Close()

	namespace := p.containerNamespace(info.State.Pid)
	p.Ctx.Logger.Debugf("Register docker container %s and PID %d", info.ID, info.State.Pid)

	var n *graph.Node
	if p.hostNs.Equal(nsHandle) {
		// The container is in net=host mode
		n = p.Ctx.RootNode
	} else {
		if n, err = p.nsProbe.Register(namespace, info.Name[1:]); err != nil {
			p.Ctx.Logger.Debugf("Failed to register probe for namespace %s: %s", namespace, err)
			return
		}

		p.Ctx.Graph.Lock()
		if err := p.Ctx.Graph.AddMetadata(n, "Manager", "docker"); err != nil {
			p.Ctx.Logger.Error(err)
		}
		p.Ctx.Graph.Unlock()
	}

	pid := int64(info.State.Pid)

	dockerMetadata := Metadata{
		ContainerID:   info.ID,
		ContainerName: info.Name[1:],
	}

	if len(info.Config.Labels) != 0 {
		dockerMetadata.Labels = graph.Metadata(common.NormalizeValue(info.Config.Labels).(map[string]interface{}))
	}

	p.Ctx.Graph.Lock()
	defer p.Ctx.Graph.Unlock()

	containerNode := p.Ctx.Graph.LookupFirstNode(graph.Metadata{"InitProcessPID": pid})
	if containerNode != nil {
		if err := p.Ctx.Graph.AddMetadata(containerNode, "Docker", dockerMetadata); err != nil {
			p.Ctx.Logger.Error(err)
		}
	} else {
		metadata := graph.Metadata{
			"Type":           "container",
			"Name":           info.Name[1:],
			"Manager":        "docker",
			"InitProcessPID": pid,
			"Docker":         dockerMetadata,
		}

		if containerNode, err = p.Ctx.Graph.NewNode(graph.GenID(), metadata); err != nil {
			p.Ctx.Logger.Error(err)
			return
		}
	}
	topology.AddOwnershipLink(p.Ctx.Graph, n, containerNode, nil)

	p.containerMap[info.ID] = containerInfo{
		Pid:  info.State.Pid,
		Node: containerNode,
	}
}

func (p *ProbeHandler) unregisterContainer(id string) {
	p.Lock()
	defer p.Unlock()

	infos, ok := p.containerMap[id]
	if !ok {
		return
	}

	p.Ctx.Graph.Lock()
	if err := p.Ctx.Graph.DelNode(infos.Node); err != nil {
		p.Ctx.Graph.Unlock()
		p.Ctx.Logger.Error(err)
		return
	}
	p.Ctx.Graph.Unlock()

	namespace := p.containerNamespace(infos.Pid)
	p.Ctx.Logger.Debugf("Stop listening for namespace %s with PID %d", namespace, infos.Pid)
	p.nsProbe.Unregister(namespace)

	delete(p.containerMap, id)
}

func (p *ProbeHandler) handleDockerEvent(event *events.Message) {
	if event.Status == "start" {
		p.registerContainer(event.ID)
	} else if event.Status == "die" {
		p.unregisterContainer(event.ID)
	}
}

// Do connects to the Docker daemon, registers the existing containers and
// start listening for Docker events
func (p *ProbeHandler) Do(ctx context.Context, wg *sync.WaitGroup) error {
	var err error

	p.Ctx.Logger.Debugf("Connecting to Docker daemon: %s", p.url)
	defaultHeaders := map[string]string{"User-Agent": fmt.Sprintf("skydive-agent-%s", sversion.Version)}
	p.client, err = client.NewClient(p.url, ClientAPIVersion, nil, defaultHeaders)
	if err != nil {
		p.Ctx.Logger.Errorf("Failed to create client to Docker daemon: %s", err)
		return err
	}
	defer p.client.Close()

	version, err := p.client.ServerVersion(ctx)
	if err != nil {
		p.Ctx.Logger.Errorf("Failed to connect to Docker daemon: %s", err)
		return err
	}

	p.Ctx.Logger.Infof("Connected to Docker %s", version.Version)

	if p.hostNs, err = netns.Get(); err != nil {
		return err
	}
	defer p.hostNs.Close()

	for id := range p.containerMap {
		p.unregisterContainer(id)
	}

	eventsFilter := filters.NewArgs()
	eventsFilter.Add("event", "start")
	eventsFilter.Add("event", "die")

	eventChan, errChan := p.client.Events(ctx, types.EventsOptions{Filters: eventsFilter})

	wg.Add(2)

	go func() {
		defer wg.Done()

		containers, err := p.client.ContainerList(ctx, types.ContainerListOptions{})
		if err != nil {
			p.Ctx.Logger.Errorf("Failed to list containers: %s", err)
			return
		}

		for _, c := range containers {
			select {
			case <-ctx.Done():
			default:
				p.registerContainer(c.ID)
			}
		}
	}()

	go func() {
		defer wg.Done()

		for {
			select {
			case err := <-errChan:
				switch {
				case err == nil || err == context.Canceled:
					return
				case err.Error() == "unexpected EOF":
					p.Ctx.Logger.Error("lost connection to Docker")
				default:
					p.Ctx.Logger.Errorf("got error while waiting for Docker event: %s", err)
				}
			case event := <-eventChan:
				p.handleDockerEvent(&event)
			}
		}
	}()

	return nil
}

// Init initializes a new topology Docker probe
func (p *ProbeHandler) Init(ctx tp.Context, bundle *probe.Bundle) (probe.Handler, error) {
	nsHandler := bundle.GetHandler("netns")
	if nsHandler == nil {
		return nil, errors.New("unable to find the netns handler")
	}

	dockerURL := ctx.Config.GetString("agent.topology.docker.url")
	netnsRunPath := ctx.Config.GetString("agent.topology.docker.netns.run_path")

	p.nsProbe = nsHandler.(*ns.ProbeHandler)
	p.url = dockerURL
	p.containerMap = make(map[string]containerInfo)
	p.Ctx = ctx

	if netnsRunPath != "" {
		p.nsProbe.Exclude(netnsRunPath + "/default")
		p.nsProbe.Watch(netnsRunPath)
	}

	return probes.NewProbeWrapper(p), nil
}
<|endoftext|>


<commit_before>package moby

import (
	"bytes"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"

	log "github.com/Sirupsen/logrus"
)

var linuxkitYaml = map[string]string{"mkimage": `
kernel:
  image: "linuxkit/kernel:4.9.x"
  cmdline: "console=ttyS0"
init:
  - linuxkit/init:14a38303ee9dcb4541c00e2b87404befc1ba2083
  - linuxkit/runc:a0f2894e50bacbd1ff82be41edff8b8e06e0b161
  - linuxkit/containerd:389e67c3c1fc009c1315f32b3e2b6659691a3ad4
onboot:
  - name: mkimage
    image: "linuxkit/mkimage:5ad60299be03008f29c5caec3c5ea4ac0387aae6"
  - name: poweroff
    image: "linuxkit/poweroff:a8f1e4ad8d459f1fdaad9e4b007512cb3b504ae8"
trust:
  org:
    - linuxkit
`}

func imageFilename(name string) string {
	yaml := linuxkitYaml[name]
	hash := sha256.Sum256([]byte(yaml))
	return filepath.Join(MobyDir, "linuxkit", name+"-"+fmt.Sprintf("%x", hash))
}

func ensureLinuxkitImage(name string) error {
	filename := imageFilename(name)
	_, err1 := os.Stat(filename + "-kernel")
	_, err2 := os.Stat(filename + "-initrd.img")
	_, err3 := os.Stat(filename + "-cmdline")
	if err1 == nil && err2 == nil && err3 == nil {
		return nil
	}
	err := os.MkdirAll(filepath.Join(MobyDir, "linuxkit"), 0755)
	if err != nil {
		return err
	}
	// TODO clean up old files
	log.Infof("Building LinuxKit image %s to generate output formats", name)

	yaml := linuxkitYaml[name]

	m, err := NewConfig([]byte(yaml))
	if err != nil {
		return err
	}
	// TODO pass through --pull to here
	buf := new(bytes.Buffer)
	Build(m, buf, false, "")
	image := buf.Bytes()
	kernel, initrd, cmdline, err := tarToInitrd(image)
	if err != nil {
		return fmt.Errorf("Error converting to initrd: %v", err)
	}
	err = writeKernelInitrd(filename, kernel, initrd, cmdline)
	if err != nil {
		return err
	}

	return nil
}

func writeKernelInitrd(filename string, kernel []byte, initrd []byte, cmdline string) error {
	err := ioutil.WriteFile(filename+"-kernel", kernel, 0600)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(filename+"-initrd.img", initrd, 0600)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(filename+"-cmdline", []byte(cmdline), 0600)
	if err != nil {
		return err
	}
	return nil
}

func outputLinuxKit(format string, filename string, kernel []byte, initrd []byte, cmdline string, size int, hyperkit bool) error {
	log.Debugf("output linuxkit generated img: %s %s size %d", format, filename, size)

	tmp, err := ioutil.TempDir(filepath.Join(MobyDir, "tmp"), "moby")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmp)

	buf, err := tarInitrdKernel(kernel, initrd, cmdline)
	if err != nil {
		return err
	}

	tardisk := filepath.Join(tmp, "tardisk")
	f, err := os.Create(tardisk)
	if err != nil {
		return err
	}
	_, err = io.Copy(f, buf)
	if err != nil {
		return err
	}
	err = f.Close()
	if err != nil {
		return err
	}

	sizeString := fmt.Sprintf("%dM", size)
	_ = os.Remove(filename)
	_, err = os.Stat(filename)
	if err == nil || !os.IsNotExist(err) {
		return fmt.Errorf("Cannot remove existing file [%s]", filename)
	}
	linuxkit, err := exec.LookPath("linuxkit")
	if err != nil {
		return fmt.Errorf("Cannot find linuxkit executable, needed to build %s output type: %v", format, err)
	}
	commandLine := []string{"-q", "run", "qemu", "-disk", fmt.Sprintf("%s,size=%s,format=%s", filename, sizeString, format), "-disk", fmt.Sprintf("%s,format=raw", tardisk), "-kernel", imageFilename("mkimage")}
	if hyperkit && format == "raw" {
		state, err := ioutil.TempDir("", "s")
		if err != nil {
			return err
		}
		defer os.RemoveAll(state)
		commandLine = []string{"-q", "run", "hyperkit", "-state", state, "-disk", fmt.Sprintf("%s,size=%s,format=%s", filename, sizeString, format), "-disk", fmt.Sprintf("%s,format=raw", tardisk), imageFilename("mkimage")}
	}
	log.Debugf("run %s: %v", linuxkit, commandLine)
	cmd := exec.Command(linuxkit, commandLine...)
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Update mkimage used for building to lastest hashes<commit_after>package moby

import (
	"bytes"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"

	log "github.com/Sirupsen/logrus"
)

var linuxkitYaml = map[string]string{"mkimage": `
kernel:
  image: linuxkit/kernel:4.9.39
  cmdline: "console=ttyS0"
init:
  - linuxkit/init:00ab58c9681a0bf42b2e35134c1ccf1591ebb64d
  - linuxkit/runc:f5960b83a8766ae083efc744fa63dbf877450e4f
onboot:
  - name: mkimage
    image: linuxkit/mkimage:a63b8ee4c5de335afc32ba850e0af319b25b96c0
  - name: poweroff
    image: linuxkit/poweroff:3845c4d64d47a1ea367806be5547e44594b0fa91
trust:
  org:
    - linuxkit
`}

func imageFilename(name string) string {
	yaml := linuxkitYaml[name]
	hash := sha256.Sum256([]byte(yaml))
	return filepath.Join(MobyDir, "linuxkit", name+"-"+fmt.Sprintf("%x", hash))
}

func ensureLinuxkitImage(name string) error {
	filename := imageFilename(name)
	_, err1 := os.Stat(filename + "-kernel")
	_, err2 := os.Stat(filename + "-initrd.img")
	_, err3 := os.Stat(filename + "-cmdline")
	if err1 == nil && err2 == nil && err3 == nil {
		return nil
	}
	err := os.MkdirAll(filepath.Join(MobyDir, "linuxkit"), 0755)
	if err != nil {
		return err
	}
	// TODO clean up old files
	log.Infof("Building LinuxKit image %s to generate output formats", name)

	yaml := linuxkitYaml[name]

	m, err := NewConfig([]byte(yaml))
	if err != nil {
		return err
	}
	// TODO pass through --pull to here
	buf := new(bytes.Buffer)
	Build(m, buf, false, "")
	image := buf.Bytes()
	kernel, initrd, cmdline, err := tarToInitrd(image)
	if err != nil {
		return fmt.Errorf("Error converting to initrd: %v", err)
	}
	err = writeKernelInitrd(filename, kernel, initrd, cmdline)
	if err != nil {
		return err
	}

	return nil
}

func writeKernelInitrd(filename string, kernel []byte, initrd []byte, cmdline string) error {
	err := ioutil.WriteFile(filename+"-kernel", kernel, 0600)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(filename+"-initrd.img", initrd, 0600)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(filename+"-cmdline", []byte(cmdline), 0600)
	if err != nil {
		return err
	}
	return nil
}

func outputLinuxKit(format string, filename string, kernel []byte, initrd []byte, cmdline string, size int, hyperkit bool) error {
	log.Debugf("output linuxkit generated img: %s %s size %d", format, filename, size)

	tmp, err := ioutil.TempDir(filepath.Join(MobyDir, "tmp"), "moby")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmp)

	buf, err := tarInitrdKernel(kernel, initrd, cmdline)
	if err != nil {
		return err
	}

	tardisk := filepath.Join(tmp, "tardisk")
	f, err := os.Create(tardisk)
	if err != nil {
		return err
	}
	_, err = io.Copy(f, buf)
	if err != nil {
		return err
	}
	err = f.Close()
	if err != nil {
		return err
	}

	sizeString := fmt.Sprintf("%dM", size)
	_ = os.Remove(filename)
	_, err = os.Stat(filename)
	if err == nil || !os.IsNotExist(err) {
		return fmt.Errorf("Cannot remove existing file [%s]", filename)
	}
	linuxkit, err := exec.LookPath("linuxkit")
	if err != nil {
		return fmt.Errorf("Cannot find linuxkit executable, needed to build %s output type: %v", format, err)
	}
	commandLine := []string{"-q", "run", "qemu", "-disk", fmt.Sprintf("%s,size=%s,format=%s", filename, sizeString, format), "-disk", fmt.Sprintf("%s,format=raw", tardisk), "-kernel", imageFilename("mkimage")}
	if hyperkit && format == "raw" {
		state, err := ioutil.TempDir("", "s")
		if err != nil {
			return err
		}
		defer os.RemoveAll(state)
		commandLine = []string{"-q", "run", "hyperkit", "-state", state, "-disk", fmt.Sprintf("%s,size=%s,format=%s", filename, sizeString, format), "-disk", fmt.Sprintf("%s,format=raw", tardisk), imageFilename("mkimage")}
	}
	log.Debugf("run %s: %v", linuxkit, commandLine)
	cmd := exec.Command(linuxkit, commandLine...)
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package fuse

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

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pfs/pfsutil"
	"golang.org/x/net/context"
)

func Mount(apiClient pfs.ApiClient, repositoryName string, commitID string, mountPoint string) error {
	if err := os.MkdirAll(mountPoint, 0777); err != nil {
		return err
	}
	c, err := fuse.Mount(
		mountPoint,
		fuse.FSName("pfs"),
		fuse.Subtype("pfs"),
		fuse.VolumeName("pfs://"+repositoryName),
	)
	if err != nil {
		return err
	}
	defer c.Close()

	if err := fs.Serve(c, &filesystem{apiClient, repositoryName, commitID}); err != nil {
		return err
	}

	// check if the mount process has an error to report
	<-c.Ready
	return c.MountError
}

func Unmount(mountPoint string) error {
	return nil
}

type filesystem struct {
	apiClient      pfs.ApiClient
	repositoryName string
	commitID       string
}

func (f *filesystem) Root() (fs.Node, error) {
	return &directory{f, "/"}, nil
}

type directory struct {
	fs   *filesystem
	path string
}

func (*directory) Attr(ctx context.Context, a *fuse.Attr) error {
	a.Inode = 1
	a.Mode = os.ModeDir | 0555
	return nil
}

func (d *directory) Lookup(ctx context.Context, name string) (fs.Node, error) {
	response, err := pfsutil.GetFileInfo(
		d.fs.apiClient,
		d.fs.repositoryName,
		d.fs.commitID,
		filepath.Join(d.path, name),
	)
	if err != nil {
		return nil, err
	}
	switch response.GetFileInfo().FileType {
	case pfs.FileType_FILE_TYPE_NONE:
		log.Print("FileType_FILE_TYPE_NONE")
		return nil, fuse.ENOENT
	case pfs.FileType_FILE_TYPE_OTHER:
		log.Print("FileType_FILE_TYPE_OTHER")
		return nil, fuse.ENOENT
	case pfs.FileType_FILE_TYPE_REGULAR:
		return &file{d.fs, filepath.Join(d.path, name), response.GetFileInfo().SizeBytes}, nil
	case pfs.FileType_FILE_TYPE_DIR:
		return &directory{d.fs, filepath.Join(d.path, name)}, nil
	default:
		return nil, fmt.Errorf("Unrecognized FileType.")
	}
}

type file struct {
	fs   *filesystem
	path string
	size uint64
}

func (f *file) Attr(ctx context.Context, a *fuse.Attr) error {
	a.Inode = 2
	a.Mode = 0444
	a.Size = f.size
	return nil
}
<commit_msg>Adds ReadDirAll for directory.<commit_after>package fuse

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

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pfs/pfsutil"
	"golang.org/x/net/context"
)

func Mount(apiClient pfs.ApiClient, repositoryName string, commitID string, mountPoint string) error {
	if err := os.MkdirAll(mountPoint, 0777); err != nil {
		return err
	}
	c, err := fuse.Mount(
		mountPoint,
		fuse.FSName("pfs"),
		fuse.Subtype("pfs"),
		fuse.VolumeName("pfs://"+repositoryName),
	)
	if err != nil {
		return err
	}
	defer c.Close()

	if err := fs.Serve(c, &filesystem{apiClient, repositoryName, commitID}); err != nil {
		return err
	}

	// check if the mount process has an error to report
	<-c.Ready
	return c.MountError
}

func Unmount(mountPoint string) error {
	return nil
}

type filesystem struct {
	apiClient      pfs.ApiClient
	repositoryName string
	commitID       string
}

func (f *filesystem) Root() (fs.Node, error) {
	return &directory{f, "/"}, nil
}

type directory struct {
	fs   *filesystem
	path string
}

func (*directory) Attr(ctx context.Context, a *fuse.Attr) error {
	a.Inode = 1
	a.Mode = os.ModeDir | 0555
	return nil
}

func nodeFromFileInfo(fs *filesystem, fileInfo *pfs.FileInfo) (fs.Node, error) {
	switch fileInfo.FileType {
	case pfs.FileType_FILE_TYPE_NONE:
		log.Print("FileType_FILE_TYPE_NONE")
		return nil, fuse.ENOENT
	case pfs.FileType_FILE_TYPE_OTHER:
		log.Print("FileType_FILE_TYPE_OTHER")
		return nil, fuse.ENOENT
	case pfs.FileType_FILE_TYPE_REGULAR:
		return &file{fs, fileInfo.Path.Path, fileInfo.SizeBytes}, nil
	case pfs.FileType_FILE_TYPE_DIR:
		return &directory{fs, fileInfo.Path.Path}, nil
	default:
		return nil, fmt.Errorf("Unrecognized FileType.")
	}
}

func (d *directory) Lookup(ctx context.Context, name string) (fs.Node, error) {
	response, err := pfsutil.GetFileInfo(
		d.fs.apiClient,
		d.fs.repositoryName,
		d.fs.commitID,
		filepath.Join(d.path, name),
	)
	if err != nil {
		return nil, err
	}
	return nodeFromFileInfo(d.fs, response.GetFileInfo())
}

func (d *directory) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	response, err := pfsutil.ListFiles(d.fs.apiClient, d.fs.repositoryName, d.fs.commitID, d.path, 0, 1)
	if err != nil {
		return nil, err
	}
	var result []fuse.Dirent
	for _, fileInfo := range response.GetFileInfo() {
		switch fileInfo.FileType {
		case pfs.FileType_FILE_TYPE_NONE:
			continue
		case pfs.FileType_FILE_TYPE_OTHER:
			continue
		case pfs.FileType_FILE_TYPE_REGULAR:
			result = append(result, fuse.Dirent{Inode: 3, Name: fileInfo.Path.Path, Type: fuse.DT_File})
		case pfs.FileType_FILE_TYPE_DIR:
			result = append(result, fuse.Dirent{Inode: 3, Name: fileInfo.Path.Path, Type: fuse.DT_Dir})
		}
	}
	return result, nil
}

type file struct {
	fs   *filesystem
	path string
	size uint64
}

func (f *file) Attr(ctx context.Context, a *fuse.Attr) error {
	a.Inode = 2
	a.Mode = 0444
	a.Size = f.size
	return nil
}
<|endoftext|>


<commit_before>package horizon

import (
	"testing"

	. "github.com/smartystreets/goconvey/convey"
	"github.com/stellar/horizon/test"
)

func TestEffectActions(t *testing.T) {
	test.LoadScenario("base")

	Convey("Effect Actions:", t, func() {
		app := NewTestApp()
		defer app.Close()
		rh := NewRequestHelper(app)

		Convey("GET /effects", func() {
			w := rh.Get("/effects?limit=20", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 11)

			// test streaming, regression for https://github.com/stellar/horizon/issues/147
			// TODO: fix goji so that HttpResponseRecorder writers are counted considered flushable
			// w = rh.Get("/effects?limit=2", test.RequestHelperStreaming)
			// t.Log(w.Body.String())
			// So(w.Code, ShouldEqual, 200)

		})

		Convey("GET /ledgers/:ledger_id/effects", func() {
			w := rh.Get("/ledgers/1/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 0)

			w = rh.Get("/ledgers/2/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 9)

			w = rh.Get("/ledgers/3/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 2)
		})

		Convey("GET /accounts/:account_id/effects", func() {
			w := rh.Get("/accounts/GBRPYHIL2CI3FNQ4BXLFMNDLFJUNPU2HY3ZMFSHONUCEOASW7QC7OX2H/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 3)

			w = rh.Get("/accounts/GA5WBPYA5Y4WAEHXWR2UKO2UO4BUGHUQ74EUPKON2QHV4WRHOIRNKKH2/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 2)

			w = rh.Get("/accounts/GCXKG6RN4ONIEPCMNFB732A436Z5PNDSRLGWK7GBLCMQLIFO4S7EYWVU/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 3)
		})

		Convey("GET /transactions/:tx_id/effects", func() {
			w := rh.Get("/transactions/2374e99349b9ef7dba9a5db3339b78fda8f34777b1af33ba468ad5c0df946d4d/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 3)
		})

		Convey("GET /operations/:op_id/effects", func() {
			w := rh.Get("/operations/8589938689/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 3)
		})
	})
}
<commit_msg>Clear up TODO<commit_after>package horizon

import (
	"testing"

	. "github.com/smartystreets/goconvey/convey"
	"github.com/stellar/horizon/test"
)

func TestEffectActions(t *testing.T) {
	test.LoadScenario("base")

	Convey("Effect Actions:", t, func() {
		app := NewTestApp()
		defer app.Close()
		rh := NewRequestHelper(app)

		Convey("GET /effects", func() {
			w := rh.Get("/effects?limit=20", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 11)

			// test streaming, regression for https://github.com/stellar/horizon/issues/147
			// TODO: fix goji so that HttpResponseRecorder writers are considered flushable, allowing us
			// to test streaming responses
			// w = rh.Get("/effects?limit=2", test.RequestHelperStreaming)
			// t.Log(w.Body.String())
			// So(w.Code, ShouldEqual, 200)

		})

		Convey("GET /ledgers/:ledger_id/effects", func() {
			w := rh.Get("/ledgers/1/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 0)

			w = rh.Get("/ledgers/2/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 9)

			w = rh.Get("/ledgers/3/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 2)
		})

		Convey("GET /accounts/:account_id/effects", func() {
			w := rh.Get("/accounts/GBRPYHIL2CI3FNQ4BXLFMNDLFJUNPU2HY3ZMFSHONUCEOASW7QC7OX2H/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 3)

			w = rh.Get("/accounts/GA5WBPYA5Y4WAEHXWR2UKO2UO4BUGHUQ74EUPKON2QHV4WRHOIRNKKH2/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 2)

			w = rh.Get("/accounts/GCXKG6RN4ONIEPCMNFB732A436Z5PNDSRLGWK7GBLCMQLIFO4S7EYWVU/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 3)
		})

		Convey("GET /transactions/:tx_id/effects", func() {
			w := rh.Get("/transactions/2374e99349b9ef7dba9a5db3339b78fda8f34777b1af33ba468ad5c0df946d4d/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 3)
		})

		Convey("GET /operations/:op_id/effects", func() {
			w := rh.Get("/operations/8589938689/effects", test.RequestHelperNoop)
			So(w.Code, ShouldEqual, 200)
			So(w.Body, ShouldBePageOf, 3)
		})
	})
}
<|endoftext|>


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

// Simple file i/o and string manipulation, to avoid
// depending on strconv and bufio and strings.

package net

import (
	"io"
	"os"
)

type file struct {
	file *os.File
	data []byte
}

func (f *file) close() { f.file.Close() }

func (f *file) getLineFromData() (s string, ok bool) {
	data := f.data
	for i := 0; i < len(data); i++ {
		if data[i] == '\n' {
			s = string(data[0:i])
			ok = true
			// move data
			i++
			n := len(data) - i
			for j := 0; j < n; j++ {
				data[j] = data[i+j]
			}
			f.data = data[0:n]
			return
		}
	}
	if len(f.data) > 0 {
		s = string(data)
		f.data = nil
		ok = true
	}
	return
}

func (f *file) readLine() (s string, ok bool) {
	if s, ok = f.getLineFromData(); ok {
		return
	}
	if len(f.data) < cap(f.data) {
		ln := len(f.data)
		n, _ := io.ReadFull(f.file, f.data[ln:cap(f.data)])
		if n >= 0 {
			f.data = f.data[0 : ln+n]
		}
	}
	s, ok = f.getLineFromData()
	return
}

func open(name string) (*file, os.Error) {
	fd, err := os.Open(name, os.O_RDONLY, 0)
	if err != nil {
		return nil, err
	}
	return &file{fd, make([]byte, 1024)[0:0]}, nil
}

func byteIndex(s string, c byte) int {
	for i := 0; i < len(s); i++ {
		if s[i] == c {
			return i
		}
	}
	return -1
}

// Count occurrences in s of any bytes in t.
func countAnyByte(s string, t string) int {
	n := 0
	for i := 0; i < len(s); i++ {
		if byteIndex(t, s[i]) >= 0 {
			n++
		}
	}
	return n
}

// Split s at any bytes in t.
func splitAtBytes(s string, t string) []string {
	a := make([]string, 1+countAnyByte(s, t))
	n := 0
	last := 0
	for i := 0; i < len(s); i++ {
		if byteIndex(t, s[i]) >= 0 {
			if last < i {
				a[n] = string(s[last:i])
				n++
			}
			last = i + 1
		}
	}
	if last < len(s) {
		a[n] = string(s[last:])
		n++
	}
	return a[0:n]
}

func getFields(s string) []string { return splitAtBytes(s, " \r\t\n") }

// Bigger than we need, not too big to worry about overflow
const big = 0xFFFFFF

// Decimal to integer starting at &s[i0].
// Returns number, new offset, success.
func dtoi(s string, i0 int) (n int, i int, ok bool) {
	n = 0
	for i = i0; i < len(s) && '0' <= s[i] && s[i] <= '9'; i++ {
		n = n*10 + int(s[i]-'0')
		if n >= big {
			return 0, i, false
		}
	}
	if i == i0 {
		return 0, i, false
	}
	return n, i, true
}

// Hexadecimal to integer starting at &s[i0].
// Returns number, new offset, success.
func xtoi(s string, i0 int) (n int, i int, ok bool) {
	n = 0
	for i = i0; i < len(s); i++ {
		if '0' <= s[i] && s[i] <= '9' {
			n *= 16
			n += int(s[i] - '0')
		} else if 'a' <= s[i] && s[i] <= 'f' {
			n *= 16
			n += int(s[i]-'a') + 10
		} else if 'A' <= s[i] && s[i] <= 'F' {
			n *= 16
			n += int(s[i]-'A') + 10
		} else {
			break
		}
		if n >= big {
			return 0, i, false
		}
	}
	if i == i0 {
		return 0, i, false
	}
	return n, i, true
}

// Integer to decimal.
func itoa(i int) string {
	var buf [30]byte
	n := len(buf)
	neg := false
	if i < 0 {
		i = -i
		neg = true
	}
	ui := uint(i)
	for ui > 0 || n == len(buf) {
		n--
		buf[n] = byte('0' + ui%10)
		ui /= 10
	}
	if neg {
		n--
		buf[n] = '-'
	}
	return string(buf[n:])
}

// Number of occurrences of b in s.
func count(s string, b byte) int {
	n := 0
	for i := 0; i < len(s); i++ {
		if s[i] == b {
			n++
		}
	}
	return n
}

// Index of rightmost occurrence of b in s.
func last(s string, b byte) int {
	i := len(s)
	for i--; i >= 0; i-- {
		if s[i] == b {
			break
		}
	}
	return i
}
<commit_msg>net: parser should handle EOF without newline properly.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Simple file i/o and string manipulation, to avoid
// depending on strconv and bufio and strings.

package net

import (
	"io"
	"os"
)

type file struct {
	file  *os.File
	data  []byte
	atEOF bool
}

func (f *file) close() { f.file.Close() }

func (f *file) getLineFromData() (s string, ok bool) {
	data := f.data
	i := 0
	for i = 0; i < len(data); i++ {
		if data[i] == '\n' {
			s = string(data[0:i])
			ok = true
			// move data
			i++
			n := len(data) - i
			copy(data[0:], data[i:])
			f.data = data[0:n]
			return
		}
	}
	if f.atEOF && len(f.data) > 0 {
		// EOF, return all we have
		s = string(data)
		f.data = f.data[0:0]
		ok = true
	}
	return
}

func (f *file) readLine() (s string, ok bool) {
	if s, ok = f.getLineFromData(); ok {
		return
	}
	if len(f.data) < cap(f.data) {
		ln := len(f.data)
		n, err := io.ReadFull(f.file, f.data[ln:cap(f.data)])
		if n >= 0 {
			f.data = f.data[0 : ln+n]
		}
		if err == os.EOF {
			f.atEOF = true
		}
	}
	s, ok = f.getLineFromData()
	return
}

func open(name string) (*file, os.Error) {
	fd, err := os.Open(name, os.O_RDONLY, 0)
	if err != nil {
		return nil, err
	}
	return &file{fd, make([]byte, 1024)[0:0], false}, nil
}

func byteIndex(s string, c byte) int {
	for i := 0; i < len(s); i++ {
		if s[i] == c {
			return i
		}
	}
	return -1
}

// Count occurrences in s of any bytes in t.
func countAnyByte(s string, t string) int {
	n := 0
	for i := 0; i < len(s); i++ {
		if byteIndex(t, s[i]) >= 0 {
			n++
		}
	}
	return n
}

// Split s at any bytes in t.
func splitAtBytes(s string, t string) []string {
	a := make([]string, 1+countAnyByte(s, t))
	n := 0
	last := 0
	for i := 0; i < len(s); i++ {
		if byteIndex(t, s[i]) >= 0 {
			if last < i {
				a[n] = string(s[last:i])
				n++
			}
			last = i + 1
		}
	}
	if last < len(s) {
		a[n] = string(s[last:])
		n++
	}
	return a[0:n]
}

func getFields(s string) []string { return splitAtBytes(s, " \r\t\n") }

// Bigger than we need, not too big to worry about overflow
const big = 0xFFFFFF

// Decimal to integer starting at &s[i0].
// Returns number, new offset, success.
func dtoi(s string, i0 int) (n int, i int, ok bool) {
	n = 0
	for i = i0; i < len(s) && '0' <= s[i] && s[i] <= '9'; i++ {
		n = n*10 + int(s[i]-'0')
		if n >= big {
			return 0, i, false
		}
	}
	if i == i0 {
		return 0, i, false
	}
	return n, i, true
}

// Hexadecimal to integer starting at &s[i0].
// Returns number, new offset, success.
func xtoi(s string, i0 int) (n int, i int, ok bool) {
	n = 0
	for i = i0; i < len(s); i++ {
		if '0' <= s[i] && s[i] <= '9' {
			n *= 16
			n += int(s[i] - '0')
		} else if 'a' <= s[i] && s[i] <= 'f' {
			n *= 16
			n += int(s[i]-'a') + 10
		} else if 'A' <= s[i] && s[i] <= 'F' {
			n *= 16
			n += int(s[i]-'A') + 10
		} else {
			break
		}
		if n >= big {
			return 0, i, false
		}
	}
	if i == i0 {
		return 0, i, false
	}
	return n, i, true
}

// Integer to decimal.
func itoa(i int) string {
	var buf [30]byte
	n := len(buf)
	neg := false
	if i < 0 {
		i = -i
		neg = true
	}
	ui := uint(i)
	for ui > 0 || n == len(buf) {
		n--
		buf[n] = byte('0' + ui%10)
		ui /= 10
	}
	if neg {
		n--
		buf[n] = '-'
	}
	return string(buf[n:])
}

// Number of occurrences of b in s.
func count(s string, b byte) int {
	n := 0
	for i := 0; i < len(s); i++ {
		if s[i] == b {
			n++
		}
	}
	return n
}

// Index of rightmost occurrence of b in s.
func last(s string, b byte) int {
	i := len(s)
	for i--; i >= 0; i-- {
		if s[i] == b {
			break
		}
	}
	return i
}
<|endoftext|>


<commit_before>package ebs

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/mitchellh/multistep"
	awscommon "github.com/mitchellh/packer/builder/amazon/common"
	"github.com/mitchellh/packer/packer"
)

type stepCreateEncryptedAMICopy struct {
	image *ec2.Image
}

func (s *stepCreateEncryptedAMICopy) Run(state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(Config)
	ec2conn := state.Get("ec2").(*ec2.EC2)
	ui := state.Get("ui").(packer.Ui)

	// Encrypt boot not set, so skip step
	if !config.AMIConfig.AMIEncryptBootVolume {
		return multistep.ActionContinue
	}

	ui.Say("Creating Encrypted AMI Copy")

	amis := state.Get("amis").(map[string]string)
	var region, id string
	if amis != nil {
		for region, id = range amis {
			break // Only get the first
		}
	}

	ui.Say(fmt.Sprintf("Copying AMI: %s(%s)", region, id))

	copyOpts := &ec2.CopyImageInput{
		Name:          &config.AMIName, // Try to overwrite existing AMI
		SourceImageId: aws.String(id),
		SourceRegion:  aws.String(region),
		Encrypted:     aws.Bool(true),
	}

	copyResp, err := ec2conn.CopyImage(copyOpts)
	if err != nil {
		err := fmt.Errorf("Error copying AMI: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// Wait for the copy to become ready
	stateChange := awscommon.StateChangeConf{
		Pending:   []string{"pending"},
		Target:    "available",
		Refresh:   awscommon.AMIStateRefreshFunc(ec2conn, *copyResp.ImageId),
		StepState: state,
	}

	ui.Say("Waiting for AMI copy to become ready...")
	if _, err := awscommon.WaitForState(&stateChange); err != nil {
		err := fmt.Errorf("Error waiting for AMI Copy: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// Get the unencrypted AMI image
	unencImagesResp, err := ec2conn.DescribeImages(&ec2.DescribeImagesInput{ImageIds: []*string{aws.String(id)}})
	if err != nil {
		err := fmt.Errorf("Error searching for AMI: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	unencImage := unencImagesResp.Images[0]

	// Remove unencrypted AMI
	ui.Say("Deregistering unecrypted AMI")
	deregisterOpts := &ec2.DeregisterImageInput{ImageId: aws.String(id)}
	if _, err := ec2conn.DeregisterImage(deregisterOpts); err != nil {
		ui.Error(fmt.Sprintf("Error deregistering AMI, may still be around: %s", err))
		return multistep.ActionHalt
	}

	// Remove associated unencrypted snapshot(s)
	ui.Say("Deleting unencrypted snapshots")

	for _, blockDevice := range unencImage.BlockDeviceMappings {
		if blockDevice.Ebs != nil {
			if blockDevice.Ebs.SnapshotId != nil {
				ui.Message(fmt.Sprintf("Snapshot ID: %s", *blockDevice.Ebs.SnapshotId))
				deleteSnapOpts := &ec2.DeleteSnapshotInput{
					SnapshotId: aws.String(*blockDevice.Ebs.SnapshotId),
				}
				if _, err := ec2conn.DeleteSnapshot(deleteSnapOpts); err != nil {
					ui.Error(fmt.Sprintf("Error deleting snapshot, may still be around: %s", err))
					return multistep.ActionHalt
				}
			}
		}
	}

	// Replace original AMI ID with Encrypted ID in state
	amis[region] = *copyResp.ImageId
	state.Put("amis", amis)

	imagesResp, err := ec2conn.DescribeImages(&ec2.DescribeImagesInput{ImageIds: []*string{copyResp.ImageId}})
	if err != nil {
		err := fmt.Errorf("Error searching for AMI: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	s.image = imagesResp.Images[0]

	return multistep.ActionContinue
}

func (s *stepCreateEncryptedAMICopy) Cleanup(state multistep.StateBag) {
	if s.image == nil {
		return
	}

	_, cancelled := state.GetOk(multistep.StateCancelled)
	_, halted := state.GetOk(multistep.StateHalted)
	if !cancelled && !halted {
		return
	}

	ec2conn := state.Get("ec2").(*ec2.EC2)
	ui := state.Get("ui").(packer.Ui)

	ui.Say("Deregistering the AMI because cancelation or error...")
	deregisterOpts := &ec2.DeregisterImageInput{ImageId: s.image.ImageId}
	if _, err := ec2conn.DeregisterImage(deregisterOpts); err != nil {
		ui.Error(fmt.Sprintf("Error deregistering AMI, may still be around: %s", err))
		return
	}
}
<commit_msg>Fix typo in encrypted ami step log message<commit_after>package ebs

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/mitchellh/multistep"
	awscommon "github.com/mitchellh/packer/builder/amazon/common"
	"github.com/mitchellh/packer/packer"
)

type stepCreateEncryptedAMICopy struct {
	image *ec2.Image
}

func (s *stepCreateEncryptedAMICopy) Run(state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(Config)
	ec2conn := state.Get("ec2").(*ec2.EC2)
	ui := state.Get("ui").(packer.Ui)

	// Encrypt boot not set, so skip step
	if !config.AMIConfig.AMIEncryptBootVolume {
		return multistep.ActionContinue
	}

	ui.Say("Creating Encrypted AMI Copy")

	amis := state.Get("amis").(map[string]string)
	var region, id string
	if amis != nil {
		for region, id = range amis {
			break // Only get the first
		}
	}

	ui.Say(fmt.Sprintf("Copying AMI: %s(%s)", region, id))

	copyOpts := &ec2.CopyImageInput{
		Name:          &config.AMIName, // Try to overwrite existing AMI
		SourceImageId: aws.String(id),
		SourceRegion:  aws.String(region),
		Encrypted:     aws.Bool(true),
	}

	copyResp, err := ec2conn.CopyImage(copyOpts)
	if err != nil {
		err := fmt.Errorf("Error copying AMI: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// Wait for the copy to become ready
	stateChange := awscommon.StateChangeConf{
		Pending:   []string{"pending"},
		Target:    "available",
		Refresh:   awscommon.AMIStateRefreshFunc(ec2conn, *copyResp.ImageId),
		StepState: state,
	}

	ui.Say("Waiting for AMI copy to become ready...")
	if _, err := awscommon.WaitForState(&stateChange); err != nil {
		err := fmt.Errorf("Error waiting for AMI Copy: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// Get the unencrypted AMI image
	unencImagesResp, err := ec2conn.DescribeImages(&ec2.DescribeImagesInput{ImageIds: []*string{aws.String(id)}})
	if err != nil {
		err := fmt.Errorf("Error searching for AMI: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	unencImage := unencImagesResp.Images[0]

	// Remove unencrypted AMI
	ui.Say("Deregistering unencrypted AMI")
	deregisterOpts := &ec2.DeregisterImageInput{ImageId: aws.String(id)}
	if _, err := ec2conn.DeregisterImage(deregisterOpts); err != nil {
		ui.Error(fmt.Sprintf("Error deregistering AMI, may still be around: %s", err))
		return multistep.ActionHalt
	}

	// Remove associated unencrypted snapshot(s)
	ui.Say("Deleting unencrypted snapshots")

	for _, blockDevice := range unencImage.BlockDeviceMappings {
		if blockDevice.Ebs != nil {
			if blockDevice.Ebs.SnapshotId != nil {
				ui.Message(fmt.Sprintf("Snapshot ID: %s", *blockDevice.Ebs.SnapshotId))
				deleteSnapOpts := &ec2.DeleteSnapshotInput{
					SnapshotId: aws.String(*blockDevice.Ebs.SnapshotId),
				}
				if _, err := ec2conn.DeleteSnapshot(deleteSnapOpts); err != nil {
					ui.Error(fmt.Sprintf("Error deleting snapshot, may still be around: %s", err))
					return multistep.ActionHalt
				}
			}
		}
	}

	// Replace original AMI ID with Encrypted ID in state
	amis[region] = *copyResp.ImageId
	state.Put("amis", amis)

	imagesResp, err := ec2conn.DescribeImages(&ec2.DescribeImagesInput{ImageIds: []*string{copyResp.ImageId}})
	if err != nil {
		err := fmt.Errorf("Error searching for AMI: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	s.image = imagesResp.Images[0]

	return multistep.ActionContinue
}

func (s *stepCreateEncryptedAMICopy) Cleanup(state multistep.StateBag) {
	if s.image == nil {
		return
	}

	_, cancelled := state.GetOk(multistep.StateCancelled)
	_, halted := state.GetOk(multistep.StateHalted)
	if !cancelled && !halted {
		return
	}

	ec2conn := state.Get("ec2").(*ec2.EC2)
	ui := state.Get("ui").(packer.Ui)

	ui.Say("Deregistering the AMI because cancelation or error...")
	deregisterOpts := &ec2.DeregisterImageInput{ImageId: s.image.ImageId}
	if _, err := ec2conn.DeregisterImage(deregisterOpts); err != nil {
		ui.Error(fmt.Sprintf("Error deregistering AMI, may still be around: %s", err))
		return
	}
}
<|endoftext|>


<commit_before>package common

import (
	"testing"
)

func TestVMXConfigPrepare(t *testing.T) {
	c := new(VMXConfig)
	c.VMXData = map[string]string{
		"one": "foo",
		"two": "bar",
	}

	errs := c.Prepare(testConfigTemplate(t), "")
	if len(errs) > 0 {
		t.Fatalf("bad: %#v", errs)
	}

	if len(c.VMXData) != 2 {
		t.Fatal("should have two items in VMXData")
	}
}
<commit_msg>Revert "Correctly call prepare function in test"<commit_after>package common

import (
	"testing"
)

func TestVMXConfigPrepare(t *testing.T) {
	c := new(VMXConfig)
	c.VMXData = map[string]string{
		"one": "foo",
		"two": "bar",
	}

	errs := c.Prepare(testConfigTemplate(t))
	if len(errs) > 0 {
		t.Fatalf("bad: %#v", errs)
	}

	if len(c.VMXData) != 2 {
		t.Fatal("should have two items in VMXData")
	}
}
<|endoftext|>