<commit_before>/* Copyright (c) 2016 Jason Ish
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
 * IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

package elasticsearch

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/jasonish/evebox/httputil"

	"bytes"
	"github.com/jasonish/evebox/log"
	"github.com/jasonish/evebox/resources"
	"github.com/jasonish/evebox/util"
	"github.com/pkg/errors"
	"io/ioutil"
	"math"
	"time"
)

type ElasticSearch struct {
	baseUrl  string
	username string
	password string

	EventSearchIndex string
	EventBaseIndex   string

	// These are filled on a ping request. Perhaps when creating an
	// instance of this object we should ping.
	VersionString string
	MajorVersion  int64
	MinorVersion  int64

	// The keyword to use for "keyword" type queries. Older versions
	// of the Logstash template used "raw", newer ones use "keyword".
	keyword string

	useIpDatatype bool

	// Set to true if keyword checks should not be done.
	noKeyword bool

	HttpClient *httputil.HttpClient
}

func New(url string) *ElasticSearch {
	// First, strip any trailing / from the URL.
	for strings.HasSuffix(url, "/") {
		url = strings.TrimSuffix(url, "/")
	}

	HttpClient := httputil.NewHttpClient()

	HttpClient.SetBaseUrl(url)

	es := &ElasticSearch{
		baseUrl:    url,
		HttpClient: HttpClient,
	}

	return es
}

func (es *ElasticSearch) DisableCertCheck(disableCertCheck bool) {
	es.HttpClient.DisableCertCheck(disableCertCheck)
}

func (es *ElasticSearch) SetUsernamePassword(username ...string) error {
	return es.HttpClient.SetUsernamePassword(username...)
}

func (es *ElasticSearch) SetEventIndex(index string) {
	if strings.HasSuffix(index, "*") {
		baseIndex := strings.TrimSuffix(strings.TrimSuffix(index, "*"), "-")
		es.EventSearchIndex = index
		es.EventBaseIndex = baseIndex
	} else {
		es.EventBaseIndex = index
		es.EventSearchIndex = fmt.Sprintf("%s-*", index)
	}
	log.Info("Event base index: %s", es.EventBaseIndex)
	log.Info("Event search index: %s", es.EventSearchIndex)
}

func (es *ElasticSearch) Decode(response *http.Response, v interface{}) error {
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	log.Debug("Decoding response (truncated at 1024 bytes): %s",
		string(body)[:int(math.Min(1024, float64(len(body))))])
	decoder := json.NewDecoder(bytes.NewReader(body))
	decoder.UseNumber()
	return decoder.Decode(v)
}

func (es *ElasticSearch) Ping() (*PingResponse, error) {
	response, err := es.HttpClient.Get("")
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	if response.StatusCode != 200 {
		return nil, DecodeResponseAsError(response)
	}

	body, err := DecodeResponse(response)
	if err != nil {
		return nil, err
	}

	pingResponse := PingResponse{body}
	es.MajorVersion, es.MinorVersion = pingResponse.ParseVersion()
	return &pingResponse, nil
}

func (es *ElasticSearch) CheckTemplate(name string) (exists bool, err error) {
	response, err := es.HttpClient.Head(fmt.Sprintf("_template/%s", name))
	if err != nil {
		return exists, err
	}
	exists = response.StatusCode == 200
	return exists, nil
}

func (es *ElasticSearch) GetTemplate(name string) (util.JsonMap, error) {
	url := fmt.Sprintf("_template/%s", name)
	log.Debug("Fetching template %s", url)
	response, err := es.HttpClient.Get(url)
	if err != nil {
		return nil, err
	}

	template := util.JsonMap{}

	if err := es.Decode(response, &template); err != nil {
		return nil, err
	}

	return template, nil
}

func (es *ElasticSearch) GetUseIpDatatype() bool {
	return es.useIpDatatype
}

// GetKeywordType is a crude way of determining if the template is using
// Logstash 5 keyword type, or Logstash 2 "raw" type.
func (es *ElasticSearch) GetKeywordType(index string) (string, error) {

	// It appears as though Filebeat indexes don't need this.
	if index == "filebeat" {
		es.noKeyword = true
		return "", nil
	}

	if index == "" {
		index = es.EventBaseIndex
	}
	template, err := es.GetTemplate(index)
	if err != nil {
		log.Warning("Failed to get template from Elastic Search, keyword resolution delayed.")
		return "", nil
	}

	version := template.GetMap(index).Get("version")
	log.Debug("Found template version %v", version)

	properties := template.GetMap(index).
		GetMap("mappings").GetMap("_default_").GetMap("properties")
	if properties != nil {
		destIpType := properties.GetMap("dest_ip").GetString("type")
		sourceIpType := properties.GetMap("src_ip").GetString("type")
		if (destIpType == "ip" && sourceIpType == "ip") {
			log.Info("Elastic Search EVE records are using IP datatype.")
			es.useIpDatatype = true
		}
	}

	dynamicTemplates := template.GetMap(index).
		GetMap("mappings").
		GetMap("_default_").
		GetMapList("dynamic_templates")
	if dynamicTemplates == nil {
		log.Warning("Failed to parse template, keyword resolution delayed.")
		log.Warning("Template: %s", util.ToJson(template))
		return "", nil
	}
	for _, entry := range dynamicTemplates {
		if entry["string_fields"] != nil {
			mappingType := entry.GetMap("string_fields").
				GetMap("mapping").
				GetMap("fields").
				GetMap("keyword").
				Get("type")
			if mappingType == "keyword" {
				return "keyword", nil
			}

			if entry.GetMap("string_fields").GetMap("mapping").GetMap("fields").GetMap("raw") != nil {
				return "raw", nil
			}
		}
	}

	log.Warning("Failed to parse template, keyword resolution delayed.")
	log.Warning("Template: %s", util.ToJson(template))
	return "", nil
}

func (es *ElasticSearch) SetKeyword(keyword string) {
	if keyword == "" {
		es.noKeyword = true
	} else {
		es.keyword = keyword
	}
}

func (es *ElasticSearch) InitKeyword() error {
	keyword, err := es.GetKeywordType(es.EventBaseIndex)
	if err != nil {
		return err
	}
	es.keyword = keyword
	log.Info("Elastic Search keyword initialized to \"%s\"", es.keyword)
	return nil
}

func (es *ElasticSearch) LoadTemplate(index string, majorVersion int64) error {

	var templateFilename string

	if majorVersion == 0 {
		// Version unknown, get it.
		pingResponse, err := es.Ping()
		if err != nil {
			log.Warning("Failed to ping Elastic Search: %v", err)
			return err
		}
		majorVersion = pingResponse.MajorVersion()
	}

	if majorVersion < 5 {
		return fmt.Errorf("Elastic Search versions less than 5 not supported.")
	} else if majorVersion == 5 {
		templateFilename = "template-es5x.json"
	} else {
		return fmt.Errorf("No template for Elastic Search with major version %d", majorVersion)
	}

	templateFile, err := resources.GetReader(fmt.Sprintf("elasticsearch/%s", templateFilename))
	if err != nil {
		return err
	}

	decoder := json.NewDecoder(templateFile)
	decoder.UseNumber()

	var template map[string]interface{}
	err = decoder.Decode(&template)
	if err != nil {
		return err
	}
	template["template"] = fmt.Sprintf("%s-*", index)

	response, err := es.HttpClient.PutJson(fmt.Sprintf("_template/%s", index), template)
	if err != nil {
		return err
	}
	if response.StatusCode != http.StatusOK {
		return DecodeResponseAsError(response)
	}
	es.HttpClient.DiscardResponse(response)

	// Success.
	return nil
}

type DatastoreError struct {
	Message string
	Cause   error
}

func (e *DatastoreError) Error() string {
	if e.Message != "" && e.Cause != nil {
		return fmt.Sprintf("%s: %s", e.Message, e.Cause.Error())
	} else if e.Message != "" {
		return e.Message
	} else if e.Cause != nil {
		return e.Cause.Error()
	}
	return ""
}

func (es *ElasticSearch) Search(query interface{}) (*Response, error) {
	if es.keyword == "" && !es.noKeyword {
		log.Warning("Search keyword not known, trying again.")
		es.InitKeyword()
	}

	path := fmt.Sprintf("%s/_search", es.EventSearchIndex)
	response, err := es.HttpClient.PostJson(path, query)
	if err != nil {
		return nil, errors.WithStack(&DatastoreError{
			Message: "Failed to connect to Elastic Search",
			Cause:   err,
		})
	}
	defer response.Body.Close()
	return DecodeResponse(response)
}

func (es *ElasticSearch) SearchScroll(body interface{}, duration string) (*Response, error) {
	path := fmt.Sprintf("%s/_search?scroll=%s", es.EventSearchIndex, duration)
	response, err := es.HttpClient.PostJson(path, body)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return nil, DecodeResponseAsError(response)
	}

	return DecodeResponse(response)
}

func (es *ElasticSearch) Scroll(scrollId string, duration string) (*Response, error) {
	body := map[string]interface{}{
		"scroll_id": scrollId,
		"scroll":    duration,
	}
	response, err := es.HttpClient.PostJson("_search/scroll", body)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	return DecodeResponse(response)
}

func (es *ElasticSearch) DeleteScroll(scrollId string) (*http.Response, error) {
	return es.HttpClient.Delete("_search/scroll", "application/json",
		strings.NewReader(scrollId))
}

func (es *ElasticSearch) PartialUpdate(index string, doctype string, id string,
	doc interface{}) (*http.Response, error) {
	body := map[string]interface{}{
		"doc": doc,
	}
	return es.HttpClient.PostJson(fmt.Sprintf("%s/%s/%s/_update?refresh=true",
		index, doctype, id), body)
}

func IsError(response *http.Response) error {
	if response.StatusCode < 400 {
		return nil
	}
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return errors.Wrap(err, "failed to read response")
	}
	return errors.Errorf("%s %s", response.Status, string(body))
}

func (es *ElasticSearch) Update(index string, docType string, docId string,
	body interface{}) (*Response, error) {
	response, err := es.HttpClient.PostJson(fmt.Sprintf("%s/%s/%s/_update?refresh=true",
		index, docType, docId), body)
	if err != nil {
		return nil, errors.Wrap(err, "http request failed")
	}
	defer response.Body.Close()
	if err := IsError(response); err != nil {
		return nil, err
	}
	return DecodeResponse(response)
}

// Refresh refreshes all indices logging any error but not returning and
// discarding the response so the caller doesn't have to.
func (es *ElasticSearch) Refresh() {
	response, err := es.HttpClient.PostString("_refresh", "application/json", "{}")
	if err != nil {
		log.Error("Failed to refresh Elastic Search: %v", err)
		return
	}
	io.Copy(ioutil.Discard, response.Body)
}

func (es *ElasticSearch) FormatKeyword(keyword string) string {
	if es.keyword == "" {
		return keyword
	}
	return fmt.Sprintf("%s.%s", keyword, es.keyword)
}

func (s *ElasticSearch) doUpdateByQuery(query interface{}) (util.JsonMap, error) {
	var response util.JsonMap
	err := s.HttpClient.PostJsonDecodeResponse(
		fmt.Sprintf("%s/_update_by_query?refresh=true&conflicts=proceed",
			s.EventSearchIndex), query, &response)
	if err != nil {
		return nil, errors.Wrap(err, "request failed")
	}

	return response, nil
}

// FormatTimestampUTC formats a time.Time into the format generally used
// by Elastic Search, in particular the @timestamp field.
func FormatTimestampUTC(timestamp time.Time) string {
	return timestamp.UTC().Format("2006-01-02T15:04:05.000Z")
}
<commit_msg>es: compile time flag to log search requests/responses<commit_after>/* Copyright (c) 2016 Jason Ish
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
 * IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

package elasticsearch

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/jasonish/evebox/httputil"

	"bytes"
	"github.com/jasonish/evebox/log"
	"github.com/jasonish/evebox/resources"
	"github.com/jasonish/evebox/util"
	"github.com/pkg/errors"
	"io/ioutil"
	"math"
	"time"
)

const LOG_REQUEST_RESPONSE = false

type ElasticSearch struct {
	baseUrl  string
	username string
	password string

	EventSearchIndex string
	EventBaseIndex   string

	// These are filled on a ping request. Perhaps when creating an
	// instance of this object we should ping.
	VersionString string
	MajorVersion  int64
	MinorVersion  int64

	// The keyword to use for "keyword" type queries. Older versions
	// of the Logstash template used "raw", newer ones use "keyword".
	keyword string

	useIpDatatype bool

	// Set to true if keyword checks should not be done.
	noKeyword bool

	HttpClient *httputil.HttpClient
}

func New(url string) *ElasticSearch {
	// First, strip any trailing / from the URL.
	for strings.HasSuffix(url, "/") {
		url = strings.TrimSuffix(url, "/")
	}

	HttpClient := httputil.NewHttpClient()

	HttpClient.SetBaseUrl(url)

	es := &ElasticSearch{
		baseUrl:    url,
		HttpClient: HttpClient,
	}

	return es
}

func (es *ElasticSearch) DisableCertCheck(disableCertCheck bool) {
	es.HttpClient.DisableCertCheck(disableCertCheck)
}

func (es *ElasticSearch) SetUsernamePassword(username ...string) error {
	return es.HttpClient.SetUsernamePassword(username...)
}

func (es *ElasticSearch) SetEventIndex(index string) {
	if strings.HasSuffix(index, "*") {
		baseIndex := strings.TrimSuffix(strings.TrimSuffix(index, "*"), "-")
		es.EventSearchIndex = index
		es.EventBaseIndex = baseIndex
	} else {
		es.EventBaseIndex = index
		es.EventSearchIndex = fmt.Sprintf("%s-*", index)
	}
	log.Info("Event base index: %s", es.EventBaseIndex)
	log.Info("Event search index: %s", es.EventSearchIndex)
}

func (es *ElasticSearch) Decode(response *http.Response, v interface{}) error {
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	log.Debug("Decoding response (truncated at 1024 bytes): %s",
		string(body)[:int(math.Min(1024, float64(len(body))))])
	decoder := json.NewDecoder(bytes.NewReader(body))
	decoder.UseNumber()
	return decoder.Decode(v)
}

func (es *ElasticSearch) Ping() (*PingResponse, error) {
	response, err := es.HttpClient.Get("")
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	if response.StatusCode != 200 {
		return nil, DecodeResponseAsError(response)
	}

	body, err := DecodeResponse(response)
	if err != nil {
		return nil, err
	}

	pingResponse := PingResponse{body}
	es.MajorVersion, es.MinorVersion = pingResponse.ParseVersion()
	return &pingResponse, nil
}

func (es *ElasticSearch) CheckTemplate(name string) (exists bool, err error) {
	response, err := es.HttpClient.Head(fmt.Sprintf("_template/%s", name))
	if err != nil {
		return exists, err
	}
	exists = response.StatusCode == 200
	return exists, nil
}

func (es *ElasticSearch) GetTemplate(name string) (util.JsonMap, error) {
	url := fmt.Sprintf("_template/%s", name)
	log.Debug("Fetching template %s", url)
	response, err := es.HttpClient.Get(url)
	if err != nil {
		return nil, err
	}

	template := util.JsonMap{}

	if err := es.Decode(response, &template); err != nil {
		return nil, err
	}

	return template, nil
}

func (es *ElasticSearch) GetUseIpDatatype() bool {
	return es.useIpDatatype
}

// GetKeywordType is a crude way of determining if the template is using
// Logstash 5 keyword type, or Logstash 2 "raw" type.
func (es *ElasticSearch) GetKeywordType(index string) (string, error) {

	// It appears as though Filebeat indexes don't need this.
	if index == "filebeat" {
		es.noKeyword = true
		return "", nil
	}

	if index == "" {
		index = es.EventBaseIndex
	}
	template, err := es.GetTemplate(index)
	if err != nil {
		log.Warning("Failed to get template from Elastic Search, keyword resolution delayed.")
		return "", nil
	}

	version := template.GetMap(index).Get("version")
	log.Debug("Found template version %v", version)

	properties := template.GetMap(index).
		GetMap("mappings").GetMap("_default_").GetMap("properties")
	if properties != nil {
		destIpType := properties.GetMap("dest_ip").GetString("type")
		sourceIpType := properties.GetMap("src_ip").GetString("type")
		if (destIpType == "ip" && sourceIpType == "ip") {
			log.Info("Elastic Search EVE records are using IP datatype.")
			es.useIpDatatype = true
		}
	}

	dynamicTemplates := template.GetMap(index).
		GetMap("mappings").
		GetMap("_default_").
		GetMapList("dynamic_templates")
	if dynamicTemplates == nil {
		log.Warning("Failed to parse template, keyword resolution delayed.")
		log.Warning("Template: %s", util.ToJson(template))
		return "", nil
	}
	for _, entry := range dynamicTemplates {
		if entry["string_fields"] != nil {
			mappingType := entry.GetMap("string_fields").
				GetMap("mapping").
				GetMap("fields").
				GetMap("keyword").
				Get("type")
			if mappingType == "keyword" {
				return "keyword", nil
			}

			if entry.GetMap("string_fields").GetMap("mapping").GetMap("fields").GetMap("raw") != nil {
				return "raw", nil
			}
		}
	}

	log.Warning("Failed to parse template, keyword resolution delayed.")
	log.Warning("Template: %s", util.ToJson(template))
	return "", nil
}

func (es *ElasticSearch) SetKeyword(keyword string) {
	if keyword == "" {
		es.noKeyword = true
	} else {
		es.keyword = keyword
	}
}

func (es *ElasticSearch) InitKeyword() error {
	keyword, err := es.GetKeywordType(es.EventBaseIndex)
	if err != nil {
		return err
	}
	es.keyword = keyword
	log.Info("Elastic Search keyword initialized to \"%s\"", es.keyword)
	return nil
}

func (es *ElasticSearch) LoadTemplate(index string, majorVersion int64) error {

	var templateFilename string

	if majorVersion == 0 {
		// Version unknown, get it.
		pingResponse, err := es.Ping()
		if err != nil {
			log.Warning("Failed to ping Elastic Search: %v", err)
			return err
		}
		majorVersion = pingResponse.MajorVersion()
	}

	if majorVersion < 5 {
		return fmt.Errorf("Elastic Search versions less than 5 not supported.")
	} else if majorVersion == 5 {
		templateFilename = "template-es5x.json"
	} else {
		return fmt.Errorf("No template for Elastic Search with major version %d", majorVersion)
	}

	templateFile, err := resources.GetReader(fmt.Sprintf("elasticsearch/%s", templateFilename))
	if err != nil {
		return err
	}

	decoder := json.NewDecoder(templateFile)
	decoder.UseNumber()

	var template map[string]interface{}
	err = decoder.Decode(&template)
	if err != nil {
		return err
	}
	template["template"] = fmt.Sprintf("%s-*", index)

	response, err := es.HttpClient.PutJson(fmt.Sprintf("_template/%s", index), template)
	if err != nil {
		return err
	}
	if response.StatusCode != http.StatusOK {
		return DecodeResponseAsError(response)
	}
	es.HttpClient.DiscardResponse(response)

	// Success.
	return nil
}

type DatastoreError struct {
	Message string
	Cause   error
}

func (e *DatastoreError) Error() string {
	if e.Message != "" && e.Cause != nil {
		return fmt.Sprintf("%s: %s", e.Message, e.Cause.Error())
	} else if e.Message != "" {
		return e.Message
	} else if e.Cause != nil {
		return e.Cause.Error()
	}
	return ""
}

func (es *ElasticSearch) Search(query interface{}) (*Response, error) {
	if es.keyword == "" && !es.noKeyword {
		log.Warning("Search keyword not known, trying again.")
		es.InitKeyword()
	}

	path := fmt.Sprintf("%s/_search", es.EventSearchIndex)

	if LOG_REQUEST_RESPONSE {
		log.Debug("REQUEST: POST %s: %s", path, util.ToJson(query))
	}

	response, err := es.HttpClient.PostJson(path, query)
	if err != nil {
		return nil, errors.WithStack(&DatastoreError{
			Message: "Failed to connect to Elastic Search",
			Cause:   err,
		})
	}

	defer response.Body.Close()
	decodedResponse, err := DecodeResponse(response)

	if LOG_REQUEST_RESPONSE {
		log.Debug("RESPONSE: POST %s: %s", path, decodedResponse.Raw)
	}

	return decodedResponse, err
}

func (es *ElasticSearch) SearchScroll(body interface{}, duration string) (*Response, error) {
	path := fmt.Sprintf("%s/_search?scroll=%s", es.EventSearchIndex, duration)
	response, err := es.HttpClient.PostJson(path, body)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return nil, DecodeResponseAsError(response)
	}

	return DecodeResponse(response)
}

func (es *ElasticSearch) Scroll(scrollId string, duration string) (*Response, error) {
	body := map[string]interface{}{
		"scroll_id": scrollId,
		"scroll":    duration,
	}
	response, err := es.HttpClient.PostJson("_search/scroll", body)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	return DecodeResponse(response)
}

func (es *ElasticSearch) DeleteScroll(scrollId string) (*http.Response, error) {
	return es.HttpClient.Delete("_search/scroll", "application/json",
		strings.NewReader(scrollId))
}

func (es *ElasticSearch) PartialUpdate(index string, doctype string, id string,
	doc interface{}) (*http.Response, error) {
	body := map[string]interface{}{
		"doc": doc,
	}
	return es.HttpClient.PostJson(fmt.Sprintf("%s/%s/%s/_update?refresh=true",
		index, doctype, id), body)
}

func IsError(response *http.Response) error {
	if response.StatusCode < 400 {
		return nil
	}
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return errors.Wrap(err, "failed to read response")
	}
	return errors.Errorf("%s %s", response.Status, string(body))
}

func (es *ElasticSearch) Update(index string, docType string, docId string,
	body interface{}) (*Response, error) {
	response, err := es.HttpClient.PostJson(fmt.Sprintf("%s/%s/%s/_update?refresh=true",
		index, docType, docId), body)
	if err != nil {
		return nil, errors.Wrap(err, "http request failed")
	}
	defer response.Body.Close()
	if err := IsError(response); err != nil {
		return nil, err
	}
	return DecodeResponse(response)
}

// Refresh refreshes all indices logging any error but not returning and
// discarding the response so the caller doesn't have to.
func (es *ElasticSearch) Refresh() {
	response, err := es.HttpClient.PostString("_refresh", "application/json", "{}")
	if err != nil {
		log.Error("Failed to refresh Elastic Search: %v", err)
		return
	}
	io.Copy(ioutil.Discard, response.Body)
}

func (es *ElasticSearch) FormatKeyword(keyword string) string {
	if es.keyword == "" {
		return keyword
	}
	return fmt.Sprintf("%s.%s", keyword, es.keyword)
}

func (s *ElasticSearch) doUpdateByQuery(query interface{}) (util.JsonMap, error) {
	var response util.JsonMap
	err := s.HttpClient.PostJsonDecodeResponse(
		fmt.Sprintf("%s/_update_by_query?refresh=true&conflicts=proceed",
			s.EventSearchIndex), query, &response)
	if err != nil {
		return nil, errors.Wrap(err, "request failed")
	}

	return response, nil
}

// FormatTimestampUTC formats a time.Time into the format generally used
// by Elastic Search, in particular the @timestamp field.
func FormatTimestampUTC(timestamp time.Time) string {
	return timestamp.UTC().Format("2006-01-02T15:04:05.000Z")
}
<|endoftext|>


<commit_before>package volumedriver

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/akutz/gofig"
	"github.com/akutz/goof"
	"github.com/akutz/gotil"
	"github.com/emccode/libstorage"
	"github.com/emccode/libstorage/api/context"
	apitypes "github.com/emccode/libstorage/api/types"
	apiutils "github.com/emccode/libstorage/api/utils"

	"github.com/emccode/rexray/daemon/module"
)

const (
	modName = "docker"
)

type mod struct {
	lsc    apitypes.Client
	ctx    apitypes.Context
	config gofig.Config
	name   string
	addr   string
	desc   string
}

var (
	separators  = regexp.MustCompile(`[ &_=+:]`)
	dashes      = regexp.MustCompile(`[\-]+`)
	illegalPath = regexp.MustCompile(`[^[:alnum:]\~\-\./]`)
)

func init() {
	module.RegisterModule(modName, newModule)
}

func newModule(c *module.Config) (module.Module, error) {

	host := strings.Trim(c.Address, " ")

	if host == "" {
		if c.Name == "default-docker" {
			host = "unix:///run/docker/plugins/rexray.sock"
		} else {
			fname := cleanName(c.Name)
			host = fmt.Sprintf("unix:///run/docker/plugins/%s.sock", fname)
		}
	}

	c.Address = host

	/*cc, err := c.Config.Copy()
	if err != nil {
		return nil, err
	}*/

	// TODO The config should be copied, but it was causing an issue with
	// duplicate, explicit values not being removed.
	cc := c.Config

	if !cc.GetBool("rexray.volume.path.disableCache") {
		cc.Set("rexray.volume.path.cache", true)
	}

	ctx := context.Background()

	return &mod{
		ctx:    ctx,
		lsc:    c.Client,
		config: cc,
		name:   c.Name,
		desc:   c.Description,
		addr:   host,
	}, nil
}

func cleanName(s string) string {
	s = strings.Trim(strings.ToLower(s), " ")
	s = separators.ReplaceAllString(s, "-")
	s = illegalPath.ReplaceAllString(s, "")
	s = dashes.ReplaceAllString(s, "-")
	return s
}

const driverName = "docker"

var (
	errMissingHost      = goof.New("Missing host parameter")
	errBadHostSpecified = goof.New("Bad host specified, ie. unix:///run/docker/plugins/rexray.sock or tcp://127.0.0.1:8080")
	errBadProtocol      = goof.New("Bad protocol specified with host, ie. unix:// or tcp://")
)

type pluginRequest struct {
	Name string            `json:"Name,omitempty"`
	Opts map[string]string `json:"Opts,omitempty"`
}

func (m *mod) Start() error {

	lsc, _, err, _ := libstorage.New(m.config)

	if err != nil {
		for k, v := range m.config.AllSettings() {
			m.ctx.Errorf("%s=%v", k, v)
		}
		m.ctx.Fatal(err)
	}
	m.lsc = lsc

	proto, addr, parseAddrErr := gotil.ParseAddress(m.Address())
	if parseAddrErr != nil {
		return parseAddrErr
	}

	if proto == "unix" {
		dir := filepath.Dir(addr)
		os.MkdirAll(dir, 0755)
	}

	const validProtoPatt = "(?i)^unix|tcp$"
	isProtoValid, matchProtoErr := regexp.MatchString(validProtoPatt, proto)
	if matchProtoErr != nil {
		return goof.WithFieldsE(goof.Fields{
			"protocol":       proto,
			"validProtoPatt": validProtoPatt,
		}, "error matching protocol", matchProtoErr)
	}
	if !isProtoValid {
		return goof.WithField("protocol", proto, "invalid protocol")
	}

	if err := os.MkdirAll("/etc/docker/plugins", 0755); err != nil {
		return err
	}

	var specPath string
	var startFunc func() error

	mux := m.buildMux()

	if proto == "unix" {
		sockFile := addr
		sockFileDir := filepath.Dir(sockFile)
		mkSockFileDirErr := os.MkdirAll(sockFileDir, 0755)
		if mkSockFileDirErr != nil {
			return mkSockFileDirErr
		}

		_ = os.RemoveAll(sockFile)

		specPath = m.Address()
		startFunc = func() error {
			l, lErr := net.Listen("unix", sockFile)
			if lErr != nil {
				return lErr
			}
			defer l.Close()
			defer os.Remove(sockFile)

			return http.Serve(l, mux)
		}
	} else {
		specPath = addr
		startFunc = func() error {
			s := &http.Server{
				Addr:           addr,
				Handler:        mux,
				ReadTimeout:    10 * time.Second,
				WriteTimeout:   10 * time.Second,
				MaxHeaderBytes: 1 << 20,
			}
			return s.ListenAndServe()
		}
	}

	go func() {
		sErr := startFunc()
		if sErr != nil {
			panic(sErr)
		}
	}()

	spec := m.config.GetString("spec")
	if spec == "" {
		if m.name == "default-docker" {
			spec = "/etc/docker/plugins/rexray.spec"
		} else {
			fname := cleanName(m.name)
			spec = fmt.Sprintf("/etc/docker/plugins/%s.spec", fname)
		}
	}

	m.ctx.WithField("path", spec).Debug("docker voldriver spec file")

	if !gotil.FileExists(spec) {
		if err := ioutil.WriteFile(spec, []byte(specPath), 0644); err != nil {
			return err
		}
	}

	return nil
}

func (m *mod) Stop() error {
	return nil
}

func (m *mod) Name() string {
	return m.name
}

func (m *mod) Description() string {
	return m.desc
}

func (m *mod) Address() string {
	return m.addr
}

func (m *mod) buildMux() *http.ServeMux {

	mux := http.NewServeMux()
	// m.ctx.WithServiceName(m.ctx.ServiceName())

	mux.HandleFunc("/Plugin.Activate", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, `{"Implements": ["VolumeDriver"]}`)
		m.ctx.Debug("/VolumeDriver.Create")
	})

	mux.HandleFunc("/VolumeDriver.Create", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Create: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Create")
		store := apiutils.NewStoreWithVars(pr.Opts)
		vtype := store.GetStringPtr("type")
		if vtype == nil {
			vtype = store.GetStringPtr("volumetype")
		}
		_, err := m.lsc.Integration().Create(
			m.ctx,
			pr.Name,
			&apitypes.VolumeCreateOpts{
				AvailabilityZone: store.GetStringPtr("availabilityZone"),
				IOPS:             store.GetInt64Ptr("iops"),
				Size:             store.GetInt64Ptr("size"),
				Type:             vtype,
				Opts:             store,
			})

		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Create: error creating volume")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, `{}`)
	})

	mux.HandleFunc("/VolumeDriver.Remove", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Remove: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Remove")

		// TODO We need the service name
		err := m.lsc.Integration().Remove(m.ctx, pr.Name, apiutils.NewStore())
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Remove: error removing volume")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, `{}`)
	})

	mux.HandleFunc("/VolumeDriver.Path", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Path: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Path")

		mountPath, err := m.lsc.Integration().Path(
			m.ctx, "", pr.Name, apiutils.NewStore())
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Path: error returning path")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, fmt.Sprintf("{\"Mountpoint\": \"%s\"}", mountPath))
	})

	mux.HandleFunc("/VolumeDriver.Mount", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Mount: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Mount")

		mountPath, _, err := m.lsc.Integration().Mount(
			m.ctx, "", pr.Name, &apitypes.VolumeMountOpts{})
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Mount: error mounting volume")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, fmt.Sprintf("{\"Mountpoint\": \"%s\"}", mountPath))
	})

	mux.HandleFunc("/VolumeDriver.Unmount", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Unmount: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Unmount")

		err := m.lsc.Integration().Unmount(
			m.ctx, "", pr.Name, apiutils.NewStore())
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Unmount: error unmounting volume")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, `{}`)
	})

	mux.HandleFunc("/VolumeDriver.Get", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Path: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Get")

		volMapping, err := m.lsc.Integration().Inspect(
			m.ctx, pr.Name, apiutils.NewStore())
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Get: error getting volume")
			return
		}

		w.Header().Set(
			"Content-Type", "application/vnd.docker.plugins.v1.2+json")
		json.NewEncoder(w).Encode(map[string]apitypes.VolumeMapping{
			"Volume": volMapping,
		})
	})

	mux.HandleFunc("/VolumeDriver.List", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.List: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.List")

		volMappings, err := m.lsc.Integration().List(m.ctx, apiutils.NewStore())
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.List: error listing volumes")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		json.NewEncoder(w).Encode(
			map[string][]apitypes.VolumeMapping{"Volumes": volMappings})
	})

	return mux
}
<commit_msg>Fixed error channel instead of error on libstorage start<commit_after>package volumedriver

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/akutz/gofig"
	"github.com/akutz/goof"
	"github.com/akutz/gotil"
	"github.com/emccode/libstorage"
	"github.com/emccode/libstorage/api/context"
	apitypes "github.com/emccode/libstorage/api/types"
	apiutils "github.com/emccode/libstorage/api/utils"

	"github.com/emccode/rexray/daemon/module"
)

const (
	modName = "docker"
)

type mod struct {
	lsc    apitypes.Client
	ctx    apitypes.Context
	config gofig.Config
	name   string
	addr   string
	desc   string
}

var (
	separators  = regexp.MustCompile(`[ &_=+:]`)
	dashes      = regexp.MustCompile(`[\-]+`)
	illegalPath = regexp.MustCompile(`[^[:alnum:]\~\-\./]`)
)

func init() {
	module.RegisterModule(modName, newModule)
}

func newModule(c *module.Config) (module.Module, error) {

	host := strings.Trim(c.Address, " ")

	if host == "" {
		if c.Name == "default-docker" {
			host = "unix:///run/docker/plugins/rexray.sock"
		} else {
			fname := cleanName(c.Name)
			host = fmt.Sprintf("unix:///run/docker/plugins/%s.sock", fname)
		}
	}

	c.Address = host

	/*cc, err := c.Config.Copy()
	if err != nil {
		return nil, err
	}*/

	// TODO The config should be copied, but it was causing an issue with
	// duplicate, explicit values not being removed.
	cc := c.Config

	if !cc.GetBool("rexray.volume.path.disableCache") {
		cc.Set("rexray.volume.path.cache", true)
	}

	ctx := context.Background()

	return &mod{
		ctx:    ctx,
		lsc:    c.Client,
		config: cc,
		name:   c.Name,
		desc:   c.Description,
		addr:   host,
	}, nil
}

func cleanName(s string) string {
	s = strings.Trim(strings.ToLower(s), " ")
	s = separators.ReplaceAllString(s, "-")
	s = illegalPath.ReplaceAllString(s, "")
	s = dashes.ReplaceAllString(s, "-")
	return s
}

const driverName = "docker"

var (
	errMissingHost      = goof.New("Missing host parameter")
	errBadHostSpecified = goof.New("Bad host specified, ie. unix:///run/docker/plugins/rexray.sock or tcp://127.0.0.1:8080")
	errBadProtocol      = goof.New("Bad protocol specified with host, ie. unix:// or tcp://")
)

type pluginRequest struct {
	Name string            `json:"Name,omitempty"`
	Opts map[string]string `json:"Opts,omitempty"`
}

func (m *mod) Start() error {

	lsc, _, _, err := libstorage.New(m.config)
	if err != nil {
		for k, v := range m.config.AllSettings() {
			m.ctx.Errorf("%s=%v", k, v)
		}
		m.ctx.Fatal(err)
	}
	m.lsc = lsc

	proto, addr, parseAddrErr := gotil.ParseAddress(m.Address())
	if parseAddrErr != nil {
		return parseAddrErr
	}

	if proto == "unix" {
		dir := filepath.Dir(addr)
		os.MkdirAll(dir, 0755)
	}

	const validProtoPatt = "(?i)^unix|tcp$"
	isProtoValid, matchProtoErr := regexp.MatchString(validProtoPatt, proto)
	if matchProtoErr != nil {
		return goof.WithFieldsE(goof.Fields{
			"protocol":       proto,
			"validProtoPatt": validProtoPatt,
		}, "error matching protocol", matchProtoErr)
	}
	if !isProtoValid {
		return goof.WithField("protocol", proto, "invalid protocol")
	}

	if err := os.MkdirAll("/etc/docker/plugins", 0755); err != nil {
		return err
	}

	var specPath string
	var startFunc func() error

	mux := m.buildMux()

	if proto == "unix" {
		sockFile := addr
		sockFileDir := filepath.Dir(sockFile)
		mkSockFileDirErr := os.MkdirAll(sockFileDir, 0755)
		if mkSockFileDirErr != nil {
			return mkSockFileDirErr
		}

		_ = os.RemoveAll(sockFile)

		specPath = m.Address()
		startFunc = func() error {
			l, lErr := net.Listen("unix", sockFile)
			if lErr != nil {
				return lErr
			}
			defer l.Close()
			defer os.Remove(sockFile)

			return http.Serve(l, mux)
		}
	} else {
		specPath = addr
		startFunc = func() error {
			s := &http.Server{
				Addr:           addr,
				Handler:        mux,
				ReadTimeout:    10 * time.Second,
				WriteTimeout:   10 * time.Second,
				MaxHeaderBytes: 1 << 20,
			}
			return s.ListenAndServe()
		}
	}

	go func() {
		sErr := startFunc()
		if sErr != nil {
			panic(sErr)
		}
	}()

	spec := m.config.GetString("spec")
	if spec == "" {
		if m.name == "default-docker" {
			spec = "/etc/docker/plugins/rexray.spec"
		} else {
			fname := cleanName(m.name)
			spec = fmt.Sprintf("/etc/docker/plugins/%s.spec", fname)
		}
	}

	m.ctx.WithField("path", spec).Debug("docker voldriver spec file")

	if !gotil.FileExists(spec) {
		if err := ioutil.WriteFile(spec, []byte(specPath), 0644); err != nil {
			return err
		}
	}

	return nil
}

func (m *mod) Stop() error {
	return nil
}

func (m *mod) Name() string {
	return m.name
}

func (m *mod) Description() string {
	return m.desc
}

func (m *mod) Address() string {
	return m.addr
}

func (m *mod) buildMux() *http.ServeMux {

	mux := http.NewServeMux()
	// m.ctx.WithServiceName(m.ctx.ServiceName())

	mux.HandleFunc("/Plugin.Activate", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, `{"Implements": ["VolumeDriver"]}`)
		m.ctx.Debug("/VolumeDriver.Create")
	})

	mux.HandleFunc("/VolumeDriver.Create", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Create: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Create")
		store := apiutils.NewStoreWithVars(pr.Opts)
		vtype := store.GetStringPtr("type")
		if vtype == nil {
			vtype = store.GetStringPtr("volumetype")
		}
		_, err := m.lsc.Integration().Create(
			m.ctx,
			pr.Name,
			&apitypes.VolumeCreateOpts{
				AvailabilityZone: store.GetStringPtr("availabilityZone"),
				IOPS:             store.GetInt64Ptr("iops"),
				Size:             store.GetInt64Ptr("size"),
				Type:             vtype,
				Opts:             store,
			})

		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Create: error creating volume")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, `{}`)
	})

	mux.HandleFunc("/VolumeDriver.Remove", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Remove: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Remove")

		// TODO We need the service name
		err := m.lsc.Integration().Remove(m.ctx, pr.Name, apiutils.NewStore())
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Remove: error removing volume")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, `{}`)
	})

	mux.HandleFunc("/VolumeDriver.Path", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Path: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Path")

		mountPath, err := m.lsc.Integration().Path(
			m.ctx, "", pr.Name, apiutils.NewStore())
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Path: error returning path")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, fmt.Sprintf("{\"Mountpoint\": \"%s\"}", mountPath))
	})

	mux.HandleFunc("/VolumeDriver.Mount", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Mount: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Mount")

		mountPath, _, err := m.lsc.Integration().Mount(
			m.ctx, "", pr.Name, &apitypes.VolumeMountOpts{})
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Mount: error mounting volume")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, fmt.Sprintf("{\"Mountpoint\": \"%s\"}", mountPath))
	})

	mux.HandleFunc("/VolumeDriver.Unmount", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Unmount: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Unmount")

		err := m.lsc.Integration().Unmount(
			m.ctx, "", pr.Name, apiutils.NewStore())
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Unmount: error unmounting volume")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		fmt.Fprintln(w, `{}`)
	})

	mux.HandleFunc("/VolumeDriver.Get", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Path: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.Get")

		volMapping, err := m.lsc.Integration().Inspect(
			m.ctx, pr.Name, apiutils.NewStore())
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.Get: error getting volume")
			return
		}

		w.Header().Set(
			"Content-Type", "application/vnd.docker.plugins.v1.2+json")
		json.NewEncoder(w).Encode(map[string]apitypes.VolumeMapping{
			"Volume": volMapping,
		})
	})

	mux.HandleFunc("/VolumeDriver.List", func(w http.ResponseWriter, r *http.Request) {
		var pr pluginRequest
		if err := json.NewDecoder(r.Body).Decode(&pr); err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.List: error decoding json")
			return
		}

		m.ctx.WithField("pluginResponse", pr).Debug("/VolumeDriver.List")

		volMappings, err := m.lsc.Integration().List(m.ctx, apiutils.NewStore())
		if err != nil {
			http.Error(w, fmt.Sprintf("{\"Error\":\"%s\"}", err.Error()), 500)
			m.ctx.WithError(err).Error("/VolumeDriver.List: error listing volumes")
			return
		}

		w.Header().Set("Content-Type", "application/vnd.docker.plugins.v1.2+json")
		json.NewEncoder(w).Encode(
			map[string][]apitypes.VolumeMapping{"Volumes": volMappings})
	})

	return mux
}
<|endoftext|>


<commit_before>package testrunner

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/onsi/ginkgo/config"
	"github.com/onsi/ginkgo/ginkgo/testsuite"
	"github.com/onsi/ginkgo/internal/remote"
	"github.com/onsi/ginkgo/reporters/stenographer"
	"github.com/onsi/ginkgo/types"
)

type TestRunner struct {
	Suite testsuite.TestSuite

	compiled              bool
	compilationTargetPath string

	numCPU         int
	parallelStream bool
	race           bool
	cover          bool
	coverPkg       string
	tags           string
	additionalArgs []string
}

func New(suite testsuite.TestSuite, numCPU int, parallelStream bool, race bool, cover bool, coverPkg string, tags string, additionalArgs []string) *TestRunner {
	runner := &TestRunner{
		Suite:          suite,
		numCPU:         numCPU,
		parallelStream: parallelStream,
		race:           race,
		cover:          cover,
		coverPkg:       coverPkg,
		tags:           tags,
		additionalArgs: additionalArgs,
	}

	if !suite.Precompiled {
		dir, err := ioutil.TempDir("", "ginkgo")
		if err != nil {
			panic(fmt.Sprintf("coulnd't create temporary directory... might be time to rm -rf:\n%s", err.Error()))
		}
		runner.compilationTargetPath = filepath.Join(dir, suite.PackageName+".test")
	}

	return runner
}

func (t *TestRunner) Compile() error {
	return t.CompileTo(t.compilationTargetPath)
}

func (t *TestRunner) CompileTo(path string) error {
	if t.compiled {
		return nil
	}

	if t.Suite.Precompiled {
		return nil
	}

	args := []string{"test", "-c", "-i", "-o", path, t.Suite.Path}
	if t.race {
		args = append(args, "-race")
	}
	if t.cover || t.coverPkg != "" {
		args = append(args, "-cover", "-covermode=atomic")
	}
	if t.coverPkg != "" {
		args = append(args, fmt.Sprintf("-coverpkg=%s", t.coverPkg))
	}
	if t.tags != "" {
		args = append(args, fmt.Sprintf("-tags=%s", t.tags))
	}

	cmd := exec.Command("go", args...)

	output, err := cmd.CombinedOutput()

	if err != nil {
		fixedOutput := fixCompilationOutput(string(output), t.Suite.Path)
		if len(output) > 0 {
			return fmt.Errorf("Failed to compile %s:\n\n%s", t.Suite.PackageName, fixedOutput)
		}
		return fmt.Errorf("Failed to compile %s", t.Suite.PackageName)
	}

	if fileExists(path) == false {
		compiledFile := filepath.Join(t.Suite.Path, t.Suite.PackageName+".test")
		if fileExists(compiledFile) {
			// seems like we are on an old go version that does not support the -o flag on go test
			// move the compiled test file to the desired location by hand
			err = os.Rename(compiledFile, path)
			if err != nil {
				// We cannot move the file, perhaps because the source and destination
				// are on different partitions. We can copy the file, however.
				err = copyFile(compiledFile, path)
				if err != nil {
					return fmt.Errorf("Failed to copy compiled file: %s", err)
				}
			}
		} else {
			return fmt.Errorf("Failed to compile %s: output file %q could not be found", t.Suite.PackageName, path)
		}
	}

	t.compiled = true

	return nil
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil || os.IsNotExist(err) == false
}

// copyFile copies the contents of the file named src to the file named
// by dst. The file will be created if it does not already exist. If the
// destination file exists, all it's contents will be replaced by the contents
// of the source file.
func copyFile(src, dst string) error {
	srcInfo, err := os.Stat(src)
	if err != nil {
		return err
	}
	mode := srcInfo.Mode()

	in, err := os.Open(src)
	if err != nil {
		return err
	}

	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}

	defer func() {
		closeErr := out.Close()
		if err == nil {
			err = closeErr
		}
	}()

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}

	err = out.Sync()
	if err != nil {
		return err
	}

	return out.Chmod(mode)
}

/*
go test -c -i spits package.test out into the cwd. there's no way to change this.

to make sure it doesn't generate conflicting .test files in the cwd, Compile() must switch the cwd to the test package.

unfortunately, this causes go test's compile output to be expressed *relative to the test package* instead of the cwd.

this makes it hard to reason about what failed, and also prevents iterm's Cmd+click from working.

fixCompilationOutput..... rewrites the output to fix the paths.

yeah......
*/
func fixCompilationOutput(output string, relToPath string) string {
	re := regexp.MustCompile(`^(\S.*\.go)\:\d+\:`)
	lines := strings.Split(output, "\n")
	for i, line := range lines {
		indices := re.FindStringSubmatchIndex(line)
		if len(indices) == 0 {
			continue
		}

		path := line[indices[2]:indices[3]]
		path = filepath.Join(relToPath, path)
		lines[i] = path + line[indices[3]:]
	}
	return strings.Join(lines, "\n")
}

func (t *TestRunner) Run() RunResult {
	if t.Suite.IsGinkgo {
		if t.numCPU > 1 {
			if t.parallelStream {
				return t.runAndStreamParallelGinkgoSuite()
			} else {
				return t.runParallelGinkgoSuite()
			}
		} else {
			return t.runSerialGinkgoSuite()
		}
	} else {
		return t.runGoTestSuite()
	}
}

func (t *TestRunner) CleanUp() {
	if t.Suite.Precompiled {
		return
	}
	os.RemoveAll(filepath.Dir(t.compilationTargetPath))
}

func (t *TestRunner) runSerialGinkgoSuite() RunResult {
	ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)
	return t.run(t.cmd(ginkgoArgs, os.Stdout, 1), nil)
}

func (t *TestRunner) runGoTestSuite() RunResult {
	return t.run(t.cmd([]string{"-test.v"}, os.Stdout, 1), nil)
}

func (t *TestRunner) runAndStreamParallelGinkgoSuite() RunResult {
	completions := make(chan RunResult)
	writers := make([]*logWriter, t.numCPU)

	server, err := remote.NewServer(t.numCPU)
	if err != nil {
		panic("Failed to start parallel spec server")
	}

	server.Start()
	defer server.Close()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU
		config.GinkgoConfig.SyncHost = server.Address()

		ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)

		writers[cpu] = newLogWriter(os.Stdout, cpu+1)

		cmd := t.cmd(ginkgoArgs, writers[cpu], cpu+1)

		server.RegisterAlive(cpu+1, func() bool {
			if cmd.ProcessState == nil {
				return true
			}
			return !cmd.ProcessState.Exited()
		})

		go t.run(cmd, completions)
	}

	res := PassingRunResult()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		res = res.Merge(<-completions)
	}

	for _, writer := range writers {
		writer.Close()
	}

	os.Stdout.Sync()

	if t.cover || t.coverPkg != "" {
		t.combineCoverprofiles()
	}

	return res
}

func (t *TestRunner) runParallelGinkgoSuite() RunResult {
	result := make(chan bool)
	completions := make(chan RunResult)
	writers := make([]*logWriter, t.numCPU)
	reports := make([]*bytes.Buffer, t.numCPU)

	stenographer := stenographer.New(!config.DefaultReporterConfig.NoColor)
	aggregator := remote.NewAggregator(t.numCPU, result, config.DefaultReporterConfig, stenographer)

	server, err := remote.NewServer(t.numCPU)
	if err != nil {
		panic("Failed to start parallel spec server")
	}
	server.RegisterReporters(aggregator)
	server.Start()
	defer server.Close()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU
		config.GinkgoConfig.SyncHost = server.Address()
		config.GinkgoConfig.StreamHost = server.Address()

		ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)

		reports[cpu] = &bytes.Buffer{}
		writers[cpu] = newLogWriter(reports[cpu], cpu+1)

		cmd := t.cmd(ginkgoArgs, writers[cpu], cpu+1)

		server.RegisterAlive(cpu+1, func() bool {
			if cmd.ProcessState == nil {
				return true
			}
			return !cmd.ProcessState.Exited()
		})

		go t.run(cmd, completions)
	}

	res := PassingRunResult()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		res = res.Merge(<-completions)
	}

	//all test processes are done, at this point
	//we should be able to wait for the aggregator to tell us that it's done

	select {
	case <-result:
		fmt.Println("")
	case <-time.After(time.Second):
		//the aggregator never got back to us!  something must have gone wrong
		fmt.Println(`
	 -------------------------------------------------------------------
	|                                                                   |
	|  Ginkgo timed out waiting for all parallel nodes to report back!  |
	|                                                                   |
	 -------------------------------------------------------------------
`)

		os.Stdout.Sync()

		for _, writer := range writers {
			writer.Close()
		}

		for _, report := range reports {
			fmt.Print(report.String())
		}

		os.Stdout.Sync()
	}

	if t.cover || t.coverPkg != "" {
		t.combineCoverprofiles()
	}

	return res
}

func (t *TestRunner) cmd(ginkgoArgs []string, stream io.Writer, node int) *exec.Cmd {
	args := []string{"--test.timeout=24h"}
	if t.cover || t.coverPkg != "" {
		coverprofile := "--test.coverprofile=" + t.Suite.PackageName + ".coverprofile"
		if t.numCPU > 1 {
			coverprofile = fmt.Sprintf("%s.%d", coverprofile, node)
		}
		args = append(args, coverprofile)
	}

	args = append(args, ginkgoArgs...)
	args = append(args, t.additionalArgs...)

	path := t.compilationTargetPath
	if t.Suite.Precompiled {
		path, _ = filepath.Abs(filepath.Join(t.Suite.Path, fmt.Sprintf("%s.test", t.Suite.PackageName)))
	}

	cmd := exec.Command(path, args...)

	cmd.Dir = t.Suite.Path
	cmd.Stderr = stream
	cmd.Stdout = stream

	return cmd
}

func (t *TestRunner) run(cmd *exec.Cmd, completions chan RunResult) RunResult {
	var res RunResult

	defer func() {
		if completions != nil {
			completions <- res
		}
	}()

	err := cmd.Start()
	if err != nil {
		fmt.Printf("Failed to run test suite!\n\t%s", err.Error())
		return res
	}

	cmd.Wait()
	exitStatus := cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus()
	res.Passed = (exitStatus == 0) || (exitStatus == types.GINKGO_FOCUS_EXIT_CODE)
	res.HasProgrammaticFocus = (exitStatus == types.GINKGO_FOCUS_EXIT_CODE)

	return res
}

func (t *TestRunner) combineCoverprofiles() {
	profiles := []string{}
	for cpu := 1; cpu <= t.numCPU; cpu++ {
		coverFile := fmt.Sprintf("%s.coverprofile.%d", t.Suite.PackageName, cpu)
		coverFile = filepath.Join(t.Suite.Path, coverFile)
		coverProfile, err := ioutil.ReadFile(coverFile)
		os.Remove(coverFile)

		if err == nil {
			profiles = append(profiles, string(coverProfile))
		}
	}

	if len(profiles) != t.numCPU {
		return
	}

	lines := map[string]int{}
	lineOrder := []string{}
	for i, coverProfile := range profiles {
		for _, line := range strings.Split(string(coverProfile), "\n")[1:] {
			if len(line) == 0 {
				continue
			}
			components := strings.Split(line, " ")
			count, _ := strconv.Atoi(components[len(components)-1])
			prefix := strings.Join(components[0:len(components)-1], " ")
			lines[prefix] += count
			if i == 0 {
				lineOrder = append(lineOrder, prefix)
			}
		}
	}

	output := []string{"mode: atomic"}
	for _, line := range lineOrder {
		output = append(output, fmt.Sprintf("%s %d", line, lines[line]))
	}
	finalOutput := strings.Join(output, "\n")
	ioutil.WriteFile(filepath.Join(t.Suite.Path, fmt.Sprintf("%s.coverprofile", t.Suite.PackageName)), []byte(finalOutput), 0666)
}
<commit_msg>fix compilation on older versions of go<commit_after>package testrunner

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/onsi/ginkgo/config"
	"github.com/onsi/ginkgo/ginkgo/testsuite"
	"github.com/onsi/ginkgo/internal/remote"
	"github.com/onsi/ginkgo/reporters/stenographer"
	"github.com/onsi/ginkgo/types"
)

type TestRunner struct {
	Suite testsuite.TestSuite

	compiled              bool
	compilationTargetPath string

	numCPU         int
	parallelStream bool
	race           bool
	cover          bool
	coverPkg       string
	tags           string
	additionalArgs []string
}

func New(suite testsuite.TestSuite, numCPU int, parallelStream bool, race bool, cover bool, coverPkg string, tags string, additionalArgs []string) *TestRunner {
	runner := &TestRunner{
		Suite:          suite,
		numCPU:         numCPU,
		parallelStream: parallelStream,
		race:           race,
		cover:          cover,
		coverPkg:       coverPkg,
		tags:           tags,
		additionalArgs: additionalArgs,
	}

	if !suite.Precompiled {
		dir, err := ioutil.TempDir("", "ginkgo")
		if err != nil {
			panic(fmt.Sprintf("coulnd't create temporary directory... might be time to rm -rf:\n%s", err.Error()))
		}
		runner.compilationTargetPath = filepath.Join(dir, suite.PackageName+".test")
	}

	return runner
}

func (t *TestRunner) Compile() error {
	return t.CompileTo(t.compilationTargetPath)
}

func (t *TestRunner) CompileTo(path string) error {
	if t.compiled {
		return nil
	}

	if t.Suite.Precompiled {
		return nil
	}

	args := []string{"test", "-c", "-i", "-o", path, t.Suite.Path}
	if t.race {
		args = append(args, "-race")
	}
	if t.cover || t.coverPkg != "" {
		args = append(args, "-cover", "-covermode=atomic")
	}
	if t.coverPkg != "" {
		args = append(args, fmt.Sprintf("-coverpkg=%s", t.coverPkg))
	}
	if t.tags != "" {
		args = append(args, fmt.Sprintf("-tags=%s", t.tags))
	}

	cmd := exec.Command("go", args...)

	output, err := cmd.CombinedOutput()

	if err != nil {
		fixedOutput := fixCompilationOutput(string(output), t.Suite.Path)
		if len(output) > 0 {
			return fmt.Errorf("Failed to compile %s:\n\n%s", t.Suite.PackageName, fixedOutput)
		}
		return fmt.Errorf("Failed to compile %s", t.Suite.PackageName)
	}

	if fileExists(path) == false {
		compiledFile := t.Suite.PackageName + ".test"
		if fileExists(compiledFile) {
			// seems like we are on an old go version that does not support the -o flag on go test
			// move the compiled test file to the desired location by hand
			err = os.Rename(compiledFile, path)
			if err != nil {
				// We cannot move the file, perhaps because the source and destination
				// are on different partitions. We can copy the file, however.
				err = copyFile(compiledFile, path)
				if err != nil {
					return fmt.Errorf("Failed to copy compiled file: %s", err)
				}
			}
		} else {
			return fmt.Errorf("Failed to compile %s: output file %q could not be found", t.Suite.PackageName, path)
		}
	}

	t.compiled = true

	return nil
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil || os.IsNotExist(err) == false
}

// copyFile copies the contents of the file named src to the file named
// by dst. The file will be created if it does not already exist. If the
// destination file exists, all it's contents will be replaced by the contents
// of the source file.
func copyFile(src, dst string) error {
	srcInfo, err := os.Stat(src)
	if err != nil {
		return err
	}
	mode := srcInfo.Mode()

	in, err := os.Open(src)
	if err != nil {
		return err
	}

	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}

	defer func() {
		closeErr := out.Close()
		if err == nil {
			err = closeErr
		}
	}()

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}

	err = out.Sync()
	if err != nil {
		return err
	}

	return out.Chmod(mode)
}

/*
go test -c -i spits package.test out into the cwd. there's no way to change this.

to make sure it doesn't generate conflicting .test files in the cwd, Compile() must switch the cwd to the test package.

unfortunately, this causes go test's compile output to be expressed *relative to the test package* instead of the cwd.

this makes it hard to reason about what failed, and also prevents iterm's Cmd+click from working.

fixCompilationOutput..... rewrites the output to fix the paths.

yeah......
*/
func fixCompilationOutput(output string, relToPath string) string {
	re := regexp.MustCompile(`^(\S.*\.go)\:\d+\:`)
	lines := strings.Split(output, "\n")
	for i, line := range lines {
		indices := re.FindStringSubmatchIndex(line)
		if len(indices) == 0 {
			continue
		}

		path := line[indices[2]:indices[3]]
		path = filepath.Join(relToPath, path)
		lines[i] = path + line[indices[3]:]
	}
	return strings.Join(lines, "\n")
}

func (t *TestRunner) Run() RunResult {
	if t.Suite.IsGinkgo {
		if t.numCPU > 1 {
			if t.parallelStream {
				return t.runAndStreamParallelGinkgoSuite()
			} else {
				return t.runParallelGinkgoSuite()
			}
		} else {
			return t.runSerialGinkgoSuite()
		}
	} else {
		return t.runGoTestSuite()
	}
}

func (t *TestRunner) CleanUp() {
	if t.Suite.Precompiled {
		return
	}
	os.RemoveAll(filepath.Dir(t.compilationTargetPath))
}

func (t *TestRunner) runSerialGinkgoSuite() RunResult {
	ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)
	return t.run(t.cmd(ginkgoArgs, os.Stdout, 1), nil)
}

func (t *TestRunner) runGoTestSuite() RunResult {
	return t.run(t.cmd([]string{"-test.v"}, os.Stdout, 1), nil)
}

func (t *TestRunner) runAndStreamParallelGinkgoSuite() RunResult {
	completions := make(chan RunResult)
	writers := make([]*logWriter, t.numCPU)

	server, err := remote.NewServer(t.numCPU)
	if err != nil {
		panic("Failed to start parallel spec server")
	}

	server.Start()
	defer server.Close()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU
		config.GinkgoConfig.SyncHost = server.Address()

		ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)

		writers[cpu] = newLogWriter(os.Stdout, cpu+1)

		cmd := t.cmd(ginkgoArgs, writers[cpu], cpu+1)

		server.RegisterAlive(cpu+1, func() bool {
			if cmd.ProcessState == nil {
				return true
			}
			return !cmd.ProcessState.Exited()
		})

		go t.run(cmd, completions)
	}

	res := PassingRunResult()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		res = res.Merge(<-completions)
	}

	for _, writer := range writers {
		writer.Close()
	}

	os.Stdout.Sync()

	if t.cover || t.coverPkg != "" {
		t.combineCoverprofiles()
	}

	return res
}

func (t *TestRunner) runParallelGinkgoSuite() RunResult {
	result := make(chan bool)
	completions := make(chan RunResult)
	writers := make([]*logWriter, t.numCPU)
	reports := make([]*bytes.Buffer, t.numCPU)

	stenographer := stenographer.New(!config.DefaultReporterConfig.NoColor)
	aggregator := remote.NewAggregator(t.numCPU, result, config.DefaultReporterConfig, stenographer)

	server, err := remote.NewServer(t.numCPU)
	if err != nil {
		panic("Failed to start parallel spec server")
	}
	server.RegisterReporters(aggregator)
	server.Start()
	defer server.Close()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU
		config.GinkgoConfig.SyncHost = server.Address()
		config.GinkgoConfig.StreamHost = server.Address()

		ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)

		reports[cpu] = &bytes.Buffer{}
		writers[cpu] = newLogWriter(reports[cpu], cpu+1)

		cmd := t.cmd(ginkgoArgs, writers[cpu], cpu+1)

		server.RegisterAlive(cpu+1, func() bool {
			if cmd.ProcessState == nil {
				return true
			}
			return !cmd.ProcessState.Exited()
		})

		go t.run(cmd, completions)
	}

	res := PassingRunResult()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		res = res.Merge(<-completions)
	}

	//all test processes are done, at this point
	//we should be able to wait for the aggregator to tell us that it's done

	select {
	case <-result:
		fmt.Println("")
	case <-time.After(time.Second):
		//the aggregator never got back to us!  something must have gone wrong
		fmt.Println(`
	 -------------------------------------------------------------------
	|                                                                   |
	|  Ginkgo timed out waiting for all parallel nodes to report back!  |
	|                                                                   |
	 -------------------------------------------------------------------
`)

		os.Stdout.Sync()

		for _, writer := range writers {
			writer.Close()
		}

		for _, report := range reports {
			fmt.Print(report.String())
		}

		os.Stdout.Sync()
	}

	if t.cover || t.coverPkg != "" {
		t.combineCoverprofiles()
	}

	return res
}

func (t *TestRunner) cmd(ginkgoArgs []string, stream io.Writer, node int) *exec.Cmd {
	args := []string{"--test.timeout=24h"}
	if t.cover || t.coverPkg != "" {
		coverprofile := "--test.coverprofile=" + t.Suite.PackageName + ".coverprofile"
		if t.numCPU > 1 {
			coverprofile = fmt.Sprintf("%s.%d", coverprofile, node)
		}
		args = append(args, coverprofile)
	}

	args = append(args, ginkgoArgs...)
	args = append(args, t.additionalArgs...)

	path := t.compilationTargetPath
	if t.Suite.Precompiled {
		path, _ = filepath.Abs(filepath.Join(t.Suite.Path, fmt.Sprintf("%s.test", t.Suite.PackageName)))
	}

	cmd := exec.Command(path, args...)

	cmd.Dir = t.Suite.Path
	cmd.Stderr = stream
	cmd.Stdout = stream

	return cmd
}

func (t *TestRunner) run(cmd *exec.Cmd, completions chan RunResult) RunResult {
	var res RunResult

	defer func() {
		if completions != nil {
			completions <- res
		}
	}()

	err := cmd.Start()
	if err != nil {
		fmt.Printf("Failed to run test suite!\n\t%s", err.Error())
		return res
	}

	cmd.Wait()
	exitStatus := cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus()
	res.Passed = (exitStatus == 0) || (exitStatus == types.GINKGO_FOCUS_EXIT_CODE)
	res.HasProgrammaticFocus = (exitStatus == types.GINKGO_FOCUS_EXIT_CODE)

	return res
}

func (t *TestRunner) combineCoverprofiles() {
	profiles := []string{}
	for cpu := 1; cpu <= t.numCPU; cpu++ {
		coverFile := fmt.Sprintf("%s.coverprofile.%d", t.Suite.PackageName, cpu)
		coverFile = filepath.Join(t.Suite.Path, coverFile)
		coverProfile, err := ioutil.ReadFile(coverFile)
		os.Remove(coverFile)

		if err == nil {
			profiles = append(profiles, string(coverProfile))
		}
	}

	if len(profiles) != t.numCPU {
		return
	}

	lines := map[string]int{}
	lineOrder := []string{}
	for i, coverProfile := range profiles {
		for _, line := range strings.Split(string(coverProfile), "\n")[1:] {
			if len(line) == 0 {
				continue
			}
			components := strings.Split(line, " ")
			count, _ := strconv.Atoi(components[len(components)-1])
			prefix := strings.Join(components[0:len(components)-1], " ")
			lines[prefix] += count
			if i == 0 {
				lineOrder = append(lineOrder, prefix)
			}
		}
	}

	output := []string{"mode: atomic"}
	for _, line := range lineOrder {
		output = append(output, fmt.Sprintf("%s %d", line, lines[line]))
	}
	finalOutput := strings.Join(output, "\n")
	ioutil.WriteFile(filepath.Join(t.Suite.Path, fmt.Sprintf("%s.coverprofile", t.Suite.PackageName)), []byte(finalOutput), 0666)
}
<|endoftext|>


<commit_before>// Copyright 2011 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.

package gcimporter

import (
	"go/build"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"code.google.com/p/go.tools/go/types"
)

var gcPath string // Go compiler path

func init() {
	// determine compiler
	var gc string
	switch runtime.GOARCH {
	case "386":
		gc = "8g"
	case "amd64":
		gc = "6g"
	case "arm":
		gc = "5g"
	default:
		gcPath = "unknown-GOARCH-compiler"
		return
	}
	gcPath = filepath.Join(build.ToolDir, gc)
}

func compile(t *testing.T, dirname, filename string) string {
	cmd := exec.Command(gcPath, filename)
	cmd.Dir = dirname
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Logf("%s", out)
		t.Fatalf("%s %s failed: %s", gcPath, filename, err)
	}
	archCh, _ := build.ArchChar(runtime.GOARCH)
	// filename should end with ".go"
	return filepath.Join(dirname, filename[:len(filename)-2]+archCh)
}

// Use the same global imports map for all tests. The effect is
// as if all tested packages were imported into a single package.
var imports = make(map[string]*types.Package)

func testPath(t *testing.T, path string) bool {
	t0 := time.Now()
	_, err := Import(imports, path)
	if err != nil {
		t.Errorf("testPath(%s): %s", path, err)
		return false
	}
	t.Logf("testPath(%s): %v", path, time.Since(t0))
	return true
}

const maxTime = 30 * time.Second

func testDir(t *testing.T, dir string, endTime time.Time) (nimports int) {
	dirname := filepath.Join(runtime.GOROOT(), "pkg", runtime.GOOS+"_"+runtime.GOARCH, dir)
	list, err := ioutil.ReadDir(dirname)
	if err != nil {
		t.Fatalf("testDir(%s): %s", dirname, err)
	}
	for _, f := range list {
		if time.Now().After(endTime) {
			t.Log("testing time used up")
			return
		}
		switch {
		case !f.IsDir():
			// try extensions
			for _, ext := range pkgExts {
				if strings.HasSuffix(f.Name(), ext) {
					name := f.Name()[0 : len(f.Name())-len(ext)] // remove extension
					if testPath(t, filepath.Join(dir, name)) {
						nimports++
					}
				}
			}
		case f.IsDir():
			nimports += testDir(t, filepath.Join(dir, f.Name()), endTime)
		}
	}
	return
}

func TestImport(t *testing.T) {
	// This package does not handle gccgo export data.
	if runtime.Compiler == "gccgo" {
		return
	}

	// On cross-compile builds, the path will not exist.
	// Need to use GOHOSTOS, which is not available.
	if _, err := os.Stat(gcPath); err != nil {
		t.Skipf("skipping test: %v", err)
	}

	if outFn := compile(t, "testdata", "exports.go"); outFn != "" {
		defer os.Remove(outFn)
	}

	nimports := 0
	if testPath(t, "./testdata/exports") {
		nimports++
	}
	nimports += testDir(t, "", time.Now().Add(maxTime)) // installed packages
	t.Logf("tested %d imports", nimports)
}

var importedObjectTests = []struct {
	name string
	want string
}{
	{"unsafe.Pointer", "type Pointer unsafe.Pointer"},
	{"math.Pi", "const Pi untyped float"},
	{"io.Reader", "type Reader interface{Read(p []byte) (n int, err error)}"},
	{"io.ReadWriter", "type ReadWriter interface{Read(p []byte) (n int, err error); Write(p []byte) (n int, err error)}"},
	{"math.Sin", "func math.Sin(x·2 float64) (_ float64)"},
	// TODO(gri) add more tests
}

func TestImportedTypes(t *testing.T) {
	// This package does not handle gccgo export data.
	if runtime.Compiler == "gccgo" {
		return
	}
	for _, test := range importedObjectTests {
		s := strings.Split(test.name, ".")
		if len(s) != 2 {
			t.Fatal("inconsistent test data")
		}
		importPath := s[0]
		objName := s[1]

		pkg, err := Import(imports, importPath)
		if err != nil {
			t.Error(err)
			continue
		}

		obj := pkg.Scope().Lookup(objName)
		if obj == nil {
			t.Errorf("%s: object not found", test.name)
			continue
		}

		got := obj.String()
		if got != test.want {
			t.Errorf("%s: got %q; want %q", test.name, got, test.want)
		}
	}
}

func TestIssue5815(t *testing.T) {
	// This package does not handle gccgo export data.
	if runtime.Compiler == "gccgo" {
		return
	}

	pkg, err := Import(make(map[string]*types.Package), "strings")
	if err != nil {
		t.Fatal(err)
	}

	scope := pkg.Scope()
	for _, name := range scope.Names() {
		obj := scope.Lookup(name)
		if obj.Pkg() == nil {
			t.Errorf("no pkg for %s", obj)
		}
		if tname, _ := obj.(*types.TypeName); tname != nil {
			named := tname.Type().(*types.Named)
			for i := 0; i < named.NumMethods(); i++ {
				m := named.Method(i)
				if m.Pkg() == nil {
					t.Errorf("no pkg for %s", m)
				}
			}
		}
	}
}
<commit_msg>go.tools/go/gcimporter: fix build by adjusting test to new print output<commit_after>// Copyright 2011 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.

package gcimporter

import (
	"go/build"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"code.google.com/p/go.tools/go/types"
)

var gcPath string // Go compiler path

func init() {
	// determine compiler
	var gc string
	switch runtime.GOARCH {
	case "386":
		gc = "8g"
	case "amd64":
		gc = "6g"
	case "arm":
		gc = "5g"
	default:
		gcPath = "unknown-GOARCH-compiler"
		return
	}
	gcPath = filepath.Join(build.ToolDir, gc)
}

func compile(t *testing.T, dirname, filename string) string {
	cmd := exec.Command(gcPath, filename)
	cmd.Dir = dirname
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Logf("%s", out)
		t.Fatalf("%s %s failed: %s", gcPath, filename, err)
	}
	archCh, _ := build.ArchChar(runtime.GOARCH)
	// filename should end with ".go"
	return filepath.Join(dirname, filename[:len(filename)-2]+archCh)
}

// Use the same global imports map for all tests. The effect is
// as if all tested packages were imported into a single package.
var imports = make(map[string]*types.Package)

func testPath(t *testing.T, path string) bool {
	t0 := time.Now()
	_, err := Import(imports, path)
	if err != nil {
		t.Errorf("testPath(%s): %s", path, err)
		return false
	}
	t.Logf("testPath(%s): %v", path, time.Since(t0))
	return true
}

const maxTime = 30 * time.Second

func testDir(t *testing.T, dir string, endTime time.Time) (nimports int) {
	dirname := filepath.Join(runtime.GOROOT(), "pkg", runtime.GOOS+"_"+runtime.GOARCH, dir)
	list, err := ioutil.ReadDir(dirname)
	if err != nil {
		t.Fatalf("testDir(%s): %s", dirname, err)
	}
	for _, f := range list {
		if time.Now().After(endTime) {
			t.Log("testing time used up")
			return
		}
		switch {
		case !f.IsDir():
			// try extensions
			for _, ext := range pkgExts {
				if strings.HasSuffix(f.Name(), ext) {
					name := f.Name()[0 : len(f.Name())-len(ext)] // remove extension
					if testPath(t, filepath.Join(dir, name)) {
						nimports++
					}
				}
			}
		case f.IsDir():
			nimports += testDir(t, filepath.Join(dir, f.Name()), endTime)
		}
	}
	return
}

func TestImport(t *testing.T) {
	// This package does not handle gccgo export data.
	if runtime.Compiler == "gccgo" {
		return
	}

	// On cross-compile builds, the path will not exist.
	// Need to use GOHOSTOS, which is not available.
	if _, err := os.Stat(gcPath); err != nil {
		t.Skipf("skipping test: %v", err)
	}

	if outFn := compile(t, "testdata", "exports.go"); outFn != "" {
		defer os.Remove(outFn)
	}

	nimports := 0
	if testPath(t, "./testdata/exports") {
		nimports++
	}
	nimports += testDir(t, "", time.Now().Add(maxTime)) // installed packages
	t.Logf("tested %d imports", nimports)
}

var importedObjectTests = []struct {
	name string
	want string
}{
	{"unsafe.Pointer", "type Pointer unsafe.Pointer"},
	{"math.Pi", "const Pi untyped float"},
	{"io.Reader", "type Reader interface{Read(p []byte) (n int, err error)}"},
	{"io.ReadWriter", "type ReadWriter interface{Read(p []byte) (n int, err error); Write(p []byte) (n int, err error)}"},
	{"math.Sin", "func Sin(x·2 float64) (_ float64)"},
	// TODO(gri) add more tests
}

func TestImportedTypes(t *testing.T) {
	// This package does not handle gccgo export data.
	if runtime.Compiler == "gccgo" {
		return
	}
	for _, test := range importedObjectTests {
		s := strings.Split(test.name, ".")
		if len(s) != 2 {
			t.Fatal("inconsistent test data")
		}
		importPath := s[0]
		objName := s[1]

		pkg, err := Import(imports, importPath)
		if err != nil {
			t.Error(err)
			continue
		}

		obj := pkg.Scope().Lookup(objName)
		if obj == nil {
			t.Errorf("%s: object not found", test.name)
			continue
		}

		got := types.ObjectString(pkg, obj)
		if got != test.want {
			t.Errorf("%s: got %q; want %q", test.name, got, test.want)
		}
	}
}

func TestIssue5815(t *testing.T) {
	// This package does not handle gccgo export data.
	if runtime.Compiler == "gccgo" {
		return
	}

	pkg, err := Import(make(map[string]*types.Package), "strings")
	if err != nil {
		t.Fatal(err)
	}

	scope := pkg.Scope()
	for _, name := range scope.Names() {
		obj := scope.Lookup(name)
		if obj.Pkg() == nil {
			t.Errorf("no pkg for %s", obj)
		}
		if tname, _ := obj.(*types.TypeName); tname != nil {
			named := tname.Type().(*types.Named)
			for i := 0; i < named.NumMethods(); i++ {
				m := named.Method(i)
				if m.Pkg() == nil {
					t.Errorf("no pkg for %s", m)
				}
			}
		}
	}
}
<|endoftext|>


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

package modelzooserver

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strings"

	"sqlflow.org/sqlflow/go/randstring"
)

func imageExistsOnRegistry(imageName, tag string) bool {
	var imageNamePart string
	var registryPart string
	registryPart = os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY")
	imageNamePart = imageName
	// if the imageName contains domain name like "hub.docker.com/group/my_model",
	// get the registry and image name from input imageName.
	if strings.Contains(imageName, ".") {
		parts := strings.Split(imageName, "/")
		registryPart = parts[0]
		imageNamePart = strings.Join(parts[1:len(parts)-1], "/")
	}
	checkURL := fmt.Sprintf("https://%s/v1/repositories/%s/tags/%s", registryPart, imageNamePart, tag)
	resp, err := http.Get(checkURL)
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return false
	}
	_, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return false
	}
	return true
}

func buildAndPushImage(dir, name, tag string, dryrun bool) error {
	// for public model zoo server, the Docker image name should not contain url prefix like
	// hub.docker.com/group/my_model_image
	if strings.Contains(name, ".") {
		return fmt.Errorf("release model definition to public model zoo server should not contain url prefix like hub.docker.com/group/my_model_image, the registry is configured at model zoo server")
	}

	cwd, err := os.Getwd()
	if err != nil {
		return err
	}
	defer os.Chdir(cwd)

	if err := os.Chdir(dir); err != nil {
		return err
	}

	cmd := exec.Command("docker", "build", ".", "-t", fmt.Sprintf("%s:%s", name, tag))
	var cmdStderr bytes.Buffer
	cmd.Stderr = &cmdStderr
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("run docker build err: %v, stderr: %s", err, cmdStderr.String())
	}

	if dryrun {
		// skip push to registry when dryrun enabled.
		return nil
	}

	dockerUsername := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY_USER")
	dockerPasswd := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY_PASS")
	if dockerUsername == "" || dockerPasswd == "" {
		return fmt.Errorf("need to set SQLFLOW_MODEL_ZOO_REGISTRY_USER and SQLFLOW_MODEL_ZOO_REGISTRY_PASS for model zoo server")
	}

	cmd = exec.Command("docker", "login", "-u", dockerUsername, "--password-stdin")
	cmd.Stderr = &cmdStderr
	cmd.Stdin = bytes.NewBufferString(dockerPasswd)
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("run docker login err: %v, stderr: %s", err, cmdStderr.String())
	}

	cmd = exec.Command("docker", "push", fmt.Sprintf("%s:%s", name, tag))
	cmd.Stderr = &cmdStderr
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("run docker push err: %v, stderr: %s", err, cmdStderr.String())
	}
	return nil
}

func buildAndPushImageKaniko(dir, name, tag string, dryrun bool) error {
	reg := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY")
	// default push to dockerhub
	if reg == "" {
		reg = "https://index.docker.io/v1/"
	}
	regUsername := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY_USER")
	regPasswd := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY_PASS")
	regEmail := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY_EMAIL")
	if regUsername == "" || regPasswd == "" || regEmail == "" {
		return fmt.Errorf("need to set SQLFLOW_MODEL_ZOO_REGISTRY_USER and SQLFLOW_MODEL_ZOO_REGISTRY_PASS for model zoo server")
	}
	regSecret := fmt.Sprintf("kaniko-regcred-%s", strings.ToLower(randstring.Generate(8)))

	// TODO(typhoonzero): support any registry servers
	// TODO(typhoonzero): use a random
	// create docker registry credentials
	cmd := fmt.Sprintf(`kubectl create secret docker-registry %s --docker-server=%s --docker-username=%s --docker-password=%s --docker-email=%s`, regSecret, reg, regUsername, regPasswd, regEmail)
	cmdList := strings.Split(cmd, " ")
	c := exec.Command(cmdList[0], cmdList[1:]...)
	if err := c.Run(); err != nil {
		return err
	}
	defer func() error {
		cmdDeleteSecret := exec.Command("kubectl", "delete", "secret", regSecret)
		if err := cmdDeleteSecret.Run(); err != nil {
			return err
		}
		return nil
	}()

	// cd into dir to tar the context
	cwd, err := os.Getwd()
	if err != nil {
		return err
	}
	defer os.Chdir(cwd)
	if err := os.Chdir(dir); err != nil {
		return err
	}

	kanikoImage := os.Getenv("SQLFLOW_MODEL_ZOO_KANIKO_IMAGE")
	if kanikoImage == "" {
		kanikoImage = "registry.cn-hangzhou.aliyuncs.com/sql-machine-learning/kaniko-executor"
	}
	kanikoPodName := fmt.Sprintf("kaniko-%s", strings.ToLower(randstring.Generate(8)))
	destination := fmt.Sprintf("%s:%s", name, tag)
	// pod tolerations/nodeSelector JSON like: '"tolerations": [...],', if empty string, tolerations will not be set.
	tolerations := os.Getenv("SQLFLOW_MODEL_ZOO_TOLERATIONS")
	nodeSelector := os.Getenv("SQLFLOW_MODEL_ZOO_NODE_SELECTOR")
	podTemplate := fmt.Sprintf(`'{
  "apiVersion": "v1",
  "spec": {
    %s
    %s
    "containers": [
    {
      "name": "%s",
      "image": "%s",
      "stdin": true,
      "stdinOnce": true,
      "args": [
        "--dockerfile=Dockerfile",
        "--context=tar://stdin",
        "--destination=%s"],
      "volumeMounts": [
        {
          "name": "kaniko-secret",
          "mountPath": "/kaniko/.docker/"
      }]
    }],
    "volumes": [
    {
      "name": "kaniko-secret",
      "secret": {
        "secretName": "%s",
        "items": [{"key": ".dockerconfigjson", "path": "config.json"}]
      }}
    ]
  }
}'`, tolerations, nodeSelector, kanikoPodName, kanikoImage, destination, regSecret)

	tarContextCmd := exec.Command("tar", "czf", "-", ".")
	// exec.Command can not handle quotes correctly, use bach -c here.
	kanikoBuildCmdStr := fmt.Sprintf(`kubectl run %s --rm --stdin=true --restart=Never --image=%s --overrides=%s`, kanikoPodName, kanikoImage, podTemplate)
	kanikoBuildCmdStdin := exec.Command("bash", "-c", kanikoBuildCmdStr)

	r, w := io.Pipe()
	tarContextCmd.Dir = dir
	tarContextCmd.Stdout = w
	kanikoBuildCmdStdin.Stdin = r
	var outputErr bytes.Buffer
	kanikoBuildCmdStdin.Stderr = &outputErr
	if err := tarContextCmd.Start(); err != nil {
		return err
	}
	if err := kanikoBuildCmdStdin.Start(); err != nil {
		return err
	}
	if err := tarContextCmd.Wait(); err != nil {
		return err
	}
	if err := w.Close(); err != nil {
		return err
	}

	if err := kanikoBuildCmdStdin.Wait(); err != nil {
		// TODO(typhoonzero): use log to output error messages.
		io.Copy(os.Stdout, &outputErr)
		return err
	}

	return nil
}
<commit_msg>kaniko force build to run in internal (#2713)<commit_after>// Copyright 2020 The SQLFlow Authors. All rights reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package modelzooserver

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strings"

	"sqlflow.org/sqlflow/go/randstring"
)

func imageExistsOnRegistry(imageName, tag string) bool {
	var imageNamePart string
	var registryPart string
	registryPart = os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY")
	imageNamePart = imageName
	// if the imageName contains domain name like "hub.docker.com/group/my_model",
	// get the registry and image name from input imageName.
	if strings.Contains(imageName, ".") {
		parts := strings.Split(imageName, "/")
		registryPart = parts[0]
		imageNamePart = strings.Join(parts[1:len(parts)-1], "/")
	}
	checkURL := fmt.Sprintf("https://%s/v1/repositories/%s/tags/%s", registryPart, imageNamePart, tag)
	resp, err := http.Get(checkURL)
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return false
	}
	_, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return false
	}
	return true
}

func buildAndPushImage(dir, name, tag string, dryrun bool) error {
	// for public model zoo server, the Docker image name should not contain url prefix like
	// hub.docker.com/group/my_model_image
	if strings.Contains(name, ".") {
		return fmt.Errorf("release model definition to public model zoo server should not contain url prefix like hub.docker.com/group/my_model_image, the registry is configured at model zoo server")
	}

	cwd, err := os.Getwd()
	if err != nil {
		return err
	}
	defer os.Chdir(cwd)

	if err := os.Chdir(dir); err != nil {
		return err
	}

	cmd := exec.Command("docker", "build", ".", "-t", fmt.Sprintf("%s:%s", name, tag))
	var cmdStderr bytes.Buffer
	cmd.Stderr = &cmdStderr
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("run docker build err: %v, stderr: %s", err, cmdStderr.String())
	}

	if dryrun {
		// skip push to registry when dryrun enabled.
		return nil
	}

	dockerUsername := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY_USER")
	dockerPasswd := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY_PASS")
	if dockerUsername == "" || dockerPasswd == "" {
		return fmt.Errorf("need to set SQLFLOW_MODEL_ZOO_REGISTRY_USER and SQLFLOW_MODEL_ZOO_REGISTRY_PASS for model zoo server")
	}

	cmd = exec.Command("docker", "login", "-u", dockerUsername, "--password-stdin")
	cmd.Stderr = &cmdStderr
	cmd.Stdin = bytes.NewBufferString(dockerPasswd)
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("run docker login err: %v, stderr: %s", err, cmdStderr.String())
	}

	cmd = exec.Command("docker", "push", fmt.Sprintf("%s:%s", name, tag))
	cmd.Stderr = &cmdStderr
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("run docker push err: %v, stderr: %s", err, cmdStderr.String())
	}
	return nil
}

func buildAndPushImageKaniko(dir, name, tag string, dryrun bool) error {
	reg := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY")
	// default push to dockerhub
	if reg == "" {
		reg = "https://index.docker.io/v1/"
	}
	regUsername := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY_USER")
	regPasswd := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY_PASS")
	regEmail := os.Getenv("SQLFLOW_MODEL_ZOO_REGISTRY_EMAIL")
	if regUsername == "" || regPasswd == "" || regEmail == "" {
		return fmt.Errorf("need to set SQLFLOW_MODEL_ZOO_REGISTRY_USER and SQLFLOW_MODEL_ZOO_REGISTRY_PASS for model zoo server")
	}
	regSecret := fmt.Sprintf("kaniko-regcred-%s", strings.ToLower(randstring.Generate(8)))

	// TODO(typhoonzero): support any registry servers
	// TODO(typhoonzero): use a random
	// create docker registry credentials
	cmd := fmt.Sprintf(`kubectl create secret docker-registry %s --docker-server=%s --docker-username=%s --docker-password=%s --docker-email=%s`, regSecret, reg, regUsername, regPasswd, regEmail)
	cmdList := strings.Split(cmd, " ")
	c := exec.Command(cmdList[0], cmdList[1:]...)
	if err := c.Run(); err != nil {
		return err
	}
	defer func() error {
		cmdDeleteSecret := exec.Command("kubectl", "delete", "secret", regSecret)
		if err := cmdDeleteSecret.Run(); err != nil {
			return err
		}
		return nil
	}()

	// cd into dir to tar the context
	cwd, err := os.Getwd()
	if err != nil {
		return err
	}
	defer os.Chdir(cwd)
	if err := os.Chdir(dir); err != nil {
		return err
	}

	kanikoImage := os.Getenv("SQLFLOW_MODEL_ZOO_KANIKO_IMAGE")
	if kanikoImage == "" {
		kanikoImage = "registry.cn-hangzhou.aliyuncs.com/sql-machine-learning/kaniko-executor"
	}
	kanikoPodName := fmt.Sprintf("kaniko-%s", strings.ToLower(randstring.Generate(8)))
	destination := fmt.Sprintf("%s:%s", name, tag)
	// pod tolerations/nodeSelector JSON like: '"tolerations": [...],', if empty string, tolerations will not be set.
	tolerations := os.Getenv("SQLFLOW_MODEL_ZOO_TOLERATIONS")
	nodeSelector := os.Getenv("SQLFLOW_MODEL_ZOO_NODE_SELECTOR")
	podTemplate := fmt.Sprintf(`'{
  "apiVersion": "v1",
  "spec": {
    %s
    %s
    "containers": [
    {
      "name": "%s",
      "image": "%s",
      "stdin": true,
      "stdinOnce": true,
      "args": [
        "--dockerfile=Dockerfile",
        "--context=tar://stdin",
        "--force",
        "--destination=%s"],
      "volumeMounts": [
        {
          "name": "kaniko-secret",
          "mountPath": "/kaniko/.docker/"
      }]
    }],
    "volumes": [
    {
      "name": "kaniko-secret",
      "secret": {
        "secretName": "%s",
        "items": [{"key": ".dockerconfigjson", "path": "config.json"}]
      }}
    ]
  }
}'`, tolerations, nodeSelector, kanikoPodName, kanikoImage, destination, regSecret)

	tarContextCmd := exec.Command("tar", "czf", "-", ".")
	// exec.Command can not handle quotes correctly, use bach -c here.
	kanikoBuildCmdStr := fmt.Sprintf(`kubectl run %s --rm --stdin=true --restart=Never --image=%s --overrides=%s`, kanikoPodName, kanikoImage, podTemplate)
	kanikoBuildCmdStdin := exec.Command("bash", "-c", kanikoBuildCmdStr)

	r, w := io.Pipe()
	tarContextCmd.Dir = dir
	tarContextCmd.Stdout = w
	kanikoBuildCmdStdin.Stdin = r
	var outputErr bytes.Buffer
	kanikoBuildCmdStdin.Stderr = &outputErr
	if err := tarContextCmd.Start(); err != nil {
		return err
	}
	if err := kanikoBuildCmdStdin.Start(); err != nil {
		return err
	}
	if err := tarContextCmd.Wait(); err != nil {
		return err
	}
	if err := w.Close(); err != nil {
		return err
	}

	if err := kanikoBuildCmdStdin.Wait(); err != nil {
		// TODO(typhoonzero): use log to output error messages.
		io.Copy(os.Stdout, &outputErr)
		return err
	}

	return nil
}
<|endoftext|>