<commit_before>package lib

import (
	"bytes"
	"compress/gzip"
	"crypto/tls"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/davecgh/go-spew/spew"
)

type redirectCaughtError struct{}

func (c redirectCaughtError) Error() string {
	return "redirect caught"
}

type proxyServer struct {
	local     string
	remote    string
	log       io.Writer
	rewriters []Rewriter
	mappings  []Mapping

	server http.Server
	client http.Client
}

func init() {
	_ = spew.UnsafeDisabled
}

func (p *proxyServer) ServeHTTP(outgoing http.ResponseWriter, incoming *http.Request) {
	var err error

	hostMappings := buildHostMappings(p.mappings, incoming.Host)
	for _, rewriter := range p.rewriters {
		rewriter.SetMappings(hostMappings)
	}

	proxyRequest, err := p.buildProxyRequest(incoming)
	var response *http.Response

	if err == nil {
		response, err = p.client.Do(proxyRequest)

		if isRedirectError(err) {
			err = nil
		}
	}

	if err != nil {
		fmt.Fprintf(p.log, "error during proxy request: %v\n", err)
		http.Error(outgoing, err.Error(), http.StatusBadGateway)
	} else {
		p.sendProxyResponse(incoming, response, outgoing)
	}
}

func isRedirectError(err error) (q bool) {
	urlError, q := err.(*url.Error)
	if !q {
		return
	}

	_, q = urlError.Err.(redirectCaughtError)

	return
}

func (p *proxyServer) buildProxyRequest(incoming *http.Request) (outgoing *http.Request, err error) {
	outgoing, err = http.NewRequest(incoming.Method, p.remote+incoming.RequestURI, incoming.Body)

	for key, values := range incoming.Header {
		for _, value := range values {
			outgoing.Header.Add(key, value)
		}
	}

	for _, rewriter := range p.rewriters {
		if rewriter, ok := rewriter.(IncomingHeaderRewriter); ok {
			rewriter.RewriteIncomingHeaders(outgoing.Header)
		}
	}

	// We also try to compress upstream communication
	outgoing.Header.Set("accept-encoding", "gzip")

	return
}

func (p *proxyServer) sendProxyResponse(request *http.Request, response *http.Response, outgoing http.ResponseWriter) {
	defer response.Body.Close()

	// Transfer headers
	outgoingHeaders := outgoing.Header()
	for key, values := range response.Header {
		for _, value := range values {
			outgoingHeaders.Add(key, value)
		}
	}

	// Content length will be recalculated as content may be rewritten
	outgoingHeaders.Del("content-length")

	// We need to remove any domain information set on the cookies
	outgoingHeaders.Del("set-cookie")
	for _, cookie := range response.Cookies() {
		cookie.Domain = ""

		http.SetCookie(outgoing, cookie)
	}

	// First rewrite pass: identify matching content rewriters and process headers
	responseRewriters := make([]ResponseRewriter, 0, len(p.rewriters))

	for _, rewriter := range p.rewriters {
		if r, ok := rewriter.(ResponseRewriter); ok {
			if r.Matches(request, response) {
				responseRewriters = append(responseRewriters, r)
			}
		}

		if r, ok := rewriter.(HeaderRewriter); ok {
			r.RewriteHeaders(outgoingHeaders)
		}
	}

	// Handle compression
	bodyWriter := outgoing.(io.Writer)
	bodyReader := response.Body.(io.Reader)

	if response.Header.Get("content-encoding") == "gzip" {
		var err error
		bodyReader, err = gzip.NewReader(bodyReader)
		if err != nil {
			http.Error(outgoing, err.Error(), http.StatusBadGateway)
			return
		}

		// We are ignoring any q-value here, so this is wrong for the case q=0
		if strings.Contains(request.Header.Get("accept-encoding"), "gzip") {
			bodyWriter = gzip.NewWriter(bodyWriter)
			outgoingHeaders.Set("content-encoding", "gzip")
		} else {
			outgoingHeaders.Del("content-encoding")
		}
	}

	// Send headers
	outgoing.WriteHeader(response.StatusCode)

	// If there are any body rewriters, we read the response into memory and process it
	if len(responseRewriters) > 0 {
		bodyData, err := ioutil.ReadAll(bodyReader)

		if err == nil {
			for _, rewriter := range responseRewriters {
				bodyData = rewriter.RewriteResponse(bodyData)
			}
		} else {
			bodyData = make([]byte, 0)
		}

		bodyReader = bytes.NewBuffer(bodyData)
	}

	// Send body
	io.Copy(bodyWriter, bodyReader)

	// Close writer if applicable (looking at you, gzip)
	if bodyWriter, ok := bodyWriter.(io.Closer); ok {
		bodyWriter.Close()
	}
}

func (p *proxyServer) Start() <-chan error {
	c := make(chan error, 1)

	go func() {
		c <- p.server.ListenAndServe()
	}()

	fmt.Fprintf(p.log, "proxying requests for %s to %s\n", p.local, p.remote)

	return c
}

func (p *proxyServer) AddRewriter(r Rewriter) {
	p.rewriters = append(p.rewriters, r)
}

func newProxyServer(m Mapping, mappings []Mapping, log io.Writer, sslAllowInsecure bool) (p *proxyServer, err error) {
	p = &proxyServer{
		local:     m.local,
		remote:    m.remote,
		log:       log,
		rewriters: make([]Rewriter, 0),
		mappings:  mappings,
	}

	p.server = http.Server{
		Addr:    p.local,
		Handler: p,
	}

	var tlsConfig *tls.Config
	if sslAllowInsecure {
		tlsConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}

	p.client = http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return redirectCaughtError{}
		},
		Transport: &http.Transport{
			// We rather handle compression ourselves
			DisableCompression: true,
			TLSClientConfig:    tlsConfig,
		},
	}

	return
}
<commit_msg>Don't pass on transfer encoding, don't decode stream unless we really need to.<commit_after>package lib

import (
	"bytes"
	"compress/gzip"
	"crypto/tls"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/davecgh/go-spew/spew"
)

type redirectCaughtError struct{}

func (c redirectCaughtError) Error() string {
	return "redirect caught"
}

type proxyServer struct {
	local     string
	remote    string
	log       io.Writer
	rewriters []Rewriter
	mappings  []Mapping

	server http.Server
	client http.Client
}

func init() {
	_ = spew.UnsafeDisabled
}

func (p *proxyServer) ServeHTTP(outgoing http.ResponseWriter, incoming *http.Request) {
	var err error

	hostMappings := buildHostMappings(p.mappings, incoming.Host)
	for _, rewriter := range p.rewriters {
		rewriter.SetMappings(hostMappings)
	}

	proxyRequest, err := p.buildProxyRequest(incoming)
	var response *http.Response

	if err == nil {
		response, err = p.client.Do(proxyRequest)

		if isRedirectError(err) {
			err = nil
		}
	}

	if err != nil {
		fmt.Fprintf(p.log, "error during proxy request: %v\n", err)
		http.Error(outgoing, err.Error(), http.StatusBadGateway)
	} else {
		p.sendProxyResponse(incoming, response, outgoing)
	}
}

func isRedirectError(err error) (q bool) {
	urlError, q := err.(*url.Error)
	if !q {
		return
	}

	_, q = urlError.Err.(redirectCaughtError)

	return
}

func (p *proxyServer) buildProxyRequest(incoming *http.Request) (outgoing *http.Request, err error) {
	outgoing, err = http.NewRequest(incoming.Method, p.remote+incoming.RequestURI, incoming.Body)

	for key, values := range incoming.Header {
		for _, value := range values {
			outgoing.Header.Add(key, value)
		}
	}

	for _, rewriter := range p.rewriters {
		if rewriter, ok := rewriter.(IncomingHeaderRewriter); ok {
			rewriter.RewriteIncomingHeaders(outgoing.Header)
		}
	}

	// We also try to compress upstream communication
	outgoing.Header.Set("accept-encoding", "gzip")

	return
}

func (p *proxyServer) sendProxyResponse(request *http.Request, response *http.Response, outgoing http.ResponseWriter) {
	defer response.Body.Close()

	// Transfer headers
	outgoingHeaders := outgoing.Header()
	for key, values := range response.Header {
		for _, value := range values {
			outgoingHeaders.Add(key, value)
		}
	}

	// Content length will be recalculated as content may be rewritten
	outgoingHeaders.Del("content-length")

	// Reset transfer-encoding
	outgoingHeaders.Del("transfer-encoding")

	// We need to remove any domain information set on the cookies
	outgoingHeaders.Del("set-cookie")
	for _, cookie := range response.Cookies() {
		cookie.Domain = ""

		http.SetCookie(outgoing, cookie)
	}

	// First rewrite pass: identify matching content rewriters and process headers
	responseRewriters := make([]ResponseRewriter, 0, len(p.rewriters))

	for _, rewriter := range p.rewriters {
		if r, ok := rewriter.(ResponseRewriter); ok {
			if r.Matches(request, response) {
				responseRewriters = append(responseRewriters, r)
			}
		}

		if r, ok := rewriter.(HeaderRewriter); ok {
			r.RewriteHeaders(outgoingHeaders)
		}
	}

	// Handle compression
	bodyWriter := outgoing.(io.Writer)
	bodyReader := response.Body.(io.Reader)

	// We are ignoring any q-value here, so this is wrong for the case q=0
	clientAcceptsGzip := strings.Contains(request.Header.Get("accept-encoding"), "gzip")

	if response.Header.Get("content-encoding") == "gzip" &&
		(len(responseRewriters) > 0 || !clientAcceptsGzip) {

		var err error
		bodyReader, err = gzip.NewReader(bodyReader)
		if err != nil {
			http.Error(outgoing, err.Error(), http.StatusBadGateway)
			return
		}

		if clientAcceptsGzip {
			bodyWriter = gzip.NewWriter(bodyWriter)
			outgoingHeaders.Set("content-encoding", "gzip")
		} else {
			outgoingHeaders.Del("content-encoding")
		}
	}

	// Send headers
	outgoing.WriteHeader(response.StatusCode)

	// If there are any body rewriters, we read the response into memory and process it
	if len(responseRewriters) > 0 {
		bodyData, err := ioutil.ReadAll(bodyReader)

		if err == nil {
			for _, rewriter := range responseRewriters {
				bodyData = rewriter.RewriteResponse(bodyData)
			}
		} else {
			bodyData = make([]byte, 0)
		}

		bodyReader = bytes.NewBuffer(bodyData)
	}

	// Send body
	io.Copy(bodyWriter, bodyReader)

	// Close writer if applicable (looking at you, gzip)
	if bodyWriter, ok := bodyWriter.(io.Closer); ok {
		bodyWriter.Close()
	}
}

func (p *proxyServer) Start() <-chan error {
	c := make(chan error, 1)

	go func() {
		c <- p.server.ListenAndServe()
	}()

	fmt.Fprintf(p.log, "proxying requests for %s to %s\n", p.local, p.remote)

	return c
}

func (p *proxyServer) AddRewriter(r Rewriter) {
	p.rewriters = append(p.rewriters, r)
}

func newProxyServer(m Mapping, mappings []Mapping, log io.Writer, sslAllowInsecure bool) (p *proxyServer, err error) {
	p = &proxyServer{
		local:     m.local,
		remote:    m.remote,
		log:       log,
		rewriters: make([]Rewriter, 0),
		mappings:  mappings,
	}

	p.server = http.Server{
		Addr:    p.local,
		Handler: p,
	}

	var tlsConfig *tls.Config
	if sslAllowInsecure {
		tlsConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}

	p.client = http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return redirectCaughtError{}
		},
		Transport: &http.Transport{
			// We rather handle compression ourselves
			DisableCompression: true,
			TLSClientConfig:    tlsConfig,
		},
	}

	return
}
<|endoftext|>


<commit_before>// Copyright 2019 The LUCI Authors.
//
// 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.

// TODO(hinoka): Remove this file after migrating to the new archival tasking pipeline.
// These settings are used for adjusting how much traffic goes into the new archival tasking pipeline.
// 3 Settings are available for the new pipeline:
// * Optimistic archival delay - Controls the delay to set for optimistic archival.
// * Optimistic archival tasking percentage - Controls what % of terminated streams are tasked to the new pipeline
// * Pessimistic archival tasking percentage - Controls what % of registered streams are tasked to the new pipeline.

package coordinator

import (
	"context"
	"fmt"
	"html/template"
	"strconv"
	"time"

	"go.chromium.org/luci/common/clock/clockflag"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/server/portal"
	"go.chromium.org/luci/server/settings"
)

const (
	settingDisabled = "disabled"
	settingEnabled  = "enabled"
	baseName        = "archivist"
)

type Settings struct {
	// OptimisticArchivalDelay controls the delay to set for optimistic archival.
	OptimisticArchivalDelay time.Duration

	// OptimisticArchivalPercent (0-100) Controls what percent of terminated streams
	// are tasked to the new pipeline
	OptimisticalArchivalPercent uint32

	// PessimisticArchivalPercent (0-100) Controls what percent of registered
	// streams are tasked to the new pipeline after 47 hours.
	PessimisticArchivalPercent uint32

	// ArchivistBatchSize controls the batchsize per lease cycle.
	ArchivistBatchSize int64

	// ArchivistLeaseTime controls the lease time per cycle.
	ArchivistLeaseTime time.Duration
}

var defaultSettings = Settings{
	OptimisticArchivalDelay: 5 * time.Minute,
	ArchivistBatchSize:      50,
	ArchivistLeaseTime:      15 * time.Minute,
}

// settingsPage is a UI page to configure a static Tumble configuration.
type settingsPage struct {
	portal.BasePage
}

func (settingsPage) Title(c context.Context) (string, error) {
	return "LogDog Archivist Settings", nil
}

func (settingsPage) Overview(c context.Context) (template.HTML, error) {
	return template.HTML(`<p>Configuration parameters for the Archivist tasking pipeline.</p>.`), nil
}

func (settingsPage) Fields(c context.Context) ([]portal.Field, error) {
	return []portal.Field{
		{
			ID: "OptimisticArchivalDelay",
			Title: "The delay, of how the optimistic archival delay. " +
				"This should be set to compensate for potential collector pipeline delays.",
			Type:        portal.FieldText,
			Placeholder: defaultSettings.OptimisticArchivalDelay.String(),
			Validator:   validateDuration,
		},
		{
			ID:          "OptimisticArchivalPercent",
			Title:       "Percentage (0-100) of tasks to go to the new optimistic pipeline.",
			Type:        portal.FieldText,
			Placeholder: "0",
			Validator:   validatePercent,
		},
		{
			ID:          "PessimisticArchivalPercent",
			Title:       "Percentage (0-100) of tasks to go to the new pessimistic pipeline (47hr delay from registration).",
			Type:        portal.FieldText,
			Placeholder: "0",
			Validator:   validatePercent,
		},
		{
			ID: "ArchivistBatchSize",
			Title: "Number of Archive Tasks to lease in each LeaseTasks request. " +
				"The theoretical max Task processed per second is this number * 10, " +
				"due to the limitations of GAE TaskQueues.",
			Type:        portal.FieldText,
			Placeholder: fmt.Sprintf("%d", defaultSettings.ArchivistBatchSize),
			Validator:   validateInt,
		},
		{
			ID:          "ArchivistLeaseTime",
			Title:       "Lease time per cycle for the Archivist.",
			Type:        portal.FieldText,
			Placeholder: defaultSettings.ArchivistLeaseTime.String(),
			Validator:   validateDuration,
		},
	}, nil
}

func (settingsPage) ReadSettings(c context.Context) (map[string]string, error) {
	var set Settings
	switch err := settings.GetUncached(c, baseName, &set); err {
	case nil:
		break
	case settings.ErrNoSettings:
		logging.WithError(err).Infof(c, "No settings available, using defaults.")
		set = defaultSettings
	default:
		return nil, err
	}

	values := map[string]string{}

	// Only render values if they differ from our default config.
	if set.OptimisticArchivalDelay != defaultSettings.OptimisticArchivalDelay {
		values["OptimisticArchivalDelay"] = set.OptimisticArchivalDelay.String()
	}
	if set.OptimisticalArchivalPercent != defaultSettings.OptimisticalArchivalPercent {
		values["OptimisticalArchivalPercent"] = fmt.Sprintf("%d", set.OptimisticalArchivalPercent)
	}
	if set.PessimisticArchivalPercent != defaultSettings.PessimisticArchivalPercent {
		values["PessimisticArchivalPercent"] = fmt.Sprintf("%d", set.PessimisticArchivalPercent)
	}
	if set.ArchivistBatchSize != defaultSettings.ArchivistBatchSize {
		values["ArchivistBatchSize"] = fmt.Sprintf("%d", set.ArchivistBatchSize)
	}
	if set.ArchivistLeaseTime != defaultSettings.ArchivistLeaseTime {
		values["ArchivistLeaseTime"] = set.ArchivistLeaseTime.String()
	}

	return values, nil
}

func (settingsPage) WriteSettings(c context.Context, values map[string]string, who, why string) error {
	// Start with our default config and shape it with populated values.
	set := defaultSettings

	if v := values["OptimisticArchivalDelay"]; v != "" {
		t, err := clockflag.ParseDuration(v)
		if err != nil {
			return fmt.Errorf("could not parse OptimisticArchivalDelay: %v", err)
		}
		set.OptimisticArchivalDelay = time.Duration(t)
	}
	if v := values["OptimisticArchivalPercent"]; v != "" {
		i, err := strconv.Atoi(v)
		if err != nil {
			return fmt.Errorf("could not parse OptimisticArchivalPercent: %v", err)
		}
		set.OptimisticalArchivalPercent = uint32(i)
	}
	if v := values["PessimisticArchivalPercent"]; v != "" {
		i, err := strconv.Atoi(v)
		if err != nil {
			return fmt.Errorf("could not parse PessimisticArchivalPercent: %v", err)
		}
		set.PessimisticArchivalPercent = uint32(i)
	}
	if v := values["ArchivistBatchSize"]; v != "" {
		i, err := strconv.Atoi(v)
		if err != nil {
			return fmt.Errorf("could not parse ArchivistBatchSize : %v", err)
		}
		set.ArchivistBatchSize = int64(i)
	}
	if v := values["ArchivistLeaseTime"]; v != "" {
		t, err := clockflag.ParseDuration(v)
		if err != nil {
			return fmt.Errorf("could not parse ArchivistLeaseTime : %v", err)
		}
		set.ArchivistLeaseTime = time.Duration(t)
	}

	return settings.SetIfChanged(c, baseName, &set, who, why)
}

// validatePercent validates a string is an integer between 0-100.
func validatePercent(v string) error {
	if v == "" {
		return nil
	}
	i, err := strconv.Atoi(v)
	if err != nil {
		return fmt.Errorf("invalid integer %q - %s", v, err)
	}
	if i < 0 || i > 100 {
		return fmt.Errorf("%d is out of range (0-100)", i)
	}
	return nil
}

// validateInt validates a string is an integer between 0-100.
func validateInt(v string) error {
	if v == "" {
		return nil
	}
	_, err := strconv.Atoi(v)
	return err
}

func validateDuration(v string) error {
	if v == "" {
		return nil
	}

	var cf clockflag.Duration
	if err := cf.Set(v); err != nil {
		return fmt.Errorf("bad duration %q - %s", v, err)
	}
	if cf <= 0 {
		return fmt.Errorf("duration %q must be positive", v)
	}
	return nil
}

// GetSettings returns the current settings.
//
// It first tries to load it from settings. If no settings is installed, or if
// there is no configuration in settings, defaultSettings is returned.
func GetSettings(c context.Context) *Settings {
	set := Settings{}
	switch err := settings.Get(c, baseName, &set); err {
	case nil:
		break
	case settings.ErrNoSettings:
		// Defaults.
		set = defaultSettings
	default:
		panic(fmt.Errorf("could not fetch Archivist settings - %s", err))
	}
	return &set
}

func init() {
	portal.RegisterPage(baseName, settingsPage{})
}
<commit_msg>[logdog] Fix settings func doc typos<commit_after>// Copyright 2019 The LUCI Authors.
//
// 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.

// TODO(hinoka): Remove this file after migrating to the new archival tasking pipeline.
// These settings are used for adjusting how much traffic goes into the new archival tasking pipeline.
// 3 Settings are available for the new pipeline:
// * Optimistic archival delay - Controls the delay to set for optimistic archival.
// * Optimistic archival tasking percentage - Controls what % of terminated streams are tasked to the new pipeline
// * Pessimistic archival tasking percentage - Controls what % of registered streams are tasked to the new pipeline.

package coordinator

import (
	"context"
	"fmt"
	"html/template"
	"strconv"
	"time"

	"go.chromium.org/luci/common/clock/clockflag"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/server/portal"
	"go.chromium.org/luci/server/settings"
)

const (
	settingDisabled = "disabled"
	settingEnabled  = "enabled"
	baseName        = "archivist"
)

type Settings struct {
	// OptimisticArchivalDelay controls the delay to set for optimistic archival.
	OptimisticArchivalDelay time.Duration

	// OptimisticArchivalPercent (0-100) Controls what percent of terminated streams
	// are tasked to the new pipeline
	OptimisticalArchivalPercent uint32

	// PessimisticArchivalPercent (0-100) Controls what percent of registered
	// streams are tasked to the new pipeline after 47 hours.
	PessimisticArchivalPercent uint32

	// ArchivistBatchSize controls the batchsize per lease cycle.
	ArchivistBatchSize int64

	// ArchivistLeaseTime controls the lease time per cycle.
	ArchivistLeaseTime time.Duration
}

var defaultSettings = Settings{
	OptimisticArchivalDelay: 5 * time.Minute,
	ArchivistBatchSize:      50,
	ArchivistLeaseTime:      15 * time.Minute,
}

// settingsPage is a UI page to configure a static Tumble configuration.
type settingsPage struct {
	portal.BasePage
}

func (settingsPage) Title(c context.Context) (string, error) {
	return "LogDog Archivist Settings", nil
}

func (settingsPage) Overview(c context.Context) (template.HTML, error) {
	return template.HTML(`<p>Configuration parameters for the Archivist tasking pipeline.</p>.`), nil
}

func (settingsPage) Fields(c context.Context) ([]portal.Field, error) {
	return []portal.Field{
		{
			ID: "OptimisticArchivalDelay",
			Title: "The delay, of how the optimistic archival delay. " +
				"This should be set to compensate for potential collector pipeline delays.",
			Type:        portal.FieldText,
			Placeholder: defaultSettings.OptimisticArchivalDelay.String(),
			Validator:   validateDuration,
		},
		{
			ID:          "OptimisticArchivalPercent",
			Title:       "Percentage (0-100) of tasks to go to the new optimistic pipeline.",
			Type:        portal.FieldText,
			Placeholder: "0",
			Validator:   validatePercent,
		},
		{
			ID:          "PessimisticArchivalPercent",
			Title:       "Percentage (0-100) of tasks to go to the new pessimistic pipeline (47hr delay from registration).",
			Type:        portal.FieldText,
			Placeholder: "0",
			Validator:   validatePercent,
		},
		{
			ID: "ArchivistBatchSize",
			Title: "Number of Archive Tasks to lease in each LeaseTasks request. " +
				"The theoretical max Task processed per second is this number * 10, " +
				"due to the limitations of GAE TaskQueues.",
			Type:        portal.FieldText,
			Placeholder: fmt.Sprintf("%d", defaultSettings.ArchivistBatchSize),
			Validator:   validateInt,
		},
		{
			ID:          "ArchivistLeaseTime",
			Title:       "Lease time per cycle for the Archivist.",
			Type:        portal.FieldText,
			Placeholder: defaultSettings.ArchivistLeaseTime.String(),
			Validator:   validateDuration,
		},
	}, nil
}

func (settingsPage) ReadSettings(c context.Context) (map[string]string, error) {
	var set Settings
	switch err := settings.GetUncached(c, baseName, &set); err {
	case nil:
		break
	case settings.ErrNoSettings:
		logging.WithError(err).Infof(c, "No settings available, using defaults.")
		set = defaultSettings
	default:
		return nil, err
	}

	values := map[string]string{}

	// Only render values if they differ from our default config.
	if set.OptimisticArchivalDelay != defaultSettings.OptimisticArchivalDelay {
		values["OptimisticArchivalDelay"] = set.OptimisticArchivalDelay.String()
	}
	if set.OptimisticalArchivalPercent != defaultSettings.OptimisticalArchivalPercent {
		values["OptimisticalArchivalPercent"] = fmt.Sprintf("%d", set.OptimisticalArchivalPercent)
	}
	if set.PessimisticArchivalPercent != defaultSettings.PessimisticArchivalPercent {
		values["PessimisticArchivalPercent"] = fmt.Sprintf("%d", set.PessimisticArchivalPercent)
	}
	if set.ArchivistBatchSize != defaultSettings.ArchivistBatchSize {
		values["ArchivistBatchSize"] = fmt.Sprintf("%d", set.ArchivistBatchSize)
	}
	if set.ArchivistLeaseTime != defaultSettings.ArchivistLeaseTime {
		values["ArchivistLeaseTime"] = set.ArchivistLeaseTime.String()
	}

	return values, nil
}

func (settingsPage) WriteSettings(c context.Context, values map[string]string, who, why string) error {
	// Start with our default config and shape it with populated values.
	set := defaultSettings

	if v := values["OptimisticArchivalDelay"]; v != "" {
		t, err := clockflag.ParseDuration(v)
		if err != nil {
			return fmt.Errorf("could not parse OptimisticArchivalDelay: %v", err)
		}
		set.OptimisticArchivalDelay = time.Duration(t)
	}
	if v := values["OptimisticArchivalPercent"]; v != "" {
		i, err := strconv.Atoi(v)
		if err != nil {
			return fmt.Errorf("could not parse OptimisticArchivalPercent: %v", err)
		}
		set.OptimisticalArchivalPercent = uint32(i)
	}
	if v := values["PessimisticArchivalPercent"]; v != "" {
		i, err := strconv.Atoi(v)
		if err != nil {
			return fmt.Errorf("could not parse PessimisticArchivalPercent: %v", err)
		}
		set.PessimisticArchivalPercent = uint32(i)
	}
	if v := values["ArchivistBatchSize"]; v != "" {
		i, err := strconv.Atoi(v)
		if err != nil {
			return fmt.Errorf("could not parse ArchivistBatchSize : %v", err)
		}
		set.ArchivistBatchSize = int64(i)
	}
	if v := values["ArchivistLeaseTime"]; v != "" {
		t, err := clockflag.ParseDuration(v)
		if err != nil {
			return fmt.Errorf("could not parse ArchivistLeaseTime : %v", err)
		}
		set.ArchivistLeaseTime = time.Duration(t)
	}

	return settings.SetIfChanged(c, baseName, &set, who, why)
}

// validatePercent validates a string is empty or an integer between 0-100.
func validatePercent(v string) error {
	if v == "" {
		return nil
	}
	i, err := strconv.Atoi(v)
	if err != nil {
		return fmt.Errorf("invalid integer %q - %s", v, err)
	}
	if i < 0 || i > 100 {
		return fmt.Errorf("%d is out of range (0-100)", i)
	}
	return nil
}

// validateInt validates that a string is empty or an integer.
func validateInt(v string) error {
	if v == "" {
		return nil
	}
	_, err := strconv.Atoi(v)
	return err
}

func validateDuration(v string) error {
	if v == "" {
		return nil
	}

	var cf clockflag.Duration
	if err := cf.Set(v); err != nil {
		return fmt.Errorf("bad duration %q - %s", v, err)
	}
	if cf <= 0 {
		return fmt.Errorf("duration %q must be positive", v)
	}
	return nil
}

// GetSettings returns the current settings.
//
// It first tries to load it from settings. If no settings is installed, or if
// there is no configuration in settings, defaultSettings is returned.
func GetSettings(c context.Context) *Settings {
	set := Settings{}
	switch err := settings.Get(c, baseName, &set); err {
	case nil:
		break
	case settings.ErrNoSettings:
		// Defaults.
		set = defaultSettings
	default:
		panic(fmt.Errorf("could not fetch Archivist settings - %s", err))
	}
	return &set
}

func init() {
	portal.RegisterPage(baseName, settingsPage{})
}
<|endoftext|>


<commit_before>// SILVER - Service Wrapper
// Auto Updater
//
// Copyright (c) 2014 PaperCut Software http://www.papercut.com/
// Use of this source code is governed by an MIT or GPL Version 2 license.
// See the project's LICENSE file for more information.
//

// TODO:
//  - move and copy ops should support find best using same logic in service.
// FUTURE:
//  - support restart replace on Windows

package main

import (
	"archive/zip"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/papercutsoftware/silver/lib/pathutils"
)

var (
	versionFile     = flag.String("f", ".version", "Set verison file")
	showVersion     = flag.Bool("v", false, "Display current installed version and exit")
	overrideVersion = flag.String("c", "", "Override current installed version")
	httpProxy       = flag.String("p", "", "Set HTTP proxy in format http://server:port")
)

type UpgradeInfo struct {
	URL        string
	Version    string
	Md5        string
	Sha1       string
	Sha256     string
	Operations []Operation
}

type Operation struct {
	Action string
	Args   []string
}

func usage() {
	exeName := filepath.Base(os.Args[0])
	fmt.Fprintf(os.Stderr, "usage: %s [flags] [update url]\n", exeName)
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {

	flag.Usage = usage
	flag.Parse()

	if *showVersion {
		fmt.Printf("Current version: %s\n", readCurrentVersion())
		return
	}

	if flag.NArg() == 0 {
		usage()
	}
	checkURL := flag.Arg(0)
	if !strings.HasPrefix(strings.ToLower(checkURL), "https") {
		fmt.Fprintf(os.Stderr, "ERROR: The update URL must be HTTPS for security reasons!\n")
		os.Exit(1)
	}

	setupHTTPProxy()
	ok, err := upgradeIfRequired(checkURL)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		os.Exit(1)
	}
	if !ok {
		fmt.Println("No upgrade required")
	} else {
		fmt.Printf("Upgrade successful at %s.\n", time.Now().Format(time.RFC822))
	}
}

func upgradeIfRequired(checkURL string) (upgraded bool, err error) {
	currentVer := readCurrentVersion()
	if len(*overrideVersion) > 0 {
		currentVer = *overrideVersion
	}

	// Ping update URL
	upgradeInfo, err := checkUpdate(checkURL, currentVer)
	if err != nil {
		return false, err
	}

	if upgradeInfo == nil || upgradeInfo.URL == "" {
		// No upgrade required
		return false, nil
	}

	// Download
	fmt.Printf("Downloading version %s update from %s ...\n",
		upgradeInfo.Version,
		upgradeInfo.URL)

	zipfile, err := download(upgradeInfo.URL)
	if err != nil {
		return false, err
	}
	defer os.Remove(zipfile)

	if size, err := fileSize(zipfile); err == nil {
		fmt.Printf("Download complete (%d bytes).\n", size)
	}

	// Validate checksum
	var fileSum string
	var requiredSum string
	switch {
	case len(upgradeInfo.Sha256) > 0:
		requiredSum = upgradeInfo.Sha256
		fileSum = checksum("sha256", zipfile)
	case len(upgradeInfo.Sha1) > 0:
		requiredSum = upgradeInfo.Sha1
		fileSum = checksum("sha1", zipfile)
	default:
		return false, errors.New("Upgrade failed: The upgrade URL did not provide a checksum!")
	}

	if fileSum != requiredSum {
		return false, errors.New("Download checksum failed!")
	}

	// Unzip
	fmt.Println("Unzipping update ...")
	err = extractZip(zipfile, ".")
	if err != nil {
		return false, err
	}
	fmt.Println("Unzip complete.")

	// Perform any operations
	for _, op := range upgradeInfo.Operations {
		action := strings.ToLower(op.Action)
		var fn func([]string) error
		switch action {
		case "exec", "run":
			fn = execOp
		case "batchrename", "batch-rename":
			fn = batchRenameOp
		case "move", "mv":
			fn = moveOp
		case "copy", "cp":
			fn = copyOp
		case "remove", "rm", "del", "delete":
			fn = removeOp
		default:
			msg := fmt.Sprintf("Invalid operation action: '%s'", action)
			return false, errors.New(msg)
		}
		fmt.Printf("Performing operation '%s (%s)' ...\n",
			action, strings.Join(op.Args, ", "))
		if err := fn(op.Args); err != nil {
			msg := fmt.Sprintf("Operation failed with error: %v", err)
			return false, errors.New(msg)
		}
	}

	// Write version file
	ioutil.WriteFile(*versionFile, []byte(upgradeInfo.Version+"\n"), 0644)

	// Request service restart by writing the reload file into our root
	ioutil.WriteFile(".reload", []byte(""), 0644)

	// Success
	return true, nil
}

func fileSize(file string) (size int64, err error) {
	f, err := os.Open(file)
	if err != nil {
		return
	}
	fi, err := f.Stat()
	if err != nil {
		return
	}
	return fi.Size(), nil
}

func checkUpdate(url string, currentVer string) (*UpgradeInfo, error) {
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+"?version="+currentVer, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "Update Check")

	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	if res.StatusCode == http.StatusNotModified {
		return nil, nil
	}

	dec := json.NewDecoder(res.Body)
	var info UpgradeInfo
	err = dec.Decode(&info)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Unable to parse JSON at %s : %v", url, err))
	}

	if info.Version != "" && info.Version == currentVer {
		// Same version!
		return nil, nil
	}

	return &info, nil
}

func download(url string) (string, error) {
	outfile, err := ioutil.TempFile("", "update-")
	if err != nil {
		return "", err
	}

	resp, err := http.Get(url)
	if err != nil {
		outfile.Close()
		os.Remove(outfile.Name())
		return "", err
	}
	defer resp.Body.Close()

	_, err = io.Copy(outfile, resp.Body)
	if err != nil {
		outfile.Close()
		os.Remove(outfile.Name())
		return "", err
	}
	outfile.Close()
	return outfile.Name(), nil
}

func checksum(hashType string, file string) string {
	var hasher hash.Hash

	switch {
	case hashType == "sha256":
		hasher = sha256.New()
	case hashType == "sha1":
		hasher = sha1.New()
	case hashType == "md5":
		hasher = md5.New()
	default:
		hasher = sha1.New()
	}
	f, err := os.Open(file)
	if err != nil {
		panic(err)
	}
	io.Copy(hasher, f)
	return fmt.Sprintf("%x", hasher.Sum(nil))
}

func extractZip(zipfile, dest string) error {
	r, err := zip.OpenReader(zipfile)
	if err != nil {
		return err
	}
	defer r.Close()
	for _, f := range r.File {
		if err := extractZipItem(f, dest); err != nil {
			return err
		}
	}
	return nil
}

func extractZipItem(f *zip.File, dest string) error {
	rc, err := f.Open()
	if err != nil {
		return err
	}
	defer rc.Close()

	path := filepath.Join(dest, f.Name)
	if f.FileInfo().IsDir() {
		os.MkdirAll(path, f.Mode())
	} else {
		f, err := os.OpenFile(
			path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
		if err != nil {
			return err
		}
		defer f.Close()

		_, err = io.Copy(f, rc)
		if err != nil {
			return err
		}
	}
	return nil
}

func readCurrentVersion() string {
	ver := "1"
	if dat, err := ioutil.ReadFile(*versionFile); err == nil {
		ver = strings.TrimSpace(string(dat))
	}
	return ver
}

func setupHTTPProxy() {
	if len(*httpProxy) > 0 {
		os.Setenv("HTTP_PROXY", *httpProxy)
		return
	}
	var proxy = ""
	if dat, err := ioutil.ReadFile("http-proxy.conf"); err == nil {
		proxy = strings.TrimSpace(string(dat))
	}
	if proxy != "" {
		os.Setenv("HTTP_PROXY", proxy)
		return
	}
}

func execOp(args []string) (err error) {
	if len(args) < 1 {
		return errors.New("Invalid exec operation format - arg expected.")
	}
	cmd := args[0]
	fmt.Printf("Running install command: %s\n", strings.Join(args, " "))
	os.Chmod(cmd, 0755)
	c := exec.Command(cmd, args[1:]...)
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr
	err = c.Run()
	return err
}

func batchRenameOp(args []string) error {
	if len(args) != 3 {
		return errors.New("Invalid rename operation format - three args expected.")
	}
	root := args[0]
	find := args[1]
	replacement := args[2]
	fmt.Printf("Running batch rename operation on root %s ('%s' => '%s')\n", root, find, replacement)
	return batchRename(root, find, replacement)
}

func batchRename(root, find, replacement string) error {
	matches, err := filepath.Glob(root)
	if err != nil {
		return err
	}
	if len(matches) == 0 {
		return nil
	}

	re, err := regexp.Compile(find)
	if err != nil {
		return err
	}

	renameCnt := 0
	visitFn := func(path string, fi os.FileInfo, errin error) error {
		name := fi.Name()
		newName := re.ReplaceAllString(name, replacement)
		if name == newName {
			return nil
		}
		if !fileExists(path) {
			return nil
		}

		newPath := filepath.Join(filepath.Dir(path), newName)
		fmt.Printf("Renaming '%s' to '%s' ...\n", path, newPath)
		err := os.Rename(path, newPath)
		if err != nil {
			return err
		}
		renameCnt++
		return nil
	}

	for _, match := range matches {
		err := filepath.Walk(match, visitFn)
		if err != nil {
			return err
		}
	}
	fmt.Printf("Renamed %d files.\n", renameCnt)
	return nil
}

func moveOp(args []string) error {
	if len(args) != 2 {
		return errors.New("Invalid copy operation format - two args expected.")
	}
	src := pathutils.FindLastFile(args[0])
	fmt.Printf("Moving '%s' to '%s'...\n", src, args[1])
	return os.Rename(src, args[1])
}

func removeOp(args []string) error {
	if len(args) != 1 {
		return errors.New("Invalid remove operation format - one arg file expected.")
	}
	path := args[0]
	matches, err := filepath.Glob(path)
	if err != nil {
		return err
	}
	removeCnt := 0
	for _, match := range matches {
		fmt.Printf("Removing '%s' ...\n", match)
		if os.RemoveAll(match) != nil {
			return err
		}
		removeCnt++
	}
	fmt.Printf("Removed %d files.\n", removeCnt)
	return nil
}

func copyOp(args []string) error {
	if len(args) != 2 {
		return errors.New("Invalid copy operation format - two args expected.")
	}
	src := pathutils.FindLastFile(args[0])
	fmt.Printf("Copying '%s' to '%s'...\n", src, args[1])
	return copyFile(src, args[1])
}

func copyFile(src, dest string) error {
	s, err := os.Open(src)
	if err != nil {
		return err
	}
	defer s.Close()
	d, err := os.Create(dest)
	if err != nil {
		return err
	}
	if _, err := io.Copy(d, s); err != nil {
		d.Close()
		return err
	}
	return d.Close()
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return !os.IsNotExist(err)
}
<commit_msg>Added '-unsafe' option to support 'http' protocol for check update url<commit_after>// SILVER - Service Wrapper
// Auto Updater
//
// Copyright (c) 2014 PaperCut Software http://www.papercut.com/
// Use of this source code is governed by an MIT or GPL Version 2 license.
// See the project's LICENSE file for more information.
//

// TODO:
//  - move and copy ops should support find best using same logic in service.
// FUTURE:
//  - support restart replace on Windows

package main

import (
	"archive/zip"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/PaperCutSoftware/silver/lib/pathutils"
)

var (
	versionFile     = flag.String("f", ".version", "Set verison file")
	showVersion     = flag.Bool("v", false, "Display current installed version and exit")
	overrideVersion = flag.String("c", "", "Override current installed version")
	httpProxy       = flag.String("p", "", "Set HTTP proxy in format http://server:port")
	unsafeHTTP      = flag.Bool("unsafe", false, "Debug Only: Support non-https update checks for testing.")
)

type UpgradeInfo struct {
	URL        string
	Version    string
	Md5        string
	Sha1       string
	Sha256     string
	Operations []Operation
}

type Operation struct {
	Action string
	Args   []string
}

func usage() {
	exeName := filepath.Base(os.Args[0])
	fmt.Fprintf(os.Stderr, "usage: %s [flags] [update url]\n", exeName)
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {

	flag.Usage = usage
	flag.Parse()

	if *showVersion {
		fmt.Printf("Current version: %s\n", readCurrentVersion())
		return
	}

	if flag.NArg() == 0 {
		usage()
	}
	checkURL := flag.Arg(0)

	if !*unsafeHTTP && !strings.HasPrefix(strings.ToLower(checkURL), "https") {
		fmt.Fprintf(os.Stderr, "ERROR: The update URL must be HTTPS for security reasons!\n")
		os.Exit(1)
	}

	setupHTTPProxy()
	ok, err := upgradeIfRequired(checkURL)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		os.Exit(1)
	}
	if !ok {
		fmt.Println("No upgrade required")
	} else {
		fmt.Printf("Upgrade successful at %s.\n", time.Now().Format(time.RFC822))
	}
}

func upgradeIfRequired(checkURL string) (upgraded bool, err error) {
	currentVer := readCurrentVersion()
	if len(*overrideVersion) > 0 {
		currentVer = *overrideVersion
	}

	// Ping update URL
	upgradeInfo, err := checkUpdate(checkURL, currentVer)
	if err != nil {
		return false, err
	}

	if upgradeInfo == nil || upgradeInfo.URL == "" {
		// No upgrade required
		return false, nil
	}

	// Download
	fmt.Printf("Downloading version %s update from %s ...\n",
		upgradeInfo.Version,
		upgradeInfo.URL)

	zipfile, err := download(upgradeInfo.URL)
	if err != nil {
		return false, err
	}
	defer os.Remove(zipfile)

	if size, err := fileSize(zipfile); err == nil {
		fmt.Printf("Download complete (%d bytes).\n", size)
	}

	// Validate checksum
	var fileSum string
	var requiredSum string
	switch {
	case len(upgradeInfo.Sha256) > 0:
		requiredSum = upgradeInfo.Sha256
		fileSum = checksum("sha256", zipfile)
	case len(upgradeInfo.Sha1) > 0:
		requiredSum = upgradeInfo.Sha1
		fileSum = checksum("sha1", zipfile)
	default:
		return false, errors.New("Upgrade failed: The upgrade URL did not provide a checksum!")
	}

	if fileSum != requiredSum {
		return false, errors.New("Download checksum failed!")
	}

	// Unzip
	fmt.Println("Unzipping update ...")
	err = extractZip(zipfile, ".")
	if err != nil {
		return false, err
	}
	fmt.Println("Unzip complete.")

	// Perform any operations
	for _, op := range upgradeInfo.Operations {
		action := strings.ToLower(op.Action)
		var fn func([]string) error
		switch action {
		case "exec", "run":
			fn = execOp
		case "batchrename", "batch-rename":
			fn = batchRenameOp
		case "move", "mv":
			fn = moveOp
		case "copy", "cp":
			fn = copyOp
		case "remove", "rm", "del", "delete":
			fn = removeOp
		default:
			msg := fmt.Sprintf("Invalid operation action: '%s'", action)
			return false, errors.New(msg)
		}
		fmt.Printf("Performing operation '%s (%s)' ...\n",
			action, strings.Join(op.Args, ", "))
		if err := fn(op.Args); err != nil {
			msg := fmt.Sprintf("Operation failed with error: %v", err)
			return false, errors.New(msg)
		}
	}

	// Write version file
	ioutil.WriteFile(*versionFile, []byte(upgradeInfo.Version+"\n"), 0644)

	// Request service restart by writing the reload file into our root
	ioutil.WriteFile(".reload", []byte(""), 0644)

	// Success
	return true, nil
}

func fileSize(file string) (size int64, err error) {
	f, err := os.Open(file)
	if err != nil {
		return 0, err
	}
	fi, err := f.Stat()
	if err != nil {
		return 0, err
	}
	return fi.Size(), nil
}

func checkUpdate(url string, currentVer string) (*UpgradeInfo, error) {
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+"?version="+currentVer, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "Update Check")

	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	if res.StatusCode == http.StatusNotModified {
		return nil, nil
	}

	dec := json.NewDecoder(res.Body)
	var info UpgradeInfo
	err = dec.Decode(&info)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Unable to parse JSON at %s : %v", url, err))
	}

	if info.Version != "" && info.Version == currentVer {
		// Same version!
		return nil, nil
	}

	return &info, nil
}

func download(url string) (string, error) {
	outfile, err := ioutil.TempFile("", "update-")
	if err != nil {
		return "", err
	}

	resp, err := http.Get(url)
	if err != nil {
		outfile.Close()
		os.Remove(outfile.Name())
		return "", err
	}
	defer resp.Body.Close()

	_, err = io.Copy(outfile, resp.Body)
	if err != nil {
		outfile.Close()
		os.Remove(outfile.Name())
		return "", err
	}
	outfile.Close()
	return outfile.Name(), nil
}

func checksum(hashType string, file string) string {
	var hasher hash.Hash

	switch {
	case hashType == "sha256":
		hasher = sha256.New()
	case hashType == "sha1":
		hasher = sha1.New()
	case hashType == "md5":
		hasher = md5.New()
	default:
		hasher = sha1.New()
	}
	f, err := os.Open(file)
	if err != nil {
		panic(err)
	}
	io.Copy(hasher, f)
	return fmt.Sprintf("%x", hasher.Sum(nil))
}

func extractZip(zipfile, dest string) error {
	r, err := zip.OpenReader(zipfile)
	if err != nil {
		return err
	}
	defer r.Close()
	for _, f := range r.File {
		if err := extractZipItem(f, dest); err != nil {
			return err
		}
	}
	return nil
}

func extractZipItem(f *zip.File, dest string) error {
	rc, err := f.Open()
	if err != nil {
		return err
	}
	defer rc.Close()

	path := filepath.Join(dest, f.Name)
	if f.FileInfo().IsDir() {
		os.MkdirAll(path, f.Mode())
	} else {
		f, err := os.OpenFile(
			path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
		if err != nil {
			return err
		}
		defer f.Close()

		_, err = io.Copy(f, rc)
		if err != nil {
			return err
		}
	}
	return nil
}

func readCurrentVersion() string {
	ver := "1"
	if dat, err := ioutil.ReadFile(*versionFile); err == nil {
		ver = strings.TrimSpace(string(dat))
	}
	return ver
}

func setupHTTPProxy() {
	if len(*httpProxy) > 0 {
		os.Setenv("HTTP_PROXY", *httpProxy)
		return
	}
	var proxy = ""
	if dat, err := ioutil.ReadFile("http-proxy.conf"); err == nil {
		proxy = strings.TrimSpace(string(dat))
	}
	if proxy != "" {
		os.Setenv("HTTP_PROXY", proxy)
		return
	}
}

func execOp(args []string) (err error) {
	if len(args) < 1 {
		return errors.New("Invalid exec operation format - arg expected.")
	}
	cmd := args[0]
	fmt.Printf("Running install command: %s\n", strings.Join(args, " "))
	os.Chmod(cmd, 0755)
	c := exec.Command(cmd, args[1:]...)
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr
	err = c.Run()
	return err
}

func batchRenameOp(args []string) error {
	if len(args) != 3 {
		return errors.New("Invalid rename operation format - three args expected.")
	}
	root := args[0]
	find := args[1]
	replacement := args[2]
	fmt.Printf("Running batch rename operation on root %s ('%s' => '%s')\n", root, find, replacement)
	return batchRename(root, find, replacement)
}

func batchRename(root, find, replacement string) error {
	matches, err := filepath.Glob(root)
	if err != nil {
		return err
	}
	if len(matches) == 0 {
		return nil
	}

	re, err := regexp.Compile(find)
	if err != nil {
		return err
	}

	renameCnt := 0
	visitFn := func(path string, fi os.FileInfo, errin error) error {
		name := fi.Name()
		newName := re.ReplaceAllString(name, replacement)
		if name == newName {
			return nil
		}
		if !fileExists(path) {
			return nil
		}

		newPath := filepath.Join(filepath.Dir(path), newName)
		fmt.Printf("Renaming '%s' to '%s' ...\n", path, newPath)
		err := os.Rename(path, newPath)
		if err != nil {
			return err
		}
		renameCnt++
		return nil
	}

	for _, match := range matches {
		err := filepath.Walk(match, visitFn)
		if err != nil {
			return err
		}
	}
	fmt.Printf("Renamed %d files.\n", renameCnt)
	return nil
}

func moveOp(args []string) error {
	if len(args) != 2 {
		return errors.New("Invalid copy operation format - two args expected.")
	}
	src := pathutils.FindLastFile(args[0])
	fmt.Printf("Moving '%s' to '%s'...\n", src, args[1])
	return os.Rename(src, args[1])
}

func removeOp(args []string) error {
	if len(args) != 1 {
		return errors.New("Invalid remove operation format - one arg file expected.")
	}
	path := args[0]
	matches, err := filepath.Glob(path)
	if err != nil {
		return err
	}
	removeCnt := 0
	for _, match := range matches {
		fmt.Printf("Removing '%s' ...\n", match)
		if os.RemoveAll(match) != nil {
			return err
		}
		removeCnt++
	}
	fmt.Printf("Removed %d files.\n", removeCnt)
	return nil
}

func copyOp(args []string) error {
	if len(args) != 2 {
		return errors.New("Invalid copy operation format - two args expected.")
	}
	src := pathutils.FindLastFile(args[0])
	fmt.Printf("Copying '%s' to '%s'...\n", src, args[1])
	return copyFile(src, args[1])
}

func copyFile(src, dest string) error {
	s, err := os.Open(src)
	if err != nil {
		return err
	}
	defer s.Close()
	d, err := os.Create(dest)
	if err != nil {
		return err
	}
	if _, err := io.Copy(d, s); err != nil {
		d.Close()
		return err
	}
	return d.Close()
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return !os.IsNotExist(err)
}
<|endoftext|>


<commit_before>// Copyright 2015 Unknwon
//
// 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 cmd

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

	"github.com/Unknwon/com"
	"github.com/codegangsta/cli"
	"gopkg.in/ini.v1"

	"github.com/peachdocs/peach/modules/bindata"
)

var New = cli.Command{
	Name:   "new",
	Usage:  "Initialize a new Peach project",
	Action: runNew,
	Flags: []cli.Flag{
		stringFlag("target, t", "my.peach", "Directory to save project files"),
		boolFlag("yes, y", "Yes to all confirmations"),
	},
}

func checkYesNo() bool {
	var choice string
	fmt.Scan(&choice)
	return strings.HasPrefix(strings.ToLower(choice), "y")
}

func toRed(str string) string {
	return fmt.Sprintf("\033[31m%s\033[0m", str)
}

func toGreen(str string) string {
	return fmt.Sprintf("\033[32m%s\033[0m", str)
}

func toYellow(str string) string {
	return fmt.Sprintf("\033[33m%s\033[0m", str)
}

func restoreAssets(target, dir string) {
	if err := bindata.RestoreAssets(target, dir); err != nil {
		fmt.Printf(toRed("✗  %v\n"), err)
		os.Exit(1)
	}
}

func runNew(ctx *cli.Context) {
	target := ctx.String("target")
	if com.IsExist(target) && !ctx.Bool("yes") {
		fmt.Printf(toYellow("Directory '%s' already exists, do you want to overwrite?[Y/n] "), target)
		if !checkYesNo() {
			os.Exit(0)
		}
	}

	fmt.Printf("➜  Creating '%s'...\n", target)
	os.MkdirAll(target, os.ModePerm)

	// Create default files.
	dirs := []string{"templates", "public"}
	for _, dir := range dirs {
		fmt.Printf("➜  Creating '%s'...\n", dir)
		os.RemoveAll(filepath.Join(target, dir))
		restoreAssets(target, dir)
	}

	// Create custom templates.
	yes := ctx.Bool("yes")
	if !yes {
		fmt.Printf(toYellow("Do you want to use custom templates?[Y/n] "))
		yes = checkYesNo()
	}

	if yes {
		fmt.Println("➜  Creating 'custom/templates'...")
		restoreAssets(filepath.Join(target, "custom"), "templates")

		// Update configuration to use custom templates.
		fmt.Println("➜  Updating custom configuration...")
		var cfg *ini.File
		var err error
		customPath := filepath.Join(target, "custom/app.ini")
		if com.IsExist(customPath) {
			cfg, err = ini.Load(customPath)
			if err != nil {
				fmt.Printf(toRed("✗  %v\n"), err)
				os.Exit(1)
			}
		} else {
			cfg = ini.Empty()
		}

		cfg.Section("page").Key("USE_CUSTOM_TPL").SetValue("true")
		if err = cfg.SaveTo(customPath); err != nil {
			fmt.Printf(toRed("✗  %v\n"), err)
			os.Exit(1)
		}
	}

	fmt.Println(toGreen("✓  Done!"))
}
<commit_msg>update - [cmd/init.go] 关键模块, 功能: 创建 peach 初始化目录.<commit_after>// Copyright 2015 Unknwon
//
// 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 cmd

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

	"github.com/Unknwon/com"
	"github.com/codegangsta/cli"
	"gopkg.in/ini.v1"

	"github.com/peachdocs/peach/modules/bindata"
)

/*
	注册的 peach 命令:
	- peach new
	- 此命令用来创建一个 peach 模板项目结构
	- 初始化的原型.
 */
var New = cli.Command{
	Name:   "new",
	Usage:  "Initialize a new Peach project",
	Action: runNew,
	Flags: []cli.Flag{
		stringFlag("target, t", "my.peach", "Directory to save project files"),
		boolFlag("yes, y", "Yes to all confirmations"),
	},
}

func checkYesNo() bool {
	var choice string
	fmt.Scan(&choice)
	return strings.HasPrefix(strings.ToLower(choice), "y")
}

func toRed(str string) string {
	return fmt.Sprintf("\033[31m%s\033[0m", str)
}

func toGreen(str string) string {
	return fmt.Sprintf("\033[32m%s\033[0m", str)
}

func toYellow(str string) string {
	return fmt.Sprintf("\033[33m%s\033[0m", str)
}

func restoreAssets(target, dir string) {
	if err := bindata.RestoreAssets(target, dir); err != nil {
		fmt.Printf(toRed("✗  %v\n"), err)
		os.Exit(1)
	}
}

/*
	关键入口:
	- 执行创建 peach 目录原型结构.
 */
func runNew(ctx *cli.Context) {
	// 参数提取
	target := ctx.String("target")
	if com.IsExist(target) && !ctx.Bool("yes") {
		fmt.Printf(toYellow("Directory '%s' already exists, do you want to overwrite?[Y/n] "), target)
		if !checkYesNo() {
			os.Exit(0)
		}
	}

	fmt.Printf("➜  Creating '%s'...\n", target)
	os.MkdirAll(target, os.ModePerm)			// 目录创建

	// Create default files.
	dirs := []string{"templates", "public"}
	for _, dir := range dirs {
		fmt.Printf("➜  Creating '%s'...\n", dir)
		os.RemoveAll(filepath.Join(target, dir))
		restoreAssets(target, dir)
	}

	// Create custom templates.
	yes := ctx.Bool("yes")
	if !yes {
		fmt.Printf(toYellow("Do you want to use custom templates?[Y/n] "))
		yes = checkYesNo()
	}

	if yes {
		fmt.Println("➜  Creating 'custom/templates'...")
		restoreAssets(filepath.Join(target, "custom"), "templates")

		// Update configuration to use custom templates.
		fmt.Println("➜  Updating custom configuration...")
		var cfg *ini.File
		var err error
		customPath := filepath.Join(target, "custom/app.ini")
		if com.IsExist(customPath) {
			cfg, err = ini.Load(customPath)
			if err != nil {
				fmt.Printf(toRed("✗  %v\n"), err)
				os.Exit(1)
			}
		} else {
			cfg = ini.Empty()
		}

		cfg.Section("page").Key("USE_CUSTOM_TPL").SetValue("true")
		if err = cfg.SaveTo(customPath); err != nil {
			fmt.Printf(toRed("✗  %v\n"), err)
			os.Exit(1)
		}
	}

	fmt.Println(toGreen("✓  Done!"))
}
<|endoftext|>


<commit_before>package storage

import (
	"io"
	"os"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
)

type NeedleMap struct {
	m needle.NeedleValueMap

	baseNeedleMapper
}

func NewCompactNeedleMap(file *os.File) *NeedleMap {
	nm := &NeedleMap{
		m: needle.NewCompactMap(),
	}
	nm.indexFile = file
	return nm
}

func NewBtreeNeedleMap(file *os.File) *NeedleMap {
	nm := &NeedleMap{
		m: needle.NewBtreeMap(),
	}
	nm.indexFile = file
	return nm
}

const (
	RowsToRead = 1024
)

func LoadCompactNeedleMap(file *os.File) (*NeedleMap, error) {
	nm := NewCompactNeedleMap(file)
	return doLoading(file, nm)
}

func LoadBtreeNeedleMap(file *os.File) (*NeedleMap, error) {
	nm := NewBtreeNeedleMap(file)
	return doLoading(file, nm)
}

func doLoading(file *os.File, nm *NeedleMap) (*NeedleMap, error) {
	e := WalkIndexFile(file, func(key uint64, offset, size uint32) error {
		if key > nm.MaximumFileKey {
			nm.MaximumFileKey = key
		}
		if offset > 0 && size != TombstoneFileSize {
			nm.FileCounter++
			nm.FileByteCounter = nm.FileByteCounter + uint64(size)
			oldOffset, oldSize := nm.m.Set(needle.Key(key), offset, size)
			glog.V(3).Infoln("reading key", key, "offset", offset*NeedlePaddingSize, "size", size, "oldSize", oldSize)
			if oldOffset > 0 && oldSize != TombstoneFileSize {
				nm.DeletionCounter++
				nm.DeletionByteCounter = nm.DeletionByteCounter + uint64(oldSize)
			}
		} else {
			oldSize := nm.m.Delete(needle.Key(key))
			glog.V(3).Infoln("removing key", key, "offset", offset*NeedlePaddingSize, "size", size, "oldSize", oldSize)
			nm.DeletionCounter++
			nm.DeletionByteCounter = nm.DeletionByteCounter + uint64(oldSize)
		}
		return nil
	})
	glog.V(1).Infof("max file key: %d for file: %s", nm.MaximumFileKey, file.Name())
	return nm, e
}

// walks through the index file, calls fn function with each key, offset, size
// stops with the error returned by the fn function
func WalkIndexFile(r *os.File, fn func(key uint64, offset, size uint32) error) error {
	var readerOffset int64
	bytes := make([]byte, 16*RowsToRead)
	count, e := r.ReadAt(bytes, readerOffset)
	glog.V(3).Infoln("file", r.Name(), "readerOffset", readerOffset, "count", count, "e", e)
	readerOffset += int64(count)
	var (
		key          uint64
		offset, size uint32
		i            int
	)

	for count > 0 && e == nil || e == io.EOF {
		for i = 0; i+16 <= count; i += 16 {
			key, offset, size = idxFileEntry(bytes[i : i+16])
			if e = fn(key, offset, size); e != nil {
				return e
			}
		}
		if e == io.EOF {
			return nil
		}
		count, e = r.ReadAt(bytes, readerOffset)
		glog.V(3).Infoln("file", r.Name(), "readerOffset", readerOffset, "count", count, "e", e)
		readerOffset += int64(count)
	}
	return e
}

func (nm *NeedleMap) Put(key uint64, offset uint32, size uint32) error {
	_, oldSize := nm.m.Set(needle.Key(key), offset, size)
	nm.logPut(key, oldSize, size)
	return nm.appendToIndexFile(key, offset, size)
}
func (nm *NeedleMap) Get(key uint64) (element *needle.NeedleValue, ok bool) {
	element, ok = nm.m.Get(needle.Key(key))
	return
}
func (nm *NeedleMap) Delete(key uint64, offset uint32) error {
	deletedBytes := nm.m.Delete(needle.Key(key))
	nm.logDelete(deletedBytes)
	return nm.appendToIndexFile(key, offset, TombstoneFileSize)
}
func (nm *NeedleMap) Close() {
	_ = nm.indexFile.Close()
}
func (nm *NeedleMap) Destroy() error {
	nm.Close()
	return os.Remove(nm.indexFile.Name())
}
<commit_msg>maybe speed up a little when loading index<commit_after>package storage

import (
	"io"
	"os"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
)

type NeedleMap struct {
	m needle.NeedleValueMap

	baseNeedleMapper
}

func NewCompactNeedleMap(file *os.File) *NeedleMap {
	nm := &NeedleMap{
		m: needle.NewCompactMap(),
	}
	nm.indexFile = file
	return nm
}

func NewBtreeNeedleMap(file *os.File) *NeedleMap {
	nm := &NeedleMap{
		m: needle.NewBtreeMap(),
	}
	nm.indexFile = file
	return nm
}

const (
	RowsToRead = 1024
)

func LoadCompactNeedleMap(file *os.File) (*NeedleMap, error) {
	nm := NewCompactNeedleMap(file)
	return doLoading(file, nm)
}

func LoadBtreeNeedleMap(file *os.File) (*NeedleMap, error) {
	nm := NewBtreeNeedleMap(file)
	return doLoading(file, nm)
}

func doLoading(file *os.File, nm *NeedleMap) (*NeedleMap, error) {
	e := WalkIndexFile(file, func(key uint64, offset, size uint32) error {
		if key > nm.MaximumFileKey {
			nm.MaximumFileKey = key
		}
		if offset > 0 && size != TombstoneFileSize {
			nm.FileCounter++
			nm.FileByteCounter = nm.FileByteCounter + uint64(size)
			oldOffset, oldSize := nm.m.Set(needle.Key(key), offset, size)
			// glog.V(3).Infoln("reading key", key, "offset", offset*NeedlePaddingSize, "size", size, "oldSize", oldSize)
			if oldOffset > 0 && oldSize != TombstoneFileSize {
				nm.DeletionCounter++
				nm.DeletionByteCounter = nm.DeletionByteCounter + uint64(oldSize)
			}
		} else {
			oldSize := nm.m.Delete(needle.Key(key))
			// glog.V(3).Infoln("removing key", key, "offset", offset*NeedlePaddingSize, "size", size, "oldSize", oldSize)
			nm.DeletionCounter++
			nm.DeletionByteCounter = nm.DeletionByteCounter + uint64(oldSize)
		}
		return nil
	})
	glog.V(1).Infof("max file key: %d for file: %s", nm.MaximumFileKey, file.Name())
	return nm, e
}

// walks through the index file, calls fn function with each key, offset, size
// stops with the error returned by the fn function
func WalkIndexFile(r *os.File, fn func(key uint64, offset, size uint32) error) error {
	var readerOffset int64
	bytes := make([]byte, 16*RowsToRead)
	count, e := r.ReadAt(bytes, readerOffset)
	glog.V(3).Infoln("file", r.Name(), "readerOffset", readerOffset, "count", count, "e", e)
	readerOffset += int64(count)
	var (
		key          uint64
		offset, size uint32
		i            int
	)

	for count > 0 && e == nil || e == io.EOF {
		for i = 0; i+16 <= count; i += 16 {
			key, offset, size = idxFileEntry(bytes[i: i+16])
			if e = fn(key, offset, size); e != nil {
				return e
			}
		}
		if e == io.EOF {
			return nil
		}
		count, e = r.ReadAt(bytes, readerOffset)
		glog.V(3).Infoln("file", r.Name(), "readerOffset", readerOffset, "count", count, "e", e)
		readerOffset += int64(count)
	}
	return e
}

func (nm *NeedleMap) Put(key uint64, offset uint32, size uint32) error {
	_, oldSize := nm.m.Set(needle.Key(key), offset, size)
	nm.logPut(key, oldSize, size)
	return nm.appendToIndexFile(key, offset, size)
}
func (nm *NeedleMap) Get(key uint64) (element *needle.NeedleValue, ok bool) {
	element, ok = nm.m.Get(needle.Key(key))
	return
}
func (nm *NeedleMap) Delete(key uint64, offset uint32) error {
	deletedBytes := nm.m.Delete(needle.Key(key))
	nm.logDelete(deletedBytes)
	return nm.appendToIndexFile(key, offset, TombstoneFileSize)
}
func (nm *NeedleMap) Close() {
	_ = nm.indexFile.Close()
}
func (nm *NeedleMap) Destroy() error {
	nm.Close()
	return os.Remove(nm.indexFile.Name())
}
<|endoftext|>


<commit_before>package build

import (
	"bytes"
	"flag"
	"fmt"
	cmdcommon "github.com/mitchellh/packer/common/command"
	"github.com/mitchellh/packer/packer"
	"log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
)

type Command byte

func (Command) Help() string {
	return strings.TrimSpace(helpText)
}

func (c Command) Run(env packer.Environment, args []string) int {
	var cfgDebug bool
	var cfgForce bool
	buildOptions := new(cmdcommon.BuildOptions)

	cmdFlags := flag.NewFlagSet("build", flag.ContinueOnError)
	cmdFlags.Usage = func() { env.Ui().Say(c.Help()) }
	cmdFlags.BoolVar(&cfgDebug, "debug", false, "debug mode for builds")
	cmdFlags.BoolVar(&cfgForce, "force", false, "force a build if artifacts exist")
	cmdcommon.BuildOptionFlags(cmdFlags, buildOptions)
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	args = cmdFlags.Args()
	if len(args) != 1 {
		cmdFlags.Usage()
		return 1
	}

	if err := buildOptions.Validate(); err != nil {
		env.Ui().Error(err.Error())
		env.Ui().Error("")
		env.Ui().Error(c.Help())
		return 1
	}

	userVars, err := buildOptions.AllUserVars()
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Error compiling user variables: %s", err))
		env.Ui().Error("")
		env.Ui().Error(c.Help())
		return 1
	}

	// Read the file into a byte array so that we can parse the template
	log.Printf("Reading template: %s", args[0])
	tpl, err := packer.ParseTemplateFile(args[0])
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Failed to parse template: %s", err))
		return 1
	}

	// The component finder for our builds
	components := &packer.ComponentFinder{
		Builder:       env.Builder,
		Hook:          env.Hook,
		PostProcessor: env.PostProcessor,
		Provisioner:   env.Provisioner,
	}

	// Go through each builder and compile the builds that we care about
	builds, err := buildOptions.Builds(tpl, components)
	if err != nil {
		env.Ui().Error(err.Error())
		return 1
	}

	if cfgDebug {
		env.Ui().Say("Debug mode enabled. Builds will not be parallelized.")
	}

	// Compile all the UIs for the builds
	colors := [5]packer.UiColor{
		packer.UiColorGreen,
		packer.UiColorCyan,
		packer.UiColorMagenta,
		packer.UiColorYellow,
		packer.UiColorBlue,
	}

	buildUis := make(map[string]packer.Ui)
	for i, b := range builds {
		ui := &packer.ColoredUi{
			Color: colors[i%len(colors)],
			Ui:    env.Ui(),
		}

		buildUis[b.Name()] = ui
		ui.Say(fmt.Sprintf("%s output will be in this color.", b.Name()))
	}

	// Add a newline between the color output and the actual output
	env.Ui().Say("")

	log.Printf("Build debug mode: %v", cfgDebug)
	log.Printf("Force build: %v", cfgForce)

	// Set the debug and force mode and prepare all the builds
	for _, b := range builds {
		log.Printf("Preparing build: %s", b.Name())
		b.SetDebug(cfgDebug)
		b.SetForce(cfgForce)
		err := b.Prepare(userVars)
		if err != nil {
			env.Ui().Error(err.Error())
			return 1
		}
	}

	// Run all the builds in parallel and wait for them to complete
	var interruptWg, wg sync.WaitGroup
	interrupted := false
	artifacts := make(map[string][]packer.Artifact)
	errors := make(map[string]error)
	for _, b := range builds {
		// Increment the waitgroup so we wait for this item to finish properly
		wg.Add(1)

		// Handle interrupts for this build
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, os.Interrupt)
		defer signal.Stop(sigCh)
		go func(b packer.Build) {
			<-sigCh
			interruptWg.Add(1)
			defer interruptWg.Done()
			interrupted = true

			log.Printf("Stopping build: %s", b.Name())
			b.Cancel()
			log.Printf("Build cancelled: %s", b.Name())
		}(b)

		// Run the build in a goroutine
		go func(b packer.Build) {
			defer wg.Done()

			name := b.Name()
			log.Printf("Starting build run: %s", name)
			ui := buildUis[name]
			runArtifacts, err := b.Run(ui, env.Cache())

			if err != nil {
				ui.Error(fmt.Sprintf("Build '%s' errored: %s", name, err))
				errors[name] = err
			} else {
				ui.Say(fmt.Sprintf("Build '%s' finished.", name))
				artifacts[name] = runArtifacts
			}
		}(b)

		if cfgDebug {
			log.Printf("Debug enabled, so waiting for build to finish: %s", b.Name())
			wg.Wait()
		}

		if interrupted {
			log.Println("Interrupted, not going to start any more builds.")
			break
		}
	}

	// Wait for both the builds to complete and the interrupt handler,
	// if it is interrupted.
	log.Printf("Waiting on builds to complete...")
	wg.Wait()

	log.Printf("Builds completed. Waiting on interrupt barrier...")
	interruptWg.Wait()

	if interrupted {
		env.Ui().Say("Cleanly cancelled builds after being interrupted.")
		return 1
	}

	if len(errors) > 0 {
		env.Ui().Machine("error-count", strconv.FormatInt(int64(len(errors)), 10))

		env.Ui().Error("\n==> Some builds didn't complete successfully and had errors:")
		for name, err := range errors {
			// Create a UI for the machine readable stuff to be targetted
			ui := &packer.TargettedUi{
				Target: name,
				Ui:     env.Ui(),
			}

			ui.Machine("error", err.Error())

			env.Ui().Error(fmt.Sprintf("--> %s: %s", name, err))
		}
	}

	if len(artifacts) > 0 {
		env.Ui().Say("\n==> Builds finished. The artifacts of successful builds are:")
		for name, buildArtifacts := range artifacts {
			// Create a UI for the machine readable stuff to be targetted
			ui := &packer.TargettedUi{
				Target: name,
				Ui:     env.Ui(),
			}

			// Machine-readable helpful
			ui.Machine("artifact-count", strconv.FormatInt(int64(len(buildArtifacts)), 10))

			for i, artifact := range buildArtifacts {
				var message bytes.Buffer
				fmt.Fprintf(&message, "--> %s: ", name)

				if artifact != nil {
					fmt.Fprintf(&message, artifact.String())
				} else {
					fmt.Fprint(&message, "<nothing>")
				}

				iStr := strconv.FormatInt(int64(i), 10)
				if artifact != nil {
					ui.Machine("artifact", iStr, "builder-id", artifact.BuilderId())
					ui.Machine("artifact", iStr, "id", artifact.Id())
					ui.Machine("artifact", iStr, "string", artifact.String())

					files := artifact.Files()
					ui.Machine("artifact",
						iStr,
						"files-count", strconv.FormatInt(int64(len(files)), 10))
					for fi, file := range files {
						fiStr := strconv.FormatInt(int64(fi), 10)
						ui.Machine("artifact", iStr, "file", fiStr, file)
					}
				} else {
					ui.Machine("artifact", iStr, "nil")
				}

				ui.Machine("artifact", iStr, "end")
				env.Ui().Say(message.String())
			}
		}
	} else {
		env.Ui().Say("\n==> Builds finished but no artifacts were created.")
	}

	if len(errors) > 0 {
		// If any errors occurred, exit with a non-zero exit status
		return 1
	}

	return 0
}

func (Command) Synopsis() string {
	return "build image(s) from template"
}
<commit_msg>command/build: get command passing<commit_after>package build

import (
	"bytes"
	"flag"
	"fmt"
	cmdcommon "github.com/mitchellh/packer/common/command"
	"github.com/mitchellh/packer/packer"
	"log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
)

type Command byte

func (Command) Help() string {
	return strings.TrimSpace(helpText)
}

func (c Command) Run(env packer.Environment, args []string) int {
	var cfgDebug bool
	var cfgForce bool
	buildOptions := new(cmdcommon.BuildOptions)

	cmdFlags := flag.NewFlagSet("build", flag.ContinueOnError)
	cmdFlags.Usage = func() { env.Ui().Say(c.Help()) }
	cmdFlags.BoolVar(&cfgDebug, "debug", false, "debug mode for builds")
	cmdFlags.BoolVar(&cfgForce, "force", false, "force a build if artifacts exist")
	cmdcommon.BuildOptionFlags(cmdFlags, buildOptions)
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	args = cmdFlags.Args()
	if len(args) != 1 {
		cmdFlags.Usage()
		return 1
	}

	if err := buildOptions.Validate(); err != nil {
		env.Ui().Error(err.Error())
		env.Ui().Error("")
		env.Ui().Error(c.Help())
		return 1
	}

	userVars, err := buildOptions.AllUserVars()
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Error compiling user variables: %s", err))
		env.Ui().Error("")
		env.Ui().Error(c.Help())
		return 1
	}

	// Read the file into a byte array so that we can parse the template
	log.Printf("Reading template: %s", args[0])
	tpl, err := packer.ParseTemplateFile(args[0])
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Failed to parse template: %s", err))
		return 1
	}

	// The component finder for our builds
	components := &packer.ComponentFinder{
		Builder:       env.Builder,
		Hook:          env.Hook,
		PostProcessor: env.PostProcessor,
		Provisioner:   env.Provisioner,
	}

	// Go through each builder and compile the builds that we care about
	builds, err := buildOptions.Builds(tpl, components)
	if err != nil {
		env.Ui().Error(err.Error())
		return 1
	}

	if cfgDebug {
		env.Ui().Say("Debug mode enabled. Builds will not be parallelized.")
	}

	// Compile all the UIs for the builds
	colors := [5]packer.UiColor{
		packer.UiColorGreen,
		packer.UiColorCyan,
		packer.UiColorMagenta,
		packer.UiColorYellow,
		packer.UiColorBlue,
	}

	buildUis := make(map[string]packer.Ui)
	for i, b := range builds {
		ui := &packer.ColoredUi{
			Color: colors[i%len(colors)],
			Ui:    env.Ui(),
		}

		buildUis[b.Name()] = ui
		ui.Say(fmt.Sprintf("%s output will be in this color.", b.Name()))
	}

	// Add a newline between the color output and the actual output
	env.Ui().Say("")

	log.Printf("Build debug mode: %v", cfgDebug)
	log.Printf("Force build: %v", cfgForce)

	// Set the debug and force mode and prepare all the builds
	for _, b := range builds {
		log.Printf("Preparing build: %s", b.Name())
		b.SetDebug(cfgDebug)
		b.SetForce(cfgForce)
		// TODO(mitchellh): Handle warnings
		_, err := b.Prepare(userVars)
		if err != nil {
			env.Ui().Error(err.Error())
			return 1
		}
	}

	// Run all the builds in parallel and wait for them to complete
	var interruptWg, wg sync.WaitGroup
	interrupted := false
	artifacts := make(map[string][]packer.Artifact)
	errors := make(map[string]error)
	for _, b := range builds {
		// Increment the waitgroup so we wait for this item to finish properly
		wg.Add(1)

		// Handle interrupts for this build
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, os.Interrupt)
		defer signal.Stop(sigCh)
		go func(b packer.Build) {
			<-sigCh
			interruptWg.Add(1)
			defer interruptWg.Done()
			interrupted = true

			log.Printf("Stopping build: %s", b.Name())
			b.Cancel()
			log.Printf("Build cancelled: %s", b.Name())
		}(b)

		// Run the build in a goroutine
		go func(b packer.Build) {
			defer wg.Done()

			name := b.Name()
			log.Printf("Starting build run: %s", name)
			ui := buildUis[name]
			runArtifacts, err := b.Run(ui, env.Cache())

			if err != nil {
				ui.Error(fmt.Sprintf("Build '%s' errored: %s", name, err))
				errors[name] = err
			} else {
				ui.Say(fmt.Sprintf("Build '%s' finished.", name))
				artifacts[name] = runArtifacts
			}
		}(b)

		if cfgDebug {
			log.Printf("Debug enabled, so waiting for build to finish: %s", b.Name())
			wg.Wait()
		}

		if interrupted {
			log.Println("Interrupted, not going to start any more builds.")
			break
		}
	}

	// Wait for both the builds to complete and the interrupt handler,
	// if it is interrupted.
	log.Printf("Waiting on builds to complete...")
	wg.Wait()

	log.Printf("Builds completed. Waiting on interrupt barrier...")
	interruptWg.Wait()

	if interrupted {
		env.Ui().Say("Cleanly cancelled builds after being interrupted.")
		return 1
	}

	if len(errors) > 0 {
		env.Ui().Machine("error-count", strconv.FormatInt(int64(len(errors)), 10))

		env.Ui().Error("\n==> Some builds didn't complete successfully and had errors:")
		for name, err := range errors {
			// Create a UI for the machine readable stuff to be targetted
			ui := &packer.TargettedUi{
				Target: name,
				Ui:     env.Ui(),
			}

			ui.Machine("error", err.Error())

			env.Ui().Error(fmt.Sprintf("--> %s: %s", name, err))
		}
	}

	if len(artifacts) > 0 {
		env.Ui().Say("\n==> Builds finished. The artifacts of successful builds are:")
		for name, buildArtifacts := range artifacts {
			// Create a UI for the machine readable stuff to be targetted
			ui := &packer.TargettedUi{
				Target: name,
				Ui:     env.Ui(),
			}

			// Machine-readable helpful
			ui.Machine("artifact-count", strconv.FormatInt(int64(len(buildArtifacts)), 10))

			for i, artifact := range buildArtifacts {
				var message bytes.Buffer
				fmt.Fprintf(&message, "--> %s: ", name)

				if artifact != nil {
					fmt.Fprintf(&message, artifact.String())
				} else {
					fmt.Fprint(&message, "<nothing>")
				}

				iStr := strconv.FormatInt(int64(i), 10)
				if artifact != nil {
					ui.Machine("artifact", iStr, "builder-id", artifact.BuilderId())
					ui.Machine("artifact", iStr, "id", artifact.Id())
					ui.Machine("artifact", iStr, "string", artifact.String())

					files := artifact.Files()
					ui.Machine("artifact",
						iStr,
						"files-count", strconv.FormatInt(int64(len(files)), 10))
					for fi, file := range files {
						fiStr := strconv.FormatInt(int64(fi), 10)
						ui.Machine("artifact", iStr, "file", fiStr, file)
					}
				} else {
					ui.Machine("artifact", iStr, "nil")
				}

				ui.Machine("artifact", iStr, "end")
				env.Ui().Say(message.String())
			}
		}
	} else {
		env.Ui().Say("\n==> Builds finished but no artifacts were created.")
	}

	if len(errors) > 0 {
		// If any errors occurred, exit with a non-zero exit status
		return 1
	}

	return 0
}

func (Command) Synopsis() string {
	return "build image(s) from template"
}
<|endoftext|>


<commit_before>package commands

import (
	"bufio"
	"fmt"
	"github.com/jingweno/gh/cmd"
	"github.com/jingweno/gh/git"
	"github.com/jingweno/gh/github"
	"github.com/jingweno/gh/utils"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"strings"
)

var cmdPullRequest = &Command{
	Run:   pullRequest,
  Usage: "pull-request [-f] [-i ISSUE] [-b BASE] [-d HEAD] [TITLE]",
	Short: "Open a pull request on GitHub",
	Long: `Opens a pull request on GitHub for the project that the "origin" remote
points to. The default head of the pull request is the current branch.
Both base and head of the pull request can be explicitly given in one of
the following formats: "branch", "owner:branch", "owner/repo:branch".
This command will abort operation if it detects that the current topic
branch has local commits that are not yet pushed to its upstream branch
on the remote. To skip this check, use -f.

If TITLE is omitted, a text editor will open in which title and body of
the pull request can be entered in the same manner as git commit message.

If instead of normal TITLE an issue number is given with -i, the pull
request will be attached to an existing GitHub issue. Alternatively, instead
of title you can paste a full URL to an issue on GitHub.
`,
}

var flagPullRequestBase, flagPullRequestHead, flagPullRequestIssue string

func init() {
	cmdPullRequest.Flag.StringVar(&flagPullRequestBase, "b", "master", "BASE")
	cmdPullRequest.Flag.StringVar(&flagPullRequestHead, "d", "", "HEAD")
	cmdPullRequest.Flag.StringVar(&flagPullRequestIssue, "i", "", "ISSUE")
}

func pullRequest(cmd *Command, args *Args) {
	var title, body string
	if args.Size() == 1 {
		title = args.First()
	}

	gh := github.New()
	repo := gh.Project.LocalRepoWith(flagPullRequestBase, flagPullRequestHead)
	if title == "" && flagPullRequestIssue == "" {
		messageFile, err := git.PullReqMsgFile()
		utils.Check(err)

		err = writePullRequestChanges(repo, messageFile)
		utils.Check(err)

		editorPath, err := git.EditorPath()
		utils.Check(err)

		err = editTitleAndBody(editorPath, messageFile)
		utils.Check(err)

		title, body, err = readTitleAndBody(messageFile)
		utils.Check(err)
	}

	if title == "" && flagPullRequestIssue == "" {
		log.Fatal("Aborting due to empty pull request title")
	}

	var pullRequestURL string
	var err error
	if title != "" {
		pullRequestURL, err = gh.CreatePullRequest(repo.Base, repo.Head, title, body)
	}
	if flagPullRequestIssue != "" {
		pullRequestURL, err = gh.CreatePullRequestForIssue(repo.Base, repo.Head, flagPullRequestIssue)
	}

	utils.Check(err)

	fmt.Println(pullRequestURL)

	os.Exit(0)
}

func writePullRequestChanges(repo *github.Repo, messageFile string) error {
	message := `
# Requesting a pull to %s from %s
#
# Write a message for this pull request. The first block
# of the text is the title and the rest is description.%s
`
	startRegexp := regexp.MustCompilePOSIX("^")
	endRegexp := regexp.MustCompilePOSIX(" +$")

	commitLogs, _ := git.Log(repo.Base, repo.Head)
	var changesMsg string
	if len(commitLogs) > 0 {
		commitLogs = strings.TrimSpace(commitLogs)
		commitLogs = startRegexp.ReplaceAllString(commitLogs, "# ")
		commitLogs = endRegexp.ReplaceAllString(commitLogs, "")
		changesMsg = `
#
# Changes:
#
%s`
		changesMsg = fmt.Sprintf(changesMsg, commitLogs)
	}

	message = fmt.Sprintf(message, repo.FullBase(), repo.FullHead(), changesMsg)

	return ioutil.WriteFile(messageFile, []byte(message), 0644)
}

func editTitleAndBody(editorPath, messageFile string) error {
	editCmd := cmd.New(editorPath)
	r := regexp.MustCompile("[mg]?vi[m]$")
	if r.MatchString(editorPath) {
		editCmd.WithArg("-c")
		editCmd.WithArg("set ft=gitcommit tw=0 wrap lbr")
	}
	editCmd.WithArg(messageFile)

	return editCmd.Exec()
}

func readTitleAndBody(messageFile string) (title, body string, err error) {
	f, err := os.Open(messageFile)
	defer f.Close()
	if err != nil {
		return "", "", err
	}

	reader := bufio.NewReader(f)

	return readTitleAndBodyFrom(reader)
}

func readTitleAndBodyFrom(reader *bufio.Reader) (title, body string, err error) {
	r := regexp.MustCompile("\\S")
	var titleParts, bodyParts []string

	line, err := readln(reader)
	for err == nil {
		if strings.HasPrefix(line, "#") {
			break
		}
		if len(bodyParts) == 0 && r.MatchString(line) {
			titleParts = append(titleParts, line)
		} else {
			bodyParts = append(bodyParts, line)
		}
		line, err = readln(reader)
	}

	title = strings.Join(titleParts, " ")
	title = strings.TrimSpace(title)

	body = strings.Join(bodyParts, "\n")
	body = strings.TrimSpace(body)

	return title, body, nil
}

func readln(r *bufio.Reader) (string, error) {
	var (
		isPrefix = true
		err      error
		line, ln []byte
	)
	for isPrefix && err == nil {
		line, isPrefix, err = r.ReadLine()
		ln = append(ln, line...)
	}

	return string(ln), err
}
<commit_msg>Remove messageFile after done<commit_after>package commands

import (
	"bufio"
	"fmt"
	"github.com/jingweno/gh/cmd"
	"github.com/jingweno/gh/git"
	"github.com/jingweno/gh/github"
	"github.com/jingweno/gh/utils"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"strings"
)

var cmdPullRequest = &Command{
	Run:   pullRequest,
	Usage: "pull-request [-f] [-i ISSUE] [-b BASE] [-d HEAD] [TITLE]",
	Short: "Open a pull request on GitHub",
	Long: `Opens a pull request on GitHub for the project that the "origin" remote
points to. The default head of the pull request is the current branch.
Both base and head of the pull request can be explicitly given in one of
the following formats: "branch", "owner:branch", "owner/repo:branch".
This command will abort operation if it detects that the current topic
branch has local commits that are not yet pushed to its upstream branch
on the remote. To skip this check, use -f.

If TITLE is omitted, a text editor will open in which title and body of
the pull request can be entered in the same manner as git commit message.

If instead of normal TITLE an issue number is given with -i, the pull
request will be attached to an existing GitHub issue. Alternatively, instead
of title you can paste a full URL to an issue on GitHub.
`,
}

var flagPullRequestBase, flagPullRequestHead, flagPullRequestIssue string

func init() {
	cmdPullRequest.Flag.StringVar(&flagPullRequestBase, "b", "master", "BASE")
	cmdPullRequest.Flag.StringVar(&flagPullRequestHead, "d", "", "HEAD")
	cmdPullRequest.Flag.StringVar(&flagPullRequestIssue, "i", "", "ISSUE")
}

func pullRequest(cmd *Command, args *Args) {
	var title, body string
	if args.Size() == 1 {
		title = args.First()
	}

	gh := github.New()
	repo := gh.Project.LocalRepoWith(flagPullRequestBase, flagPullRequestHead)
	if title == "" && flagPullRequestIssue == "" {
		messageFile, err := git.PullReqMsgFile()
		utils.Check(err)

		defer removeFile(messageFile)

		err = writePullRequestChanges(repo, messageFile)
		utils.Check(err)

		editorPath, err := git.EditorPath()
		utils.Check(err)

		err = editTitleAndBody(editorPath, messageFile)
		utils.Check(err)

		title, body, err = readTitleAndBody(messageFile)
		utils.Check(err)
	}

	if title == "" && flagPullRequestIssue == "" {
		log.Fatal("Aborting due to empty pull request title")
	}

	var pullRequestURL string
	var err error
	if title != "" {
		pullRequestURL, err = gh.CreatePullRequest(repo.Base, repo.Head, title, body)
	}
	if flagPullRequestIssue != "" {
		pullRequestURL, err = gh.CreatePullRequestForIssue(repo.Base, repo.Head, flagPullRequestIssue)
	}

	utils.Check(err)

	fmt.Println(pullRequestURL)

	os.Exit(0)
}

func writePullRequestChanges(repo *github.Repo, messageFile string) error {
	message := `
# Requesting a pull to %s from %s
#
# Write a message for this pull request. The first block
# of the text is the title and the rest is description.%s
`
	startRegexp := regexp.MustCompilePOSIX("^")
	endRegexp := regexp.MustCompilePOSIX(" +$")

	commitLogs, _ := git.Log(repo.Base, repo.Head)
	var changesMsg string
	if len(commitLogs) > 0 {
		commitLogs = strings.TrimSpace(commitLogs)
		commitLogs = startRegexp.ReplaceAllString(commitLogs, "# ")
		commitLogs = endRegexp.ReplaceAllString(commitLogs, "")
		changesMsg = `
#
# Changes:
#
%s`
		changesMsg = fmt.Sprintf(changesMsg, commitLogs)
	}

	message = fmt.Sprintf(message, repo.FullBase(), repo.FullHead(), changesMsg)

	return ioutil.WriteFile(messageFile, []byte(message), 0644)
}

func editTitleAndBody(editorPath, messageFile string) error {
	editCmd := cmd.New(editorPath)
	r := regexp.MustCompile("[mg]?vi[m]$")
	if r.MatchString(editorPath) {
		editCmd.WithArg("-c")
		editCmd.WithArg("set ft=gitcommit tw=0 wrap lbr")
	}
	editCmd.WithArg(messageFile)

	return editCmd.Exec()
}

func readTitleAndBody(messageFile string) (title, body string, err error) {
	f, err := os.Open(messageFile)
	defer f.Close()
	if err != nil {
		return "", "", err
	}

	reader := bufio.NewReader(f)

	return readTitleAndBodyFrom(reader)
}

func readTitleAndBodyFrom(reader *bufio.Reader) (title, body string, err error) {
	r := regexp.MustCompile("\\S")
	var titleParts, bodyParts []string

	line, err := readln(reader)
	for err == nil {
		if strings.HasPrefix(line, "#") {
			break
		}
		if len(bodyParts) == 0 && r.MatchString(line) {
			titleParts = append(titleParts, line)
		} else {
			bodyParts = append(bodyParts, line)
		}
		line, err = readln(reader)
	}

	title = strings.Join(titleParts, " ")
	title = strings.TrimSpace(title)

	body = strings.Join(bodyParts, "\n")
	body = strings.TrimSpace(body)

	return title, body, nil
}

func readln(r *bufio.Reader) (string, error) {
	var (
		isPrefix = true
		err      error
		line, ln []byte
	)
	for isPrefix && err == nil {
		line, isPrefix, err = r.ReadLine()
		ln = append(ln, line...)
	}

	return string(ln), err
}

func removeFile(file string) {
	err := os.Remove(file)
	utils.Check(err)
}
<|endoftext|>