<commit_before>package main

import (
	"time"

	"labix.org/v2/mgo/bson"
)

// Sane defaults to prevent spamming the system.
var (
	EmailLimitPerRun    = 5000
	DeleteVMLimitPerRun = 5000
)

var VMDeletionWarning1 = &Warning{
	ID: "vmDeletionWarning-1",

	Description: "Find users inactive for > 20 days, send email",

	Select: []bson.M{
		bson.M{"lastLoginDate": moreThanDaysQuery(20)},
		bson.M{"inactive.warning": bson.M{"$exists": false}},
	},

	ExemptCheckers: []*ExemptChecker{
		IsUserPaid, IsUserNotConfirmed, IsUserVMsEmpty, IsUserKodingEmployee,
	},

	Action: SendEmail,
}

var VMDeletionWarning2 = &Warning{
	ID: "vmDeletionWarning-2",

	Description: "Find users inactive for > 24 days, send email",

	PreviousWarning: VMDeletionWarning1,

	IntervalSinceLastWarning: time.Hour * 24 * 4, // 4 days since last warning

	Select: []bson.M{
		bson.M{"lastLoginDate": moreThanDaysQuery(24)},
		bson.M{"inactive.warning": VMDeletionWarning1.ID},
	},

	ExemptCheckers: []*ExemptChecker{
		IsTooSoon, IsUserPaid, IsUserNotConfirmed, IsUserVMsEmpty, IsUserKodingEmployee,
	},

	Action: SendEmail,
}

var DeleteInactiveUserVM = &Warning{
	ID: "deleteInactiveUserVm",

	Description: "Find users inactive for > 29 days, deleted ALL their vms",

	PreviousWarning: VMDeletionWarning2,

	IntervalSinceLastWarning: time.Hour * 24 * 4, // 4 days since last warning

	Select: []bson.M{
		bson.M{"lastLoginDate": moreThanDaysQuery(29)},
		bson.M{"inactive.warning": VMDeletionWarning2.ID},
	},

	ExemptCheckers: []*ExemptChecker{
		IsTooSoon, IsUserPaid, IsUserVMsEmpty, IsUserKodingEmployee,
	},

	Action: DeleteVMs,
}

var DeleteBlockedUserVM = &Warning{
	ID: "deleteBlockedUserVm",

	Description: "Find blocked users inactive > 14 days, delete ALL their vms",

	Select: []bson.M{
		bson.M{"lastLoginDate": moreThanDaysQuery(14)},
		bson.M{"inactive.warning": bson.M{"$exists": false}},
		bson.M{"status": "blocked"},
	},

	ExemptCheckers: []*ExemptChecker{IsUserVMsEmpty},

	Action: DeleteVMs,
}
<commit_msg>janitor: remove unused code<commit_after>package main

import (
	"time"

	"labix.org/v2/mgo/bson"
)

var VMDeletionWarning1 = &Warning{
	ID: "vmDeletionWarning-1",

	Description: "Find users inactive for > 20 days, send email",

	Select: []bson.M{
		bson.M{"lastLoginDate": moreThanDaysQuery(20)},
		bson.M{"inactive.warning": bson.M{"$exists": false}},
	},

	ExemptCheckers: []*ExemptChecker{
		IsUserPaid, IsUserNotConfirmed, IsUserVMsEmpty, IsUserKodingEmployee,
	},

	Action: SendEmail,
}

var VMDeletionWarning2 = &Warning{
	ID: "vmDeletionWarning-2",

	Description: "Find users inactive for > 24 days, send email",

	PreviousWarning: VMDeletionWarning1,

	IntervalSinceLastWarning: time.Hour * 24 * 4, // 4 days since last warning

	Select: []bson.M{
		bson.M{"lastLoginDate": moreThanDaysQuery(24)},
		bson.M{"inactive.warning": VMDeletionWarning1.ID},
	},

	ExemptCheckers: []*ExemptChecker{
		IsTooSoon, IsUserPaid, IsUserNotConfirmed, IsUserVMsEmpty, IsUserKodingEmployee,
	},

	Action: SendEmail,
}

var DeleteInactiveUserVM = &Warning{
	ID: "deleteInactiveUserVm",

	Description: "Find users inactive for > 29 days, deleted ALL their vms",

	PreviousWarning: VMDeletionWarning2,

	IntervalSinceLastWarning: time.Hour * 24 * 4, // 4 days since last warning

	Select: []bson.M{
		bson.M{"lastLoginDate": moreThanDaysQuery(29)},
		bson.M{"inactive.warning": VMDeletionWarning2.ID},
	},

	ExemptCheckers: []*ExemptChecker{
		IsTooSoon, IsUserPaid, IsUserVMsEmpty, IsUserKodingEmployee,
	},

	Action: DeleteVMs,
}

var DeleteBlockedUserVM = &Warning{
	ID: "deleteBlockedUserVm",

	Description: "Find blocked users inactive > 14 days, delete ALL their vms",

	Select: []bson.M{
		bson.M{"lastLoginDate": moreThanDaysQuery(14)},
		bson.M{"inactive.warning": bson.M{"$exists": false}},
		bson.M{"status": "blocked"},
	},

	ExemptCheckers: []*ExemptChecker{IsUserVMsEmpty},

	Action: DeleteVMs,
}
<|endoftext|>


<commit_before>package models

import (
	"errors"
	"socialapi/db"
	"time"
)

type ChannelMessage struct {
	// unique identifier of the channel message
	Id int64

	// Body of the mesage
	Body string

	// type of the message
	Type int

	// Creator of the channel message
	CreatorId int64

	// Creation date of the message
	CreatedAt time.Time

	// Modification date of the message
	ModifiedAt time.Time
}

const (
	POST int = iota
	JOIN
	LEAVE
	CHAT
)

func NewChannelMessage() *ChannelMessage {
	now := time.Now().UTC()
	return &ChannelMessage{
		Type:       POST,
		CreatedAt:  now,
		ModifiedAt: now,
	}
}

func (c *ChannelMessage) Fetch() error {
	if err := db.DB.First(c, c.Id).Error; err != nil {
		return err
	}
	return nil
}

func (c *ChannelMessage) Update() error {
	if c.Id == 0 {
		return errors.New("Channel message id is not set")
	}
	return c.Save()
}

func (c *ChannelMessage) Save() error {

	now := time.Now().UTC()
	// created at shouldnt be updated
	c.CreatedAt = now
	c.ModifiedAt = now

	if err := db.DB.Save(c).Error; err != nil {
		return err
	}
	return nil
}

func (c *ChannelMessage) Delete() error {
	if c.Id == 0 {
		return errors.New("Channel message id is not set")
	}

	if err := db.DB.Delete(c).Error; err != nil {
		return err
	}

	return nil
}
<commit_msg>Social: use helper functions for model operations<commit_after>package models

import (
	"errors"
	"socialapi/db"
	"time"
)

type ChannelMessage struct {
	// unique identifier of the channel message
	Id int64

	// Body of the mesage
	Body string

	// type of the message
	Type int

	// Creator of the channel message
	AccountId int64

	// Creation date of the message
	CreatedAt time.Time

	// Modification date of the message
	UpdatedAt time.Time
}

const (
	POST int = iota
	JOIN
	LEAVE
	CHAT
)

func NewChannelMessage() *ChannelMessage {
	now := time.Now().UTC()
	return &ChannelMessage{
		Type:      POST,
		CreatedAt: now,
		UpdatedAt: now,
	}
}

func (c ChannelMessage) TableName() string {
	return "channel_message"
}

func (c *ChannelMessage) Fetch() error {
	if c.Id == 0 {
		return errors.New("Channel message id is not set")
	}
	if err := First(c, c.Id); err != nil {
		return err
	}
	return nil
}

func (c *ChannelMessage) Update() error {
	if c.Id == 0 {
		return errors.New("Channel message id is not set")
	}

	if err := db.DB.
		Table(c.TableName()).
		Where(c.Id).
		Update("body", c.Body).
		Error; err != nil {
		return err
	}

	return nil
}

func (c *ChannelMessage) Save() error {

	now := time.Now().UTC()
	// created at shouldnt be updated
	c.CreatedAt = now
	c.UpdatedAt = now

	if err := Save(c); err != nil {
		return err
	}
	return nil
}

func (c *ChannelMessage) Delete() error {
	if c.Id == 0 {
		return errors.New("Channel message id is not set")
	}

	if err := Delete(c); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package google

import (
	"fmt"
	"log"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"

	"google.golang.org/api/compute/v1"
)

func resourceComputeGlobalAddress() *schema.Resource {
	return &schema.Resource{
		Create: resourceComputeGlobalAddressCreate,
		Read:   resourceComputeGlobalAddressRead,
		Delete: resourceComputeGlobalAddressDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"ip_version": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				ForceNew:     true,
				ValidateFunc: validation.StringInSlice([]string{"IPV4", "IPV6"}, false),
			},

			"project": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"address": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"self_link": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceComputeGlobalAddressCreate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	project, err := getProject(d, config)
	if err != nil {
		return err
	}

	// Build the address parameter
	addr := &compute.Address{
		Name:      d.Get("name").(string),
		IpVersion: d.Get("ip_version").(string),
	}

	op, err := config.clientCompute.GlobalAddresses.Insert(project, addr).Do()
	if err != nil {
		return fmt.Errorf("Error creating address: %s", err)
	}

	// It probably maybe worked, so store the ID now
	d.SetId(addr.Name)

	err = computeSharedOperationWait(config.clientCompute, op, project, "Creating Global Address")
	if err != nil {
		return err
	}

	return resourceComputeGlobalAddressRead(d, meta)
}

func resourceComputeGlobalAddressRead(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	project, err := getProject(d, config)
	if err != nil {
		return err
	}

	addr, err := config.clientCompute.GlobalAddresses.Get(project, d.Id()).Do()
	if err != nil {
		return handleNotFoundError(err, d, fmt.Sprintf("Global Address %q", d.Get("name").(string)))
	}

	d.Set("name", addr.Name)
	d.Set("ip_version", addr.IpVersion)
	d.Set("address", addr.Address)
	d.Set("project", project)
	d.Set("self_link", ConvertSelfLinkToV1(addr.SelfLink))

	return nil
}

func resourceComputeGlobalAddressDelete(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	project, err := getProject(d, config)
	if err != nil {
		return err
	}

	// Delete the address
	log.Printf("[DEBUG] address delete request")
	op, err := config.clientCompute.GlobalAddresses.Delete(project, d.Id()).Do()
	if err != nil {
		return fmt.Errorf("Error deleting address: %s", err)
	}

	err = computeSharedOperationWait(config.clientCompute, op, project, "Deleting Global Address")
	if err != nil {
		return err
	}

	d.SetId("")
	return nil
}
<commit_msg>add empty value for ip version for global address (#799)<commit_after>package google

import (
	"fmt"
	"log"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"

	"google.golang.org/api/compute/v1"
)

func resourceComputeGlobalAddress() *schema.Resource {
	return &schema.Resource{
		Create: resourceComputeGlobalAddressCreate,
		Read:   resourceComputeGlobalAddressRead,
		Delete: resourceComputeGlobalAddressDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"ip_version": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				ForceNew:     true,
				ValidateFunc: validation.StringInSlice([]string{"IPV4", "IPV6", ""}, false),
			},

			"project": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"address": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"self_link": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceComputeGlobalAddressCreate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	project, err := getProject(d, config)
	if err != nil {
		return err
	}

	// Build the address parameter
	addr := &compute.Address{
		Name:      d.Get("name").(string),
		IpVersion: d.Get("ip_version").(string),
	}

	op, err := config.clientCompute.GlobalAddresses.Insert(project, addr).Do()
	if err != nil {
		return fmt.Errorf("Error creating address: %s", err)
	}

	// It probably maybe worked, so store the ID now
	d.SetId(addr.Name)

	err = computeSharedOperationWait(config.clientCompute, op, project, "Creating Global Address")
	if err != nil {
		return err
	}

	return resourceComputeGlobalAddressRead(d, meta)
}

func resourceComputeGlobalAddressRead(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	project, err := getProject(d, config)
	if err != nil {
		return err
	}

	addr, err := config.clientCompute.GlobalAddresses.Get(project, d.Id()).Do()
	if err != nil {
		return handleNotFoundError(err, d, fmt.Sprintf("Global Address %q", d.Get("name").(string)))
	}

	d.Set("name", addr.Name)
	d.Set("ip_version", addr.IpVersion)
	d.Set("address", addr.Address)
	d.Set("project", project)
	d.Set("self_link", ConvertSelfLinkToV1(addr.SelfLink))

	return nil
}

func resourceComputeGlobalAddressDelete(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	project, err := getProject(d, config)
	if err != nil {
		return err
	}

	// Delete the address
	log.Printf("[DEBUG] address delete request")
	op, err := config.clientCompute.GlobalAddresses.Delete(project, d.Id()).Do()
	if err != nil {
		return fmt.Errorf("Error deleting address: %s", err)
	}

	err = computeSharedOperationWait(config.clientCompute, op, project, "Deleting Global Address")
	if err != nil {
		return err
	}

	d.SetId("")
	return nil
}
<|endoftext|>


<commit_before>package aggregator

import (
	"sync"
	"time"

	"github.com/percona/go-mysql/event"
	"github.com/percona/percona-toolkit/src/go/mongolib/fingerprinter"
	"github.com/percona/percona-toolkit/src/go/mongolib/proto"
	"github.com/percona/percona-toolkit/src/go/mongolib/stats"
	pc "github.com/percona/pmm/proto/config"
	"github.com/percona/pmm/proto/qan"
	"github.com/percona/qan-agent/qan/analyzer/report"
)

func New(timeStart time.Time, config pc.QAN) *Aggregator {
	fp := fingerprinter.NewFingerprinter(fingerprinter.DEFAULT_KEY_FILTERS)
	s := stats.New(fp)

	d := time.Duration(config.Interval) * time.Second

	// truncate to the interval e.g 12:15:35 with 1 minute interval it will be 12:15:00
	timeStart = timeStart.UTC().Truncate(d)
	// create ending time by adding interval
	timeEnd := timeStart.Add(d)

	return &Aggregator{
		// dependencies
		timeStart: timeStart,
		config:    config,

		// internal
		timeEnd: timeEnd,
		d:       d,
		stats:   s,
	}
}

type Aggregator struct {
	// dependencies
	config pc.QAN

	// interval
	timeStart time.Time
	timeEnd   time.Time
	d         time.Duration

	stats *stats.Stats

	// state
	sync.RWMutex // Lock() to protect internal consistency of the service
}

func (self *Aggregator) Add(doc proto.SystemProfile) (*qan.Report, error) {
	ts := doc.Ts.UTC()

	// skip old metrics
	if ts.Before(self.timeStart) {
		return nil, nil
	}

	return self.NextReport(ts), self.stats.Add(doc)
}

func (self *Aggregator) NextReport(ts time.Time) *qan.Report {
	// if time is greater than interval then we are done with this interval
	if !ts.Before(self.timeEnd) {
		// reset stats
		defer self.resetStats(ts)

		// let's check if we have anything to send
		if len(self.stats.Queries()) > 0 {
			// create result
			result := self.createResult()

			// translate result into report
			return report.MakeReport(self.config, self.timeStart, self.timeEnd, nil, result)
		}
	}

	// we are not done with the interval so no report yet
	return nil
}

func (self *Aggregator) TimeStart() time.Time {
	self.RLock()
	defer self.RUnlock()
	return self.timeStart
}

func (self *Aggregator) TimeEnd() time.Time {
	self.RLock()
	defer self.RUnlock()
	return self.timeEnd
}

func (self *Aggregator) resetStats(ts time.Time) {
	// reset stats
	self.stats.Reset()

	// update time intervals
	self.timeStart = ts.Truncate(self.d)
	self.timeEnd = self.timeStart.Add(self.d)
}

func (self *Aggregator) createResult() *report.Result {
	queries := self.stats.Queries()
	global := event.NewClass("", "", false)
	queryStats := queries.CalcQueriesStats(int64(self.config.Interval))
	classes := []*event.Class{}
	for _, queryInfo := range queryStats {
		class := event.NewClass(queryInfo.ID, queryInfo.Fingerprint, self.config.ExampleQueries)
		if self.config.ExampleQueries {
			class.Example = &event.Example{
				QueryTime: queryInfo.QueryTime.Total,
				Db:        queryInfo.Namespace,
				Query:     queryInfo.Query,
			}
		}

		metrics := event.NewMetrics()

		metrics.TimeMetrics["Query_time"] = newEventTimeStats(queryInfo.QueryTime)

		// @todo we map below metrics to MySQL equivalents according to PMM-830
		metrics.NumberMetrics["Bytes_sent"] = newEventNumberStats(queryInfo.ResponseLength)
		metrics.NumberMetrics["Rows_sent"] = newEventNumberStats(queryInfo.Returned)
		metrics.NumberMetrics["Rows_examined"] = newEventNumberStats(queryInfo.Scanned)

		class.Metrics = metrics
		class.TotalQueries = uint(queryInfo.Count)
		class.UniqueQueries = 1
		classes = append(classes, class)

		// Add the class to the global metrics.
		global.AddClass(class)
	}

	return &report.Result{
		Global: global,
		Class:  classes,
	}

}

func newEventNumberStats(s stats.Statistics) *event.NumberStats {
	return &event.NumberStats{
		Sum: uint64(s.Total),
		Min: uint64(s.Min),
		Avg: uint64(s.Avg),
		Med: uint64(s.Median),
		P95: uint64(s.Pct95),
		Max: uint64(s.Max),
	}
}

func newEventTimeStats(s stats.Statistics) *event.TimeStats {
	return &event.TimeStats{
		Sum: s.Total,
		Min: s.Min,
		Avg: s.Avg,
		Med: s.Median,
		P95: s.Pct95,
		Max: s.Max,
	}
}
<commit_msg>cleanups<commit_after>package aggregator

import (
	"time"

	"github.com/percona/go-mysql/event"
	"github.com/percona/percona-toolkit/src/go/mongolib/fingerprinter"
	"github.com/percona/percona-toolkit/src/go/mongolib/proto"
	"github.com/percona/percona-toolkit/src/go/mongolib/stats"
	pc "github.com/percona/pmm/proto/config"
	"github.com/percona/pmm/proto/qan"
	"github.com/percona/qan-agent/qan/analyzer/report"
)

// New returns configured *Aggregator
func New(timeStart time.Time, config pc.QAN) *Aggregator {
	aggregator := &Aggregator{
		config: config,
	}

	// create duration from interval
	aggregator.d = time.Duration(config.Interval) * time.Second

	// create mongolib stats
	fp := fingerprinter.NewFingerprinter(fingerprinter.DEFAULT_KEY_FILTERS)
	aggregator.stats = stats.New(fp)

	// create new interval
	aggregator.newInterval(timeStart)

	return aggregator
}

// Aggregator aggregates system.profile document
type Aggregator struct {
	// dependencies
	config pc.QAN

	// interval
	timeStart time.Time
	timeEnd   time.Time
	d         time.Duration
	stats     *stats.Stats
}

// Add aggregates new system.profile document and returns report if it's ready
func (self *Aggregator) Add(doc proto.SystemProfile) (*qan.Report, error) {
	ts := doc.Ts.UTC()

	// skip old metrics
	if ts.Before(self.timeStart) {
		return nil, nil
	}

	return self.Interval(ts), self.stats.Add(doc)
}

// Interval sets interval if necessary and returns *qan.Report for old interval if not empty
func (self *Aggregator) Interval(ts time.Time) *qan.Report {
	// if time is before interval end then we are still in the same interval, nothing to do
	if ts.Before(self.timeEnd) {
		return nil
	}

	// create new interval
	defer self.newInterval(ts)

	// let's check if we have anything to send for current interval
	if len(self.stats.Queries()) == 0 {
		// if there are no queries then we don't create report #PMM-927
		return nil
	}

	// create result
	result := self.createResult()

	// translate result into report and return it
	return report.MakeReport(self.config, self.timeStart, self.timeEnd, nil, result)
}

// TimeStart returns start time for current interval
func (self *Aggregator) TimeStart() time.Time {
	return self.timeStart
}

// TimeEnd returns end time for current interval
func (self *Aggregator) TimeEnd() time.Time {
	return self.timeEnd
}

func (self *Aggregator) newInterval(ts time.Time) {
	// reset stats
	self.stats.Reset()

	// truncate to the duration e.g 12:15:35 with 1 minute duration it will be 12:15:00
	self.timeStart = ts.UTC().Truncate(self.d)
	// create ending time by adding interval
	self.timeEnd = self.timeStart.Add(self.d)
}

func (self *Aggregator) createResult() *report.Result {
	queries := self.stats.Queries()
	global := event.NewClass("", "", false)
	queryStats := queries.CalcQueriesStats(int64(self.config.Interval))
	classes := []*event.Class{}
	for _, queryInfo := range queryStats {
		class := event.NewClass(queryInfo.ID, queryInfo.Fingerprint, self.config.ExampleQueries)
		if self.config.ExampleQueries {
			class.Example = &event.Example{
				QueryTime: queryInfo.QueryTime.Total,
				Db:        queryInfo.Namespace,
				Query:     queryInfo.Query,
			}
		}

		metrics := event.NewMetrics()

		metrics.TimeMetrics["Query_time"] = newEventTimeStats(queryInfo.QueryTime)

		// @todo we map below metrics to MySQL equivalents according to PMM-830
		metrics.NumberMetrics["Bytes_sent"] = newEventNumberStats(queryInfo.ResponseLength)
		metrics.NumberMetrics["Rows_sent"] = newEventNumberStats(queryInfo.Returned)
		metrics.NumberMetrics["Rows_examined"] = newEventNumberStats(queryInfo.Scanned)

		class.Metrics = metrics
		class.TotalQueries = uint(queryInfo.Count)
		class.UniqueQueries = 1
		classes = append(classes, class)

		// Add the class to the global metrics.
		global.AddClass(class)
	}

	return &report.Result{
		Global: global,
		Class:  classes,
	}

}

func newEventNumberStats(s stats.Statistics) *event.NumberStats {
	return &event.NumberStats{
		Sum: uint64(s.Total),
		Min: uint64(s.Min),
		Avg: uint64(s.Avg),
		Med: uint64(s.Median),
		P95: uint64(s.Pct95),
		Max: uint64(s.Max),
	}
}

func newEventTimeStats(s stats.Statistics) *event.TimeStats {
	return &event.TimeStats{
		Sum: s.Total,
		Min: s.Min,
		Avg: s.Avg,
		Med: s.Median,
		P95: s.Pct95,
		Max: s.Max,
	}
}
<|endoftext|>


<commit_before>package kubernetes

import (
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
	api "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	pkgApi "k8s.io/apimachinery/pkg/types"
	kubernetes "k8s.io/client-go/kubernetes"
)

func resourceKubernetesService() *schema.Resource {
	return &schema.Resource{
		Create: resourceKubernetesServiceCreate,
		Read:   resourceKubernetesServiceRead,
		Exists: resourceKubernetesServiceExists,
		Update: resourceKubernetesServiceUpdate,
		Delete: resourceKubernetesServiceDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"metadata": namespacedMetadataSchema("service", true),
			"spec": {
				Type:        schema.TypeList,
				Description: "Spec defines the behavior of a service. https://github.com/kubernetes/community/blob/master/contributors/devel/api-conventions.md#spec-and-status",
				Required:    true,
				MaxItems:    1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"cluster_ip": {
							Type:        schema.TypeString,
							Description: "The IP address of the service. It is usually assigned randomly by the master. If an address is specified manually and is not in use by others, it will be allocated to the service; otherwise, creation of the service will fail. `None` can be specified for headless services when proxying is not required. Ignored if type is `ExternalName`. More info: http://kubernetes.io/docs/user-guide/services#virtual-ips-and-service-proxies",
							Optional:    true,
							ForceNew:    true,
							Computed:    true,
						},
						"external_ips": {
							Type:        schema.TypeSet,
							Description: "A list of IP addresses for which nodes in the cluster will also accept traffic for this service. These IPs are not managed by Kubernetes. The user is responsible for ensuring that traffic arrives at a node with this IP.  A common example is external load-balancers that are not part of the Kubernetes system.",
							Optional:    true,
							Elem:        &schema.Schema{Type: schema.TypeString},
							Set:         schema.HashString,
						},
						"external_name": {
							Type:        schema.TypeString,
							Description: "The external reference that kubedns or equivalent will return as a CNAME record for this service. No proxying will be involved. Must be a valid DNS name and requires `type` to be `ExternalName`.",
							Optional:    true,
						},
						"external_traffic_policy": {
							Type:         schema.TypeString,
							Description:  "Denotes if this Service desires to route external traffic to node-local or cluster-wide endpoints. `Local` preserves the client source IP and avoids a second hop for LoadBalancer and Nodeport type services, but risks potentially imbalanced traffic spreading. `Cluster` obscures the client source IP and may cause a second hop to another node, but should have good overall load-spreading. More info: https://kubernetes.io/docs/tutorials/services/source-ip/",
							Optional:     true,
							ValidateFunc: validation.StringInSlice([]string{"Local", "Cluster"}, false),
						},
						"load_balancer_ip": {
							Type:        schema.TypeString,
							Description: "Only applies to `type = LoadBalancer`. LoadBalancer will get created with the IP specified in this field. This feature depends on whether the underlying cloud-provider supports specifying this field when a load balancer is created. This field will be ignored if the cloud-provider does not support the feature.",
							Optional:    true,
						},
						"load_balancer_source_ranges": {
							Type:        schema.TypeSet,
							Description: "If specified and supported by the platform, this will restrict traffic through the cloud-provider load-balancer will be restricted to the specified client IPs. This field will be ignored if the cloud-provider does not support the feature. More info: http://kubernetes.io/docs/user-guide/services-firewalls",
							Optional:    true,
							Elem:        &schema.Schema{Type: schema.TypeString},
							Set:         schema.HashString,
						},
						"port": {
							Type:        schema.TypeList,
							Description: "The list of ports that are exposed by this service. More info: http://kubernetes.io/docs/user-guide/services#virtual-ips-and-service-proxies",
							Optional:    true,
							MinItems:    1,
							Elem: &schema.Resource{
								Schema: map[string]*schema.Schema{
									"name": {
										Type:        schema.TypeString,
										Description: "The name of this port within the service. All ports within the service must have unique names. Optional if only one ServicePort is defined on this service.",
										Optional:    true,
									},
									"node_port": {
										Type:        schema.TypeInt,
										Description: "The port on each node on which this service is exposed when `type` is `NodePort` or `LoadBalancer`. Usually assigned by the system. If specified, it will be allocated to the service if unused or else creation of the service will fail. Default is to auto-allocate a port if the `type` of this service requires one. More info: http://kubernetes.io/docs/user-guide/services#type--nodeport",
										Computed:    true,
										Optional:    true,
									},
									"port": {
										Type:        schema.TypeInt,
										Description: "The port that will be exposed by this service.",
										Required:    true,
									},
									"protocol": {
										Type:        schema.TypeString,
										Description: "The IP protocol for this port. Supports `TCP` and `UDP`. Default is `TCP`.",
										Optional:    true,
										Default:     "TCP",
									},
									"target_port": {
										Type:        schema.TypeString,
										Description: "Number or name of the port to access on the pods targeted by the service. Number must be in the range 1 to 65535. This field is ignored for services with `cluster_ip = \"None\"`. More info: http://kubernetes.io/docs/user-guide/services#defining-a-service",
										Optional:    true,
										Computed:    true,
									},
								},
							},
						},
						"publish_not_ready_addresses": {
							Type:        schema.TypeBool,
							Optional:    true,
							Default:     false,
							Description: "When set to true, indicates that DNS implementations must publish the `notReadyAddresses` of subsets for the Endpoints associated with the Service. The default value is `false`. The primary use case for setting this field is to use a StatefulSet's Headless Service to propagate `SRV` records for its Pods without respect to their readiness for purpose of peer discovery.",
						},
						"selector": {
							Type:        schema.TypeMap,
							Description: "Route service traffic to pods with label keys and values matching this selector. Only applies to types `ClusterIP`, `NodePort`, and `LoadBalancer`. More info: http://kubernetes.io/docs/user-guide/services#overview",
							Optional:    true,
						},
						"session_affinity": {
							Type:        schema.TypeString,
							Description: "Used to maintain session affinity. Supports `ClientIP` and `None`. Defaults to `None`. More info: http://kubernetes.io/docs/user-guide/services#virtual-ips-and-service-proxies",
							Optional:    true,
							Default:     "None",
						},
						"type": {
							Type:        schema.TypeString,
							Description: "Determines how the service is exposed. Defaults to `ClusterIP`. Valid options are `ExternalName`, `ClusterIP`, `NodePort`, and `LoadBalancer`. `ExternalName` maps to the specified `external_name`. More info: http://kubernetes.io/docs/user-guide/services#overview",
							Optional:    true,
							Default:     "ClusterIP",
						},
					},
				},
			},
			"load_balancer_ingress": {
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"ip": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"hostname": {
							Type:     schema.TypeString,
							Computed: true,
						},
					},
				},
			},
		},
	}
}

func resourceKubernetesServiceCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*kubernetes.Clientset)

	metadata := expandMetadata(d.Get("metadata").([]interface{}))
	svc := api.Service{
		ObjectMeta: metadata,
		Spec:       expandServiceSpec(d.Get("spec").([]interface{})),
	}
	log.Printf("[INFO] Creating new service: %#v", svc)
	out, err := conn.CoreV1().Services(metadata.Namespace).Create(&svc)
	if err != nil {
		return err
	}
	log.Printf("[INFO] Submitted new service: %#v", out)
	d.SetId(buildId(out.ObjectMeta))

	if out.Spec.Type == api.ServiceTypeLoadBalancer {
		log.Printf("[DEBUG] Waiting for load balancer to assign IP/hostname")

		err = resource.Retry(10*time.Minute, func() *resource.RetryError {
			svc, err := conn.CoreV1().Services(out.Namespace).Get(out.Name, meta_v1.GetOptions{})
			if err != nil {
				log.Printf("[DEBUG] Received error: %#v", err)
				return resource.NonRetryableError(err)
			}

			lbIngress := svc.Status.LoadBalancer.Ingress

			log.Printf("[INFO] Received service status: %#v", svc.Status)
			if len(lbIngress) > 0 {
				return nil
			}

			return resource.RetryableError(fmt.Errorf(
				"Waiting for service %q to assign IP/hostname for a load balancer", d.Id()))
		})
		if err != nil {
			lastWarnings, wErr := getLastWarningsForObject(conn, out.ObjectMeta, "Service", 3)
			if wErr != nil {
				return wErr
			}
			return fmt.Errorf("%s%s", err, stringifyEvents(lastWarnings))
		}
	}

	return resourceKubernetesServiceRead(d, meta)
}

func resourceKubernetesServiceRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*kubernetes.Clientset)

	namespace, name, err := idParts(d.Id())
	if err != nil {
		return err
	}

	log.Printf("[INFO] Reading service %s", name)
	svc, err := conn.CoreV1().Services(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		log.Printf("[DEBUG] Received error: %#v", err)
		return err
	}
	log.Printf("[INFO] Received service: %#v", svc)
	err = d.Set("metadata", flattenMetadata(svc.ObjectMeta, d))
	if err != nil {
		return err
	}

	err = d.Set("load_balancer_ingress", flattenLoadBalancerIngress(svc.Status.LoadBalancer.Ingress))
	if err != nil {
		return err
	}

	flattened := flattenServiceSpec(svc.Spec)
	log.Printf("[DEBUG] Flattened service spec: %#v", flattened)
	err = d.Set("spec", flattened)
	if err != nil {
		return err
	}

	return nil
}

func resourceKubernetesServiceUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*kubernetes.Clientset)

	namespace, name, err := idParts(d.Id())
	if err != nil {
		return err
	}

	ops := patchMetadata("metadata.0.", "/metadata/", d)
	if d.HasChange("spec") {
		serverVersion, err := conn.ServerVersion()
		if err != nil {
			return err
		}
		diffOps, err := patchServiceSpec("spec.0.", "/spec/", d, serverVersion)
		if err != nil {
			return err
		}
		ops = append(ops, diffOps...)
	}
	data, err := ops.MarshalJSON()
	if err != nil {
		return fmt.Errorf("Failed to marshal update operations: %s", err)
	}
	log.Printf("[INFO] Updating service %q: %v", name, string(data))
	out, err := conn.CoreV1().Services(namespace).Patch(name, pkgApi.JSONPatchType, data)
	if err != nil {
		return fmt.Errorf("Failed to update service: %s", err)
	}
	log.Printf("[INFO] Submitted updated service: %#v", out)
	d.SetId(buildId(out.ObjectMeta))

	return resourceKubernetesServiceRead(d, meta)
}

func resourceKubernetesServiceDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*kubernetes.Clientset)

	namespace, name, err := idParts(d.Id())
	if err != nil {
		return err
	}

	log.Printf("[INFO] Deleting service: %#v", name)
	err = conn.CoreV1().Services(namespace).Delete(name, &meta_v1.DeleteOptions{})
	if err != nil {
		return err
	}

	log.Printf("[INFO] Service %s deleted", name)

	d.SetId("")
	return nil
}

func resourceKubernetesServiceExists(d *schema.ResourceData, meta interface{}) (bool, error) {
	conn := meta.(*kubernetes.Clientset)

	namespace, name, err := idParts(d.Id())
	if err != nil {
		return false, err
	}

	log.Printf("[INFO] Checking service %s", name)
	_, err = conn.CoreV1().Services(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		if statusErr, ok := err.(*errors.StatusError); ok && statusErr.ErrStatus.Code == 404 {
			return false, nil
		}
		log.Printf("[DEBUG] Received error: %#v", err)
	}
	return true, err
}
<commit_msg>Set default for 'external_traffic_policy'<commit_after>package kubernetes

import (
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
	api "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	pkgApi "k8s.io/apimachinery/pkg/types"
	kubernetes "k8s.io/client-go/kubernetes"
)

func resourceKubernetesService() *schema.Resource {
	return &schema.Resource{
		Create: resourceKubernetesServiceCreate,
		Read:   resourceKubernetesServiceRead,
		Exists: resourceKubernetesServiceExists,
		Update: resourceKubernetesServiceUpdate,
		Delete: resourceKubernetesServiceDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"metadata": namespacedMetadataSchema("service", true),
			"spec": {
				Type:        schema.TypeList,
				Description: "Spec defines the behavior of a service. https://github.com/kubernetes/community/blob/master/contributors/devel/api-conventions.md#spec-and-status",
				Required:    true,
				MaxItems:    1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"cluster_ip": {
							Type:        schema.TypeString,
							Description: "The IP address of the service. It is usually assigned randomly by the master. If an address is specified manually and is not in use by others, it will be allocated to the service; otherwise, creation of the service will fail. `None` can be specified for headless services when proxying is not required. Ignored if type is `ExternalName`. More info: http://kubernetes.io/docs/user-guide/services#virtual-ips-and-service-proxies",
							Optional:    true,
							ForceNew:    true,
							Computed:    true,
						},
						"external_ips": {
							Type:        schema.TypeSet,
							Description: "A list of IP addresses for which nodes in the cluster will also accept traffic for this service. These IPs are not managed by Kubernetes. The user is responsible for ensuring that traffic arrives at a node with this IP.  A common example is external load-balancers that are not part of the Kubernetes system.",
							Optional:    true,
							Elem:        &schema.Schema{Type: schema.TypeString},
							Set:         schema.HashString,
						},
						"external_name": {
							Type:        schema.TypeString,
							Description: "The external reference that kubedns or equivalent will return as a CNAME record for this service. No proxying will be involved. Must be a valid DNS name and requires `type` to be `ExternalName`.",
							Optional:    true,
						},
						"external_traffic_policy": {
							Type:         schema.TypeString,
							Description:  "Denotes if this Service desires to route external traffic to node-local or cluster-wide endpoints. `Local` preserves the client source IP and avoids a second hop for LoadBalancer and Nodeport type services, but risks potentially imbalanced traffic spreading. `Cluster` obscures the client source IP and may cause a second hop to another node, but should have good overall load-spreading. More info: https://kubernetes.io/docs/tutorials/services/source-ip/",
							Optional:     true,
							Default:      "Cluster",
							ValidateFunc: validation.StringInSlice([]string{"Local", "Cluster"}, false),
						},
						"load_balancer_ip": {
							Type:        schema.TypeString,
							Description: "Only applies to `type = LoadBalancer`. LoadBalancer will get created with the IP specified in this field. This feature depends on whether the underlying cloud-provider supports specifying this field when a load balancer is created. This field will be ignored if the cloud-provider does not support the feature.",
							Optional:    true,
						},
						"load_balancer_source_ranges": {
							Type:        schema.TypeSet,
							Description: "If specified and supported by the platform, this will restrict traffic through the cloud-provider load-balancer will be restricted to the specified client IPs. This field will be ignored if the cloud-provider does not support the feature. More info: http://kubernetes.io/docs/user-guide/services-firewalls",
							Optional:    true,
							Elem:        &schema.Schema{Type: schema.TypeString},
							Set:         schema.HashString,
						},
						"port": {
							Type:        schema.TypeList,
							Description: "The list of ports that are exposed by this service. More info: http://kubernetes.io/docs/user-guide/services#virtual-ips-and-service-proxies",
							Optional:    true,
							MinItems:    1,
							Elem: &schema.Resource{
								Schema: map[string]*schema.Schema{
									"name": {
										Type:        schema.TypeString,
										Description: "The name of this port within the service. All ports within the service must have unique names. Optional if only one ServicePort is defined on this service.",
										Optional:    true,
									},
									"node_port": {
										Type:        schema.TypeInt,
										Description: "The port on each node on which this service is exposed when `type` is `NodePort` or `LoadBalancer`. Usually assigned by the system. If specified, it will be allocated to the service if unused or else creation of the service will fail. Default is to auto-allocate a port if the `type` of this service requires one. More info: http://kubernetes.io/docs/user-guide/services#type--nodeport",
										Computed:    true,
										Optional:    true,
									},
									"port": {
										Type:        schema.TypeInt,
										Description: "The port that will be exposed by this service.",
										Required:    true,
									},
									"protocol": {
										Type:        schema.TypeString,
										Description: "The IP protocol for this port. Supports `TCP` and `UDP`. Default is `TCP`.",
										Optional:    true,
										Default:     "TCP",
									},
									"target_port": {
										Type:        schema.TypeString,
										Description: "Number or name of the port to access on the pods targeted by the service. Number must be in the range 1 to 65535. This field is ignored for services with `cluster_ip = \"None\"`. More info: http://kubernetes.io/docs/user-guide/services#defining-a-service",
										Optional:    true,
										Computed:    true,
									},
								},
							},
						},
						"publish_not_ready_addresses": {
							Type:        schema.TypeBool,
							Optional:    true,
							Default:     false,
							Description: "When set to true, indicates that DNS implementations must publish the `notReadyAddresses` of subsets for the Endpoints associated with the Service. The default value is `false`. The primary use case for setting this field is to use a StatefulSet's Headless Service to propagate `SRV` records for its Pods without respect to their readiness for purpose of peer discovery.",
						},
						"selector": {
							Type:        schema.TypeMap,
							Description: "Route service traffic to pods with label keys and values matching this selector. Only applies to types `ClusterIP`, `NodePort`, and `LoadBalancer`. More info: http://kubernetes.io/docs/user-guide/services#overview",
							Optional:    true,
						},
						"session_affinity": {
							Type:        schema.TypeString,
							Description: "Used to maintain session affinity. Supports `ClientIP` and `None`. Defaults to `None`. More info: http://kubernetes.io/docs/user-guide/services#virtual-ips-and-service-proxies",
							Optional:    true,
							Default:     "None",
						},
						"type": {
							Type:        schema.TypeString,
							Description: "Determines how the service is exposed. Defaults to `ClusterIP`. Valid options are `ExternalName`, `ClusterIP`, `NodePort`, and `LoadBalancer`. `ExternalName` maps to the specified `external_name`. More info: http://kubernetes.io/docs/user-guide/services#overview",
							Optional:    true,
							Default:     "ClusterIP",
						},
					},
				},
			},
			"load_balancer_ingress": {
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"ip": {
							Type:     schema.TypeString,
							Computed: true,
						},
						"hostname": {
							Type:     schema.TypeString,
							Computed: true,
						},
					},
				},
			},
		},
	}
}

func resourceKubernetesServiceCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*kubernetes.Clientset)

	metadata := expandMetadata(d.Get("metadata").([]interface{}))
	svc := api.Service{
		ObjectMeta: metadata,
		Spec:       expandServiceSpec(d.Get("spec").([]interface{})),
	}
	log.Printf("[INFO] Creating new service: %#v", svc)
	out, err := conn.CoreV1().Services(metadata.Namespace).Create(&svc)
	if err != nil {
		return err
	}
	log.Printf("[INFO] Submitted new service: %#v", out)
	d.SetId(buildId(out.ObjectMeta))

	if out.Spec.Type == api.ServiceTypeLoadBalancer {
		log.Printf("[DEBUG] Waiting for load balancer to assign IP/hostname")

		err = resource.Retry(10*time.Minute, func() *resource.RetryError {
			svc, err := conn.CoreV1().Services(out.Namespace).Get(out.Name, meta_v1.GetOptions{})
			if err != nil {
				log.Printf("[DEBUG] Received error: %#v", err)
				return resource.NonRetryableError(err)
			}

			lbIngress := svc.Status.LoadBalancer.Ingress

			log.Printf("[INFO] Received service status: %#v", svc.Status)
			if len(lbIngress) > 0 {
				return nil
			}

			return resource.RetryableError(fmt.Errorf(
				"Waiting for service %q to assign IP/hostname for a load balancer", d.Id()))
		})
		if err != nil {
			lastWarnings, wErr := getLastWarningsForObject(conn, out.ObjectMeta, "Service", 3)
			if wErr != nil {
				return wErr
			}
			return fmt.Errorf("%s%s", err, stringifyEvents(lastWarnings))
		}
	}

	return resourceKubernetesServiceRead(d, meta)
}

func resourceKubernetesServiceRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*kubernetes.Clientset)

	namespace, name, err := idParts(d.Id())
	if err != nil {
		return err
	}

	log.Printf("[INFO] Reading service %s", name)
	svc, err := conn.CoreV1().Services(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		log.Printf("[DEBUG] Received error: %#v", err)
		return err
	}
	log.Printf("[INFO] Received service: %#v", svc)
	err = d.Set("metadata", flattenMetadata(svc.ObjectMeta, d))
	if err != nil {
		return err
	}

	err = d.Set("load_balancer_ingress", flattenLoadBalancerIngress(svc.Status.LoadBalancer.Ingress))
	if err != nil {
		return err
	}

	flattened := flattenServiceSpec(svc.Spec)
	log.Printf("[DEBUG] Flattened service spec: %#v", flattened)
	err = d.Set("spec", flattened)
	if err != nil {
		return err
	}

	return nil
}

func resourceKubernetesServiceUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*kubernetes.Clientset)

	namespace, name, err := idParts(d.Id())
	if err != nil {
		return err
	}

	ops := patchMetadata("metadata.0.", "/metadata/", d)
	if d.HasChange("spec") {
		serverVersion, err := conn.ServerVersion()
		if err != nil {
			return err
		}
		diffOps, err := patchServiceSpec("spec.0.", "/spec/", d, serverVersion)
		if err != nil {
			return err
		}
		ops = append(ops, diffOps...)
	}
	data, err := ops.MarshalJSON()
	if err != nil {
		return fmt.Errorf("Failed to marshal update operations: %s", err)
	}
	log.Printf("[INFO] Updating service %q: %v", name, string(data))
	out, err := conn.CoreV1().Services(namespace).Patch(name, pkgApi.JSONPatchType, data)
	if err != nil {
		return fmt.Errorf("Failed to update service: %s", err)
	}
	log.Printf("[INFO] Submitted updated service: %#v", out)
	d.SetId(buildId(out.ObjectMeta))

	return resourceKubernetesServiceRead(d, meta)
}

func resourceKubernetesServiceDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*kubernetes.Clientset)

	namespace, name, err := idParts(d.Id())
	if err != nil {
		return err
	}

	log.Printf("[INFO] Deleting service: %#v", name)
	err = conn.CoreV1().Services(namespace).Delete(name, &meta_v1.DeleteOptions{})
	if err != nil {
		return err
	}

	log.Printf("[INFO] Service %s deleted", name)

	d.SetId("")
	return nil
}

func resourceKubernetesServiceExists(d *schema.ResourceData, meta interface{}) (bool, error) {
	conn := meta.(*kubernetes.Clientset)

	namespace, name, err := idParts(d.Id())
	if err != nil {
		return false, err
	}

	log.Printf("[INFO] Checking service %s", name)
	_, err = conn.CoreV1().Services(namespace).Get(name, meta_v1.GetOptions{})
	if err != nil {
		if statusErr, ok := err.(*errors.StatusError); ok && statusErr.ErrStatus.Code == 404 {
			return false, nil
		}
		log.Printf("[DEBUG] Received error: %#v", err)
	}
	return true, err
}
<|endoftext|>


<commit_before>package bc

import (
	"bytes"
	"database/sql/driver"
	"encoding/hex"
	"fmt"
	"io"
	"strconv"

	"chain/crypto/sha3pool"
	"chain/encoding/blockchain"
	"chain/errors"
)

// CurrentTransactionVersion is the current latest
// supported transaction version.
const CurrentTransactionVersion = 1

// Tx holds a transaction along with its hash.
type Tx struct {
	TxData
	Hash Hash
}

func (tx *Tx) UnmarshalText(p []byte) error {
	if err := tx.TxData.UnmarshalText(p); err != nil {
		return err
	}

	tx.Hash = tx.TxData.Hash()
	return nil
}

// NewTx returns a new Tx containing data and its hash.
// If you have already computed the hash, use struct literal
// notation to make a Tx object directly.
func NewTx(data TxData) *Tx {
	return &Tx{
		TxData: data,
		Hash:   data.Hash(),
	}
}

// These flags are part of the wire protocol;
// they must not change.
const (
	SerWitness uint8 = 1 << iota
	SerPrevout
	SerMetadata

	// Bit mask for accepted serialization flags.
	// All other flag bits must be 0.
	SerValid    = 0x7
	serRequired = 0x7 // we support only this combination of flags
)

// TxData encodes a transaction in the blockchain.
// Most users will want to use Tx instead;
// it includes the hash.
type TxData struct {
	Version       uint64
	Inputs        []*TxInput
	Outputs       []*TxOutput
	MinTime       uint64
	MaxTime       uint64
	ReferenceData []byte
}

// Outpoint defines a bitcoin data type that is used to track previous
// transaction outputs.
type Outpoint struct {
	Hash  Hash   `json:"hash"`
	Index uint32 `json:"index"`
}

// HasIssuance returns true if this transaction has an issuance input.
func (tx *TxData) HasIssuance() bool {
	for _, in := range tx.Inputs {
		if in.IsIssuance() {
			return true
		}
	}
	return false
}

func (tx *TxData) UnmarshalText(p []byte) error {
	b := make([]byte, hex.DecodedLen(len(p)))
	_, err := hex.Decode(b, p)
	if err != nil {
		return err
	}
	return tx.readFrom(bytes.NewReader(b))
}

func (tx *TxData) Scan(val interface{}) error {
	b, ok := val.([]byte)
	if !ok {
		return errors.New("Scan must receive a byte slice")
	}
	return tx.readFrom(bytes.NewReader(b))
}

func (tx *TxData) Value() (driver.Value, error) {
	b := new(bytes.Buffer)
	_, err := tx.WriteTo(b)
	if err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

func (tx *TxData) readFrom(r io.Reader) error {
	var serflags [1]byte
	_, err := io.ReadFull(r, serflags[:])
	if err != nil {
		return errors.Wrap(err, "reading serialization flags")
	}
	if err == nil && serflags[0] != serRequired {
		return fmt.Errorf("unsupported serflags %#x", serflags[0])
	}

	tx.Version, _, err = blockchain.ReadVarint63(r)
	if err != nil {
		return errors.Wrap(err, "reading transaction version")
	}

	// Common fields
	all := tx.Version == 1
	_, err = blockchain.ReadExtensibleString(r, all, func(r io.Reader) error {
		tx.MinTime, _, err = blockchain.ReadVarint63(r)
		if err != nil {
			return errors.Wrap(err, "reading transaction mintime")
		}
		tx.MaxTime, _, err = blockchain.ReadVarint63(r)
		return errors.Wrap(err, "reading transaction maxtime")
	})
	if err != nil {
		return errors.Wrap(err, "reading transaction common fields")
	}

	// Common witness, empty in v1
	_, _, err = blockchain.ReadVarstr31(r)
	if err != nil {
		return errors.Wrap(err, "reading transaction common witness")
	}

	n, _, err := blockchain.ReadVarint31(r)
	if err != nil {
		return errors.Wrap(err, "reading number of transaction inputs")
	}
	for ; n > 0; n-- {
		ti := new(TxInput)
		err = ti.readFrom(r, tx.Version)
		if err != nil {
			return errors.Wrapf(err, "reading input %d", len(tx.Inputs))
		}
		tx.Inputs = append(tx.Inputs, ti)
	}

	n, _, err = blockchain.ReadVarint31(r)
	if err != nil {
		return errors.Wrap(err, "reading number of transaction outputs")
	}
	for ; n > 0; n-- {
		to := new(TxOutput)
		err = to.readFrom(r, tx.Version)
		if err != nil {
			return errors.Wrapf(err, "reading output %d", len(tx.Outputs))
		}
		tx.Outputs = append(tx.Outputs, to)
	}

	tx.ReferenceData, _, err = blockchain.ReadVarstr31(r)
	return errors.Wrap(err, "reading transaction reference data")
}

func (p *Outpoint) readFrom(r io.Reader) (int, error) {
	n1, err := io.ReadFull(r, p.Hash[:])
	if err != nil {
		return n1, err
	}
	var n2 int
	p.Index, n2, err = blockchain.ReadVarint31(r)
	return n1 + n2, err
}

// Hash computes the hash of the transaction with reference data fields
// replaced by their hashes,
// and stores the result in Hash.
func (tx *TxData) Hash() Hash {
	h := sha3pool.Get256()
	tx.writeTo(h, 0) // error is impossible
	var v Hash
	h.Read(v[:])
	sha3pool.Put256(h)
	return v
}

// WitnessHash is the combined hash of the
// transactions hash and signature data hash.
// It is used to compute the TxRoot of a block.
func (tx *Tx) WitnessHash() (hash Hash, err error) {
	hasher := sha3pool.Get256()
	defer sha3pool.Put256(hasher)

	hasher.Write(tx.Hash[:])

	blockchain.WriteVarint31(hasher, uint64(len(tx.Inputs))) // TODO(bobg): check and return error
	for _, txin := range tx.Inputs {
		h, err := txin.witnessHash()
		if err != nil {
			return hash, err
		}
		hasher.Write(h[:])
	}

	blockchain.WriteVarint31(hasher, uint64(len(tx.Outputs))) // TODO(bobg): check and return error
	for _, txout := range tx.Outputs {
		h := txout.witnessHash()
		hasher.Write(h[:])
	}

	hasher.Read(hash[:])
	return hash, nil
}

func (tx *TxData) IssuanceHash(n int) (h Hash, err error) {
	if n < 0 || n >= len(tx.Inputs) {
		return h, fmt.Errorf("no input %d", n)
	}
	ii, ok := tx.Inputs[n].TypedInput.(*IssuanceInput)
	if !ok {
		return h, fmt.Errorf("not an issuance input")
	}
	buf := sha3pool.Get256()
	defer sha3pool.Put256(buf)

	_, err = blockchain.WriteVarstr31(buf, ii.Nonce)
	if err != nil {
		return h, err
	}
	assetID := ii.AssetID()
	buf.Write(assetID[:])
	_, err = blockchain.WriteVarint63(buf, tx.MinTime)
	if err != nil {
		return h, err
	}
	_, err = blockchain.WriteVarint63(buf, tx.MaxTime)
	if err != nil {
		return h, err
	}
	buf.Read(h[:])
	return h, nil
}

// HashForSig generates the hash required for the specified input's
// signature.
func (tx *TxData) HashForSig(idx uint32) Hash {
	return NewSigHasher(tx).Hash(idx)
}

func (tx *TxData) MarshalText() ([]byte, error) {
	var buf bytes.Buffer
	tx.WriteTo(&buf) // error is impossible
	b := make([]byte, hex.EncodedLen(buf.Len()))
	hex.Encode(b, buf.Bytes())
	return b, nil
}

// WriteTo writes tx to w.
func (tx *TxData) WriteTo(w io.Writer) (int64, error) {
	ew := errors.NewWriter(w)
	err := tx.writeTo(ew, serRequired)
	if err != nil {
		return ew.Written(), ew.Err()
	}
	return ew.Written(), ew.Err()
}

func (tx *TxData) writeTo(w io.Writer, serflags byte) error {
	_, err := w.Write([]byte{serflags})
	if err != nil {
		return errors.Wrap(err, "writing serialization flags")
	}

	_, err = blockchain.WriteVarint63(w, tx.Version)
	if err != nil {
		return errors.Wrap(err, "writing transaction version")
	}

	// common fields
	_, err = blockchain.WriteExtensibleString(w, func(w io.Writer) error {
		_, err := blockchain.WriteVarint63(w, tx.MinTime)
		if err != nil {
			return errors.Wrap(err, "writing transaction min time")
		}
		_, err = blockchain.WriteVarint63(w, tx.MaxTime)
		return errors.Wrap(err, "writing transaction max time")
	})
	if err != nil {
		return errors.Wrap(err, "writing common fields")
	}

	// common witness, empty in v1
	_, err = blockchain.WriteVarstr31(w, []byte{})
	if err != nil {
		return errors.Wrap(err, "writing common witness")
	}

	_, err = blockchain.WriteVarint31(w, uint64(len(tx.Inputs)))
	if err != nil {
		return errors.Wrap(err, "writing tx input count")
	}
	for i, ti := range tx.Inputs {
		err = ti.writeTo(w, serflags)
		if err != nil {
			return errors.Wrapf(err, "writing tx input %d", i)
		}
	}

	_, err = blockchain.WriteVarint31(w, uint64(len(tx.Outputs)))
	if err != nil {
		return errors.Wrap(err, "writing tx output count")
	}
	for i, to := range tx.Outputs {
		err = to.writeTo(w, serflags)
		if err != nil {
			return errors.Wrapf(err, "writing tx output %d", i)
		}
	}

	return writeRefData(w, tx.ReferenceData, serflags)
}

// String returns the Outpoint in the human-readable form "hash:index".
func (p Outpoint) String() string {
	return p.Hash.String() + ":" + strconv.FormatUint(uint64(p.Index), 10)
}

// WriteTo writes p to w.
func (p *Outpoint) WriteTo(w io.Writer) (int64, error) {
	n, err := w.Write(p.Hash[:])
	if err != nil {
		return int64(n), err
	}
	n2, err := blockchain.WriteVarint31(w, uint64(p.Index))
	return int64(n + n2), err
}

func writeRefData(w io.Writer, data []byte, serflags byte) error {
	if serflags&SerMetadata != 0 {
		_, err := blockchain.WriteVarstr31(w, data) // TODO(bobg): check and return error
		return err
	}
	return writeFastHash(w, data)
}
<commit_msg>protocol/bc: refactor common witness serialization<commit_after>package bc

import (
	"bytes"
	"database/sql/driver"
	"encoding/hex"
	"fmt"
	"io"
	"strconv"

	"chain/crypto/sha3pool"
	"chain/encoding/blockchain"
	"chain/errors"
)

// CurrentTransactionVersion is the current latest
// supported transaction version.
const CurrentTransactionVersion = 1

// Tx holds a transaction along with its hash.
type Tx struct {
	TxData
	Hash Hash
}

func (tx *Tx) UnmarshalText(p []byte) error {
	if err := tx.TxData.UnmarshalText(p); err != nil {
		return err
	}

	tx.Hash = tx.TxData.Hash()
	return nil
}

// NewTx returns a new Tx containing data and its hash.
// If you have already computed the hash, use struct literal
// notation to make a Tx object directly.
func NewTx(data TxData) *Tx {
	return &Tx{
		TxData: data,
		Hash:   data.Hash(),
	}
}

// These flags are part of the wire protocol;
// they must not change.
const (
	SerWitness uint8 = 1 << iota
	SerPrevout
	SerMetadata

	// Bit mask for accepted serialization flags.
	// All other flag bits must be 0.
	SerValid    = 0x7
	serRequired = 0x7 // we support only this combination of flags
)

// TxData encodes a transaction in the blockchain.
// Most users will want to use Tx instead;
// it includes the hash.
type TxData struct {
	Version       uint64
	Inputs        []*TxInput
	Outputs       []*TxOutput
	MinTime       uint64
	MaxTime       uint64
	ReferenceData []byte
}

// Outpoint defines a bitcoin data type that is used to track previous
// transaction outputs.
type Outpoint struct {
	Hash  Hash   `json:"hash"`
	Index uint32 `json:"index"`
}

// HasIssuance returns true if this transaction has an issuance input.
func (tx *TxData) HasIssuance() bool {
	for _, in := range tx.Inputs {
		if in.IsIssuance() {
			return true
		}
	}
	return false
}

func (tx *TxData) UnmarshalText(p []byte) error {
	b := make([]byte, hex.DecodedLen(len(p)))
	_, err := hex.Decode(b, p)
	if err != nil {
		return err
	}
	return tx.readFrom(bytes.NewReader(b))
}

func (tx *TxData) Scan(val interface{}) error {
	b, ok := val.([]byte)
	if !ok {
		return errors.New("Scan must receive a byte slice")
	}
	return tx.readFrom(bytes.NewReader(b))
}

func (tx *TxData) Value() (driver.Value, error) {
	b := new(bytes.Buffer)
	_, err := tx.WriteTo(b)
	if err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

func (tx *TxData) readFrom(r io.Reader) error {
	var serflags [1]byte
	_, err := io.ReadFull(r, serflags[:])
	if err != nil {
		return errors.Wrap(err, "reading serialization flags")
	}
	if err == nil && serflags[0] != serRequired {
		return fmt.Errorf("unsupported serflags %#x", serflags[0])
	}

	tx.Version, _, err = blockchain.ReadVarint63(r)
	if err != nil {
		return errors.Wrap(err, "reading transaction version")
	}

	// Common fields
	all := tx.Version == 1
	_, err = blockchain.ReadExtensibleString(r, all, func(r io.Reader) error {
		tx.MinTime, _, err = blockchain.ReadVarint63(r)
		if err != nil {
			return errors.Wrap(err, "reading transaction mintime")
		}
		tx.MaxTime, _, err = blockchain.ReadVarint63(r)
		return errors.Wrap(err, "reading transaction maxtime")
	})
	if err != nil {
		return errors.Wrap(err, "reading transaction common fields")
	}

	// Common witness
	_, err = blockchain.ReadExtensibleString(r, false, func(r io.Reader) error {
		return tx.readCommonWitness(r)
	})
	if err != nil {
		return errors.Wrap(err, "reading transaction common witness")
	}

	n, _, err := blockchain.ReadVarint31(r)
	if err != nil {
		return errors.Wrap(err, "reading number of transaction inputs")
	}
	for ; n > 0; n-- {
		ti := new(TxInput)
		err = ti.readFrom(r, tx.Version)
		if err != nil {
			return errors.Wrapf(err, "reading input %d", len(tx.Inputs))
		}
		tx.Inputs = append(tx.Inputs, ti)
	}

	n, _, err = blockchain.ReadVarint31(r)
	if err != nil {
		return errors.Wrap(err, "reading number of transaction outputs")
	}
	for ; n > 0; n-- {
		to := new(TxOutput)
		err = to.readFrom(r, tx.Version)
		if err != nil {
			return errors.Wrapf(err, "reading output %d", len(tx.Outputs))
		}
		tx.Outputs = append(tx.Outputs, to)
	}

	tx.ReferenceData, _, err = blockchain.ReadVarstr31(r)
	return errors.Wrap(err, "reading transaction reference data")
}

// does not read the enclosing extensible string
func (tx *TxData) readCommonWitness(r io.Reader) error {
	return nil
}

func (p *Outpoint) readFrom(r io.Reader) (int, error) {
	n1, err := io.ReadFull(r, p.Hash[:])
	if err != nil {
		return n1, err
	}
	var n2 int
	p.Index, n2, err = blockchain.ReadVarint31(r)
	return n1 + n2, err
}

// Hash computes the hash of the transaction with reference data fields
// replaced by their hashes,
// and stores the result in Hash.
func (tx *TxData) Hash() Hash {
	h := sha3pool.Get256()
	tx.writeTo(h, 0) // error is impossible
	var v Hash
	h.Read(v[:])
	sha3pool.Put256(h)
	return v
}

// WitnessHash is the combined hash of the
// transactions hash and signature data hash.
// It is used to compute the TxRoot of a block.
func (tx *Tx) WitnessHash() (hash Hash, err error) {
	hasher := sha3pool.Get256()
	defer sha3pool.Put256(hasher)

	hasher.Write(tx.Hash[:])

	blockchain.WriteVarint31(hasher, uint64(len(tx.Inputs))) // TODO(bobg): check and return error
	for _, txin := range tx.Inputs {
		h, err := txin.witnessHash()
		if err != nil {
			return hash, err
		}
		hasher.Write(h[:])
	}

	blockchain.WriteVarint31(hasher, uint64(len(tx.Outputs))) // TODO(bobg): check and return error
	for _, txout := range tx.Outputs {
		h := txout.witnessHash()
		hasher.Write(h[:])
	}

	hasher.Read(hash[:])
	return hash, nil
}

func (tx *TxData) IssuanceHash(n int) (h Hash, err error) {
	if n < 0 || n >= len(tx.Inputs) {
		return h, fmt.Errorf("no input %d", n)
	}
	ii, ok := tx.Inputs[n].TypedInput.(*IssuanceInput)
	if !ok {
		return h, fmt.Errorf("not an issuance input")
	}
	buf := sha3pool.Get256()
	defer sha3pool.Put256(buf)

	_, err = blockchain.WriteVarstr31(buf, ii.Nonce)
	if err != nil {
		return h, err
	}
	assetID := ii.AssetID()
	buf.Write(assetID[:])
	_, err = blockchain.WriteVarint63(buf, tx.MinTime)
	if err != nil {
		return h, err
	}
	_, err = blockchain.WriteVarint63(buf, tx.MaxTime)
	if err != nil {
		return h, err
	}
	buf.Read(h[:])
	return h, nil
}

// HashForSig generates the hash required for the specified input's
// signature.
func (tx *TxData) HashForSig(idx uint32) Hash {
	return NewSigHasher(tx).Hash(idx)
}

func (tx *TxData) MarshalText() ([]byte, error) {
	var buf bytes.Buffer
	tx.WriteTo(&buf) // error is impossible
	b := make([]byte, hex.EncodedLen(buf.Len()))
	hex.Encode(b, buf.Bytes())
	return b, nil
}

// WriteTo writes tx to w.
func (tx *TxData) WriteTo(w io.Writer) (int64, error) {
	ew := errors.NewWriter(w)
	err := tx.writeTo(ew, serRequired)
	if err != nil {
		return ew.Written(), ew.Err()
	}
	return ew.Written(), ew.Err()
}

func (tx *TxData) writeTo(w io.Writer, serflags byte) error {
	_, err := w.Write([]byte{serflags})
	if err != nil {
		return errors.Wrap(err, "writing serialization flags")
	}

	_, err = blockchain.WriteVarint63(w, tx.Version)
	if err != nil {
		return errors.Wrap(err, "writing transaction version")
	}

	// common fields
	_, err = blockchain.WriteExtensibleString(w, func(w io.Writer) error {
		_, err := blockchain.WriteVarint63(w, tx.MinTime)
		if err != nil {
			return errors.Wrap(err, "writing transaction min time")
		}
		_, err = blockchain.WriteVarint63(w, tx.MaxTime)
		return errors.Wrap(err, "writing transaction max time")
	})
	if err != nil {
		return errors.Wrap(err, "writing common fields")
	}

	// common witness
	_, err = blockchain.WriteExtensibleString(w, func(w io.Writer) error {
		return tx.writeCommonWitness(w)
	})
	if err != nil {
		return errors.Wrap(err, "writing common witness")
	}

	_, err = blockchain.WriteVarint31(w, uint64(len(tx.Inputs)))
	if err != nil {
		return errors.Wrap(err, "writing tx input count")
	}
	for i, ti := range tx.Inputs {
		err = ti.writeTo(w, serflags)
		if err != nil {
			return errors.Wrapf(err, "writing tx input %d", i)
		}
	}

	_, err = blockchain.WriteVarint31(w, uint64(len(tx.Outputs)))
	if err != nil {
		return errors.Wrap(err, "writing tx output count")
	}
	for i, to := range tx.Outputs {
		err = to.writeTo(w, serflags)
		if err != nil {
			return errors.Wrapf(err, "writing tx output %d", i)
		}
	}

	return writeRefData(w, tx.ReferenceData, serflags)
}

// does not write the enclosing extensible string
func (tx *TxData) writeCommonWitness(w io.Writer) error {
	return nil
}

// String returns the Outpoint in the human-readable form "hash:index".
func (p Outpoint) String() string {
	return p.Hash.String() + ":" + strconv.FormatUint(uint64(p.Index), 10)
}

// WriteTo writes p to w.
func (p *Outpoint) WriteTo(w io.Writer) (int64, error) {
	n, err := w.Write(p.Hash[:])
	if err != nil {
		return int64(n), err
	}
	n2, err := blockchain.WriteVarint31(w, uint64(p.Index))
	return int64(n + n2), err
}

func writeRefData(w io.Writer, data []byte, serflags byte) error {
	if serflags&SerMetadata != 0 {
		_, err := blockchain.WriteVarstr31(w, data) // TODO(bobg): check and return error
		return err
	}
	return writeFastHash(w, data)
}
<|endoftext|>


<commit_before>package hcsshim

import (
	"errors"
	"os"
	"path/filepath"
	"syscall"

	"github.com/Microsoft/go-winio"
)

type baseLayerWriter struct {
	root         string
	f            *os.File
	bw           *winio.BackupFileWriter
	err          error
	hasUtilityVM bool
	dirInfo      []dirInfo
}

type dirInfo struct {
	path     string
	fileInfo winio.FileBasicInfo
}

func (w *baseLayerWriter) closeCurrentFile() error {
	if w.f != nil {
		err := w.bw.Close()
		err2 := w.f.Close()
		w.f = nil
		w.bw = nil
		if err != nil {
			return err
		}
		if err2 != nil {
			return err2
		}
	}
	return nil
}

func (w *baseLayerWriter) Add(name string, fileInfo *winio.FileBasicInfo) (err error) {
	defer func() {
		if err != nil {
			w.err = err
		}
	}()

	err = w.closeCurrentFile()
	if err != nil {
		return err
	}

	if filepath.ToSlash(name) == `UtilityVM/Files` {
		w.hasUtilityVM = true
	}

	path := filepath.Join(w.root, name)
	path, err = makeLongAbsPath(path)
	if err != nil {
		return err
	}

	var f *os.File
	defer func() {
		if f != nil {
			f.Close()
		}
	}()

	createmode := uint32(syscall.CREATE_NEW)
	if fileInfo.FileAttributes&syscall.FILE_ATTRIBUTE_DIRECTORY != 0 {
		err := os.Mkdir(path, 0)
		if err != nil && !os.IsExist(err) {
			return err
		}
		createmode = syscall.OPEN_EXISTING
		w.dirInfo = append(w.dirInfo, dirInfo{path, *fileInfo})
	}

	mode := uint32(syscall.GENERIC_READ | syscall.GENERIC_WRITE | winio.WRITE_DAC | winio.WRITE_OWNER | winio.ACCESS_SYSTEM_SECURITY)
	f, err = winio.OpenForBackup(path, mode, syscall.FILE_SHARE_READ, createmode)
	if err != nil {
		return err
	}

	err = winio.SetFileBasicInfo(f, fileInfo)
	if err != nil {
		return err
	}

	w.f = f
	w.bw = winio.NewBackupFileWriter(f, true)
	f = nil
	return nil
}

func (w *baseLayerWriter) AddLink(name string, target string) (err error) {
	defer func() {
		if err != nil {
			w.err = err
		}
	}()

	err = w.closeCurrentFile()
	if err != nil {
		return err
	}

	linkpath, err := makeLongAbsPath(filepath.Join(w.root, name))
	if err != nil {
		return err
	}

	linktarget, err := makeLongAbsPath(filepath.Join(w.root, target))
	if err != nil {
		return err
	}

	return os.Link(linktarget, linkpath)
}

func (w *baseLayerWriter) Remove(name string) error {
	return errors.New("base layer cannot have tombstones")
}

func (w *baseLayerWriter) Write(b []byte) (int, error) {
	n, err := w.bw.Write(b)
	if err != nil {
		w.err = err
	}
	return n, err
}

func (w *baseLayerWriter) Close() error {
	err := w.closeCurrentFile()
	if err != nil {
		return err
	}
	if w.err == nil {
		// Restore the file times of all the directories, since they may have
		// been modified by creating child directories.
		for i := range w.dirInfo {
			di := &w.dirInfo[len(w.dirInfo)-i-1]
			f, err := winio.OpenForBackup(di.path, uint32(syscall.GENERIC_READ|syscall.GENERIC_WRITE), syscall.FILE_SHARE_READ, syscall.OPEN_EXISTING)
			if err != nil {
				return err
			}

			err = winio.SetFileBasicInfo(f, &di.fileInfo)
			f.Close()
			if err != nil {
				return err
			}
		}

		err = ProcessBaseLayer(w.root)
		if err != nil {
			return err
		}

		if w.hasUtilityVM {
			err = ProcessUtilityVMImage(filepath.Join(w.root, "UtilityVM"))
			if err != nil {
				return err
			}
		}
	}
	return w.err
}
<commit_msg>Only update directory modified times if not a reparse point<commit_after>package hcsshim

import (
	"errors"
	"os"
	"path/filepath"
	"syscall"

	"github.com/Microsoft/go-winio"
)

type baseLayerWriter struct {
	root         string
	f            *os.File
	bw           *winio.BackupFileWriter
	err          error
	hasUtilityVM bool
	dirInfo      []dirInfo
}

type dirInfo struct {
	path     string
	fileInfo winio.FileBasicInfo
}

func (w *baseLayerWriter) closeCurrentFile() error {
	if w.f != nil {
		err := w.bw.Close()
		err2 := w.f.Close()
		w.f = nil
		w.bw = nil
		if err != nil {
			return err
		}
		if err2 != nil {
			return err2
		}
	}
	return nil
}

func (w *baseLayerWriter) Add(name string, fileInfo *winio.FileBasicInfo) (err error) {
	defer func() {
		if err != nil {
			w.err = err
		}
	}()

	err = w.closeCurrentFile()
	if err != nil {
		return err
	}

	if filepath.ToSlash(name) == `UtilityVM/Files` {
		w.hasUtilityVM = true
	}

	path := filepath.Join(w.root, name)
	path, err = makeLongAbsPath(path)
	if err != nil {
		return err
	}

	var f *os.File
	defer func() {
		if f != nil {
			f.Close()
		}
	}()

	createmode := uint32(syscall.CREATE_NEW)
	if fileInfo.FileAttributes&syscall.FILE_ATTRIBUTE_DIRECTORY != 0 {
		err := os.Mkdir(path, 0)
		if err != nil && !os.IsExist(err) {
			return err
		}
		createmode = syscall.OPEN_EXISTING
		if fileInfo.FileAttributes&syscall.FILE_ATTRIBUTE_REPARSE_POINT == 0 {
			w.dirInfo = append(w.dirInfo, dirInfo{path, *fileInfo})
		}
	}

	mode := uint32(syscall.GENERIC_READ | syscall.GENERIC_WRITE | winio.WRITE_DAC | winio.WRITE_OWNER | winio.ACCESS_SYSTEM_SECURITY)
	f, err = winio.OpenForBackup(path, mode, syscall.FILE_SHARE_READ, createmode)
	if err != nil {
		return makeError(err, "Failed to OpenForBackup", path)
	}

	err = winio.SetFileBasicInfo(f, fileInfo)
	if err != nil {
		return makeError(err, "Failed to SetFileBasicInfo", path)
	}

	w.f = f
	w.bw = winio.NewBackupFileWriter(f, true)
	f = nil
	return nil
}

func (w *baseLayerWriter) AddLink(name string, target string) (err error) {
	defer func() {
		if err != nil {
			w.err = err
		}
	}()

	err = w.closeCurrentFile()
	if err != nil {
		return err
	}

	linkpath, err := makeLongAbsPath(filepath.Join(w.root, name))
	if err != nil {
		return err
	}

	linktarget, err := makeLongAbsPath(filepath.Join(w.root, target))
	if err != nil {
		return err
	}

	return os.Link(linktarget, linkpath)
}

func (w *baseLayerWriter) Remove(name string) error {
	return errors.New("base layer cannot have tombstones")
}

func (w *baseLayerWriter) Write(b []byte) (int, error) {
	n, err := w.bw.Write(b)
	if err != nil {
		w.err = err
	}
	return n, err
}

func (w *baseLayerWriter) Close() error {
	err := w.closeCurrentFile()
	if err != nil {
		return err
	}
	if w.err == nil {
		// Restore the file times of all the directories, since they may have
		// been modified by creating child directories.
		for i := range w.dirInfo {
			di := &w.dirInfo[len(w.dirInfo)-i-1]
			f, err := winio.OpenForBackup(di.path, uint32(syscall.GENERIC_READ|syscall.GENERIC_WRITE), syscall.FILE_SHARE_READ, syscall.OPEN_EXISTING)
			if err != nil {
				return makeError(err, "Failed to OpenForBackup", di.path)
			}

			err = winio.SetFileBasicInfo(f, &di.fileInfo)
			f.Close()
			if err != nil {
				return makeError(err, "Failed to SetFileBasicInfo", di.path)
			}
		}

		err = ProcessBaseLayer(w.root)
		if err != nil {
			return err
		}

		if w.hasUtilityVM {
			err = ProcessUtilityVMImage(filepath.Join(w.root, "UtilityVM"))
			if err != nil {
				return err
			}
		}
	}
	return w.err
}
<|endoftext|>


<commit_before>package basicauth

import (
	"fmt"
	"net/http"
)

var (
	// DEBUG is debug mode.
	// Set to true to enable debug messages.
	DEBUG bool = false
)

const (
	DefBasicRealmStr = "Please input username and password" // Default "Basic realm"
)

// Arguments is used to customize Basic Auth UI(Ex: "Basic realm" string
type Arguments struct {
	BasicRealmStr string // "Basic realm" string
}

// BasicAuth provides function to process HTTP Basic Auth
type BasicAuth struct {
	UserName string    // User Name
	Password string    // Password
	Args     Arguments // Arguments to customize Basic Auth UI(Ex: "Basic realm" string).
}

// New creates a new BasicAuth with default arguments.
func New(username, password string) (ba *BasicAuth) {
	return &BasicAuth{username, password, Arguments{DefBasicRealmStr}}
}

// NewWithArgs creates a new BasicAuth with user's arguments(Ex: "Basic realm" string).
//
// See Arguments for more.
func NewWithArgs(username, password string, args Arguments) (ba *BasicAuth) {
	return &BasicAuth{username, password, args}
}

// IsOK checks HTTP Basic Auth.
func (ba *BasicAuth) IsOK(w http.ResponseWriter, r *http.Request) bool {
	username, password, ok := r.BasicAuth()
	if !ok {
		if DEBUG {
			fmt.Printf("r.BasicAuth() failed.\n")
		}
		goto end
	}

	if username != ba.UserName || password != ba.Password {
		ok = false
		goto end
	}

end:
	if !ok {
		w.Header().Set("WWW-Authenticate", fmt.Sprintf(`Basic realm="%v"`, ba.Args.BasicRealmStr))
		w.WriteHeader(401)
	}
	return ok
}
<commit_msg>Omit type in 'var xx type = xx' to remove golint warnings.<commit_after>package basicauth

import (
	"fmt"
	"net/http"
)

var (
	// DEBUG is debug mode.
	// Set to true to enable debug messages.
	DEBUG = false
)

const (
	DefBasicRealmStr = "Please input username and password" // Default "Basic realm"
)

// Arguments is used to customize Basic Auth UI(Ex: "Basic realm" string
type Arguments struct {
	BasicRealmStr string // "Basic realm" string
}

// BasicAuth provides function to process HTTP Basic Auth
type BasicAuth struct {
	UserName string    // User Name
	Password string    // Password
	Args     Arguments // Arguments to customize Basic Auth UI(Ex: "Basic realm" string).
}

// New creates a new BasicAuth with default arguments.
func New(username, password string) (ba *BasicAuth) {
	return &BasicAuth{username, password, Arguments{DefBasicRealmStr}}
}

// NewWithArgs creates a new BasicAuth with user's arguments(Ex: "Basic realm" string).
//
// See Arguments for more.
func NewWithArgs(username, password string, args Arguments) (ba *BasicAuth) {
	return &BasicAuth{username, password, args}
}

// IsOK checks HTTP Basic Auth.
func (ba *BasicAuth) IsOK(w http.ResponseWriter, r *http.Request) bool {
	username, password, ok := r.BasicAuth()
	if !ok {
		if DEBUG {
			fmt.Printf("r.BasicAuth() failed.\n")
		}
		goto end
	}

	if username != ba.UserName || password != ba.Password {
		ok = false
		goto end
	}

end:
	if !ok {
		w.Header().Set("WWW-Authenticate", fmt.Sprintf(`Basic realm="%v"`, ba.Args.BasicRealmStr))
		w.WriteHeader(401)
	}
	return ok
}
<|endoftext|>


<commit_before>package pstoreds

import (
	"context"
	"io/ioutil"
	"os"
	"testing"
	"time"

	ds "github.com/ipfs/go-datastore"
	badger "github.com/ipfs/go-ds-badger"
	leveldb "github.com/ipfs/go-ds-leveldb"

	pstore "github.com/libp2p/go-libp2p-core/peerstore"
	pt "github.com/libp2p/go-libp2p-peerstore/test"
)

type datastoreFactory func(tb testing.TB) (ds.Batching, func())

var dstores = map[string]datastoreFactory{
	//"Badger": badgerStore,
	"Leveldb": leveldbStore,
}

func TestDsPeerstore(t *testing.T) {
	for name, dsFactory := range dstores {
		t.Run(name, func(t *testing.T) {
			pt.TestPeerstore(t, peerstoreFactory(t, dsFactory, DefaultOpts()))
		})
	}
}

func TestDsAddrBook(t *testing.T) {
	for name, dsFactory := range dstores {
		t.Run(name+" Cacheful", func(t *testing.T) {
			opts := DefaultOpts()
			opts.GCPurgeInterval = 1 * time.Second
			opts.CacheSize = 1024

			pt.TestAddrBook(t, addressBookFactory(t, dsFactory, opts))
		})

		t.Run(name+" Cacheless", func(t *testing.T) {
			opts := DefaultOpts()
			opts.GCPurgeInterval = 1 * time.Second
			opts.CacheSize = 0

			pt.TestAddrBook(t, addressBookFactory(t, dsFactory, opts))
		})
	}
}

func TestDsKeyBook(t *testing.T) {
	for name, dsFactory := range dstores {
		t.Run(name, func(t *testing.T) {
			pt.TestKeyBook(t, keyBookFactory(t, dsFactory, DefaultOpts()))
		})
	}
}

func BenchmarkDsKeyBook(b *testing.B) {
	for name, dsFactory := range dstores {
		b.Run(name, func(b *testing.B) {
			pt.BenchmarkKeyBook(b, keyBookFactory(b, dsFactory, DefaultOpts()))
		})
	}
}

func BenchmarkDsPeerstore(b *testing.B) {
	caching := DefaultOpts()
	caching.CacheSize = 1024

	cacheless := DefaultOpts()
	cacheless.CacheSize = 0

	for name, dsFactory := range dstores {
		b.Run(name, func(b *testing.B) {
			pt.BenchmarkPeerstore(b, peerstoreFactory(b, dsFactory, caching), "Caching")
		})
		b.Run(name, func(b *testing.B) {
			pt.BenchmarkPeerstore(b, peerstoreFactory(b, dsFactory, cacheless), "Cacheless")
		})
	}
}

// Doesn't work on 32bit because badger.
//lint:ignore U1000 disabled for now
func badgerStore(tb testing.TB) (ds.Batching, func()) {
	dataPath, err := ioutil.TempDir(os.TempDir(), "badger")
	if err != nil {
		tb.Fatal(err)
	}
	store, err := badger.NewDatastore(dataPath, nil)
	if err != nil {
		tb.Fatal(err)
	}
	closer := func() {
		store.Close()
		os.RemoveAll(dataPath)
	}
	return store, closer
}

func leveldbStore(tb testing.TB) (ds.Batching, func()) {
	dataPath, err := ioutil.TempDir(os.TempDir(), "leveldb")
	if err != nil {
		tb.Fatal(err)
	}
	store, err := leveldb.NewDatastore(dataPath, nil)
	if err != nil {
		tb.Fatal(err)
	}
	closer := func() {
		store.Close()
		os.RemoveAll(dataPath)
	}
	return store, closer
}

func peerstoreFactory(tb testing.TB, storeFactory datastoreFactory, opts Options) pt.PeerstoreFactory {
	return func() (pstore.Peerstore, func()) {
		store, storeCloseFn := storeFactory(tb)
		ps, err := NewPeerstore(context.Background(), store, opts)
		if err != nil {
			tb.Fatal(err)
		}
		closer := func() {
			ps.Close()
			storeCloseFn()
		}
		return ps, closer
	}
}

func addressBookFactory(tb testing.TB, storeFactory datastoreFactory, opts Options) pt.AddrBookFactory {
	return func() (pstore.AddrBook, func()) {
		store, closeFunc := storeFactory(tb)
		ab, err := NewAddrBook(context.Background(), store, opts)
		if err != nil {
			tb.Fatal(err)
		}
		closer := func() {
			ab.Close()
			closeFunc()
		}
		return ab, closer
	}
}

func keyBookFactory(tb testing.TB, storeFactory datastoreFactory, opts Options) pt.KeyBookFactory {
	return func() (pstore.KeyBook, func()) {
		store, storeCloseFn := storeFactory(tb)
		kb, err := NewKeyBook(context.Background(), store, opts)
		if err != nil {
			tb.Fatal(err)
		}
		return kb, storeCloseFn
	}
}
<commit_msg>fix: avoid using the disk in datastore tests<commit_after>package pstoreds

import (
	"context"
	"io/ioutil"
	"os"
	"testing"
	"time"

	ds "github.com/ipfs/go-datastore"
	badger "github.com/ipfs/go-ds-badger"
	leveldb "github.com/ipfs/go-ds-leveldb"

	pstore "github.com/libp2p/go-libp2p-core/peerstore"
	pt "github.com/libp2p/go-libp2p-peerstore/test"
)

type datastoreFactory func(tb testing.TB) (ds.Batching, func())

var dstores = map[string]datastoreFactory{
	//"Badger": badgerStore,
	"Leveldb": leveldbStore,
}

func TestDsPeerstore(t *testing.T) {
	for name, dsFactory := range dstores {
		t.Run(name, func(t *testing.T) {
			pt.TestPeerstore(t, peerstoreFactory(t, dsFactory, DefaultOpts()))
		})
	}
}

func TestDsAddrBook(t *testing.T) {
	for name, dsFactory := range dstores {
		t.Run(name+" Cacheful", func(t *testing.T) {
			opts := DefaultOpts()
			opts.GCPurgeInterval = 1 * time.Second
			opts.CacheSize = 1024

			pt.TestAddrBook(t, addressBookFactory(t, dsFactory, opts))
		})

		t.Run(name+" Cacheless", func(t *testing.T) {
			opts := DefaultOpts()
			opts.GCPurgeInterval = 1 * time.Second
			opts.CacheSize = 0

			pt.TestAddrBook(t, addressBookFactory(t, dsFactory, opts))
		})
	}
}

func TestDsKeyBook(t *testing.T) {
	for name, dsFactory := range dstores {
		t.Run(name, func(t *testing.T) {
			pt.TestKeyBook(t, keyBookFactory(t, dsFactory, DefaultOpts()))
		})
	}
}

func BenchmarkDsKeyBook(b *testing.B) {
	for name, dsFactory := range dstores {
		b.Run(name, func(b *testing.B) {
			pt.BenchmarkKeyBook(b, keyBookFactory(b, dsFactory, DefaultOpts()))
		})
	}
}

func BenchmarkDsPeerstore(b *testing.B) {
	caching := DefaultOpts()
	caching.CacheSize = 1024

	cacheless := DefaultOpts()
	cacheless.CacheSize = 0

	for name, dsFactory := range dstores {
		b.Run(name, func(b *testing.B) {
			pt.BenchmarkPeerstore(b, peerstoreFactory(b, dsFactory, caching), "Caching")
		})
		b.Run(name, func(b *testing.B) {
			pt.BenchmarkPeerstore(b, peerstoreFactory(b, dsFactory, cacheless), "Cacheless")
		})
	}
}

// Doesn't work on 32bit because badger.
//lint:ignore U1000 disabled for now
func badgerStore(tb testing.TB) (ds.Batching, func()) {
	dataPath, err := ioutil.TempDir(os.TempDir(), "badger")
	if err != nil {
		tb.Fatal(err)
	}
	store, err := badger.NewDatastore(dataPath, nil)
	if err != nil {
		tb.Fatal(err)
	}
	closer := func() {
		store.Close()
		os.RemoveAll(dataPath)
	}
	return store, closer
}

func leveldbStore(tb testing.TB) (ds.Batching, func()) {
	// Intentionally test in-memory because disks suck, especially in CI.
	store, err := leveldb.NewDatastore("", nil)
	if err != nil {
		tb.Fatal(err)
	}
	closer := func() {
		store.Close()
	}
	return store, closer
}

func peerstoreFactory(tb testing.TB, storeFactory datastoreFactory, opts Options) pt.PeerstoreFactory {
	return func() (pstore.Peerstore, func()) {
		store, storeCloseFn := storeFactory(tb)
		ps, err := NewPeerstore(context.Background(), store, opts)
		if err != nil {
			tb.Fatal(err)
		}
		closer := func() {
			ps.Close()
			storeCloseFn()
		}
		return ps, closer
	}
}

func addressBookFactory(tb testing.TB, storeFactory datastoreFactory, opts Options) pt.AddrBookFactory {
	return func() (pstore.AddrBook, func()) {
		store, closeFunc := storeFactory(tb)
		ab, err := NewAddrBook(context.Background(), store, opts)
		if err != nil {
			tb.Fatal(err)
		}
		closer := func() {
			ab.Close()
			closeFunc()
		}
		return ab, closer
	}
}

func keyBookFactory(tb testing.TB, storeFactory datastoreFactory, opts Options) pt.KeyBookFactory {
	return func() (pstore.KeyBook, func()) {
		store, storeCloseFn := storeFactory(tb)
		kb, err := NewKeyBook(context.Background(), store, opts)
		if err != nil {
			tb.Fatal(err)
		}
		return kb, storeCloseFn
	}
}
<|endoftext|>


<commit_before>package utils

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/naoina/toml"
)

// Config is a project-wide struct that holds configuration information
type Config struct {
	sourceFile string
	Core       struct {
		Debug         bool
		SpecialDebug  bool
		TempDir       string
		ModuleDir     string
		RemoteBaseDir string
		KeepTempFiles bool
		HaltOnError   bool
	}
	SSH struct {
		Username   string
		Password   string
		PrivateKey string
		Timeout    string
	}
	Hosts    []*ConfigHost
	HostsMap map[string]*ConfigHost
}

type ConfigHost struct {
	Name          string            `json:"name"`
	Address       string            `json:"address"`
	Username      string            `json:"username"`
	Password      string            `json:"-"`
	PrivateKey    string            `json:"-"`
	Disable       bool              `json:"disabled"`
	SSHConnection *ssh.Client       `json:"-"`
	SSHConfig     *ssh.ClientConfig `json:"-"`
	appConfig     *Config
}

func (c *ConfigHost) ConnectSSH(clientConfig *ssh.ClientConfig) error {
	if c.SSHConfig == nil {
		// Make a copy for this host
		hostSSHConfig := &ssh.ClientConfig{}
		*hostSSHConfig = *clientConfig

		// Replace the global connection options with host specific ones if given
		if c.Password != "" || c.PrivateKey != "" {
			hostSSHConfig.Auth = make([]ssh.AuthMethod, 0, 1)

			if c.Password != "" {
				hostSSHConfig.Auth = append(hostSSHConfig.Auth, ssh.Password(c.Password))
			}

			if c.PrivateKey != "" {
				sshPrivateKey, err := ioutil.ReadFile(c.PrivateKey)
				if err != nil {
					return err
				}

				signer, err := ssh.ParsePrivateKey(sshPrivateKey)
				if err != nil {
					return err
				}
				hostSSHConfig.Auth = append(hostSSHConfig.Auth, ssh.PublicKeys(signer))
			}
		}

		if c.Username != "" {
			hostSSHConfig.User = c.Username
		}

		c.SSHConfig = hostSSHConfig
	}

	if c.SSHConfig.User == "" || len(c.SSHConfig.Auth) == 0 {
		return fmt.Errorf("Username or password not given for host %s", c.Name)
	}

	var err error
	c.SSHConnection, err = ssh.Dial("tcp", c.Address+":22", c.SSHConfig)
	if err != nil {
		if c.appConfig.Core.Debug {
			return fmt.Errorf("Login failed on %s: %s", c.Name, err.Error())
		}
		return fmt.Errorf("Login failed on %s. Check username or password: %s", c.Name, err)
	}
	return nil
}

// FindConfigFile will locate the a configuration file by looking at the following places
// and choosing the first: SATURN_CONFIG env variable, $PWD/config.toml, $PWD/config/config.toml,
// $HOME/.saturn/config.toml, and /etc/saturn/config.toml.
func FindConfigFile() string {
	if os.Getenv("SATURN_CONFIG") != "" && FileExists(os.Getenv("SATURN_CONFIG")) {
		return os.Getenv("SATURN_CONFIG")
	} else if FileExists("./config.toml") {
		return "./config.toml"
	} else if FileExists("./config/config.toml") {
		return "./config/config.toml"
	} else if FileExists(os.ExpandEnv("$HOME/.saturn/config.toml")) {
		return os.ExpandEnv("$HOME/.saturn/config.toml")
	} else if FileExists("/etc/saturn/config.toml") {
		return "/etc/saturn/config.toml"
	}
	return ""
}

// NewConfig will load a Config object using the given TOML file.
func NewConfig(configFile string) (conf *Config, err error) {
	defer func() {
		if r := recover(); r != nil {
			switch x := r.(type) {
			case string:
				err = errors.New(x)
			case error:
				err = x
			default:
				err = errors.New("Unknown panic")
			}
		}
	}()

	if configFile == "" {
		configFile = "config.toml"
	}

	f, err := os.Open(configFile)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	buf, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	var con Config
	if err := toml.Unmarshal(buf, &con); err != nil {
		return nil, err
	}

	if len(con.Hosts) == 0 {
		return nil, errors.New("No hosts defined")
	}

	con.sourceFile = configFile

	con.HostsMap = make(map[string]*ConfigHost)
	for _, host := range con.Hosts {
		if _, exists := con.HostsMap[host.Name]; exists {
			return nil, fmt.Errorf("Host %s duplicated in configuration", host.Name)
		}
		host.appConfig = &con
		con.HostsMap[host.Name] = host
	}

	return setConfigDefaults(&con)
}

func setConfigDefaults(c *Config) (*Config, error) {
	// Anything not set here implies its zero value is the default

	c.Core.TempDir = setStringOrDefault(c.Core.TempDir, "./tmp")
	c.Core.ModuleDir = setStringOrDefault(c.Core.ModuleDir, "./modules")
	c.Core.RemoteBaseDir = setStringOrDefault(c.Core.RemoteBaseDir, "$HOME")
	c.SSH.Username = setStringOrDefault(c.SSH.Username, "root")
	c.SSH.Timeout = setStringOrDefault(c.SSH.Timeout, "20s")
	if _, err := time.ParseDuration(c.SSH.Timeout); err == nil {
		c.SSH.Timeout = "20s"
	}
	return c, nil
}

// Given string s, if it is empty, return v else return s.
func setStringOrDefault(s, v string) string {
	if s == "" {
		return v
	}
	return s
}

// Given integer s, if it is 0, return v else return s.
func setIntOrDefault(s, v int) int {
	if s == 0 {
		return v
	}
	return s
}

func FileExists(file string) bool {
	_, err := os.Stat(file)
	return !os.IsNotExist(err)
}<commit_msg>Use environment user as default SSH username<commit_after>package utils

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/naoina/toml"
)

// Config is a project-wide struct that holds configuration information
type Config struct {
	sourceFile string
	Core       struct {
		Debug         bool
		SpecialDebug  bool
		TempDir       string
		ModuleDir     string
		RemoteBaseDir string
		KeepTempFiles bool
		HaltOnError   bool
	}
	SSH struct {
		Username   string
		Password   string
		PrivateKey string
		Timeout    string
	}
	Hosts    []*ConfigHost
	HostsMap map[string]*ConfigHost
}

type ConfigHost struct {
	Name          string            `json:"name"`
	Address       string            `json:"address"`
	Username      string            `json:"username"`
	Password      string            `json:"-"`
	PrivateKey    string            `json:"-"`
	Disable       bool              `json:"disabled"`
	SSHConnection *ssh.Client       `json:"-"`
	SSHConfig     *ssh.ClientConfig `json:"-"`
	appConfig     *Config
}

func (c *ConfigHost) ConnectSSH(clientConfig *ssh.ClientConfig) error {
	if c.SSHConfig == nil {
		// Make a copy for this host
		hostSSHConfig := &ssh.ClientConfig{}
		*hostSSHConfig = *clientConfig

		// Replace the global connection options with host specific ones if given
		if c.Password != "" || c.PrivateKey != "" {
			hostSSHConfig.Auth = make([]ssh.AuthMethod, 0, 1)

			if c.Password != "" {
				hostSSHConfig.Auth = append(hostSSHConfig.Auth, ssh.Password(c.Password))
			}

			if c.PrivateKey != "" {
				sshPrivateKey, err := ioutil.ReadFile(c.PrivateKey)
				if err != nil {
					return err
				}

				signer, err := ssh.ParsePrivateKey(sshPrivateKey)
				if err != nil {
					return err
				}
				hostSSHConfig.Auth = append(hostSSHConfig.Auth, ssh.PublicKeys(signer))
			}
		}

		if c.Username != "" {
			hostSSHConfig.User = c.Username
		}

		c.SSHConfig = hostSSHConfig
	}

	if c.SSHConfig.User == "" || len(c.SSHConfig.Auth) == 0 {
		return fmt.Errorf("Username or password not given for host %s", c.Name)
	}

	var err error
	c.SSHConnection, err = ssh.Dial("tcp", c.Address+":22", c.SSHConfig)
	if err != nil {
		if c.appConfig.Core.Debug {
			return fmt.Errorf("Login failed on %s: %s", c.Name, err.Error())
		}
		return fmt.Errorf("Login failed on %s. Check username or password: %s", c.Name, err)
	}
	return nil
}

// FindConfigFile will locate the a configuration file by looking at the following places
// and choosing the first: SATURN_CONFIG env variable, $PWD/config.toml, $PWD/config/config.toml,
// $HOME/.saturn/config.toml, and /etc/saturn/config.toml.
func FindConfigFile() string {
	if os.Getenv("SATURN_CONFIG") != "" && FileExists(os.Getenv("SATURN_CONFIG")) {
		return os.Getenv("SATURN_CONFIG")
	} else if FileExists("./config.toml") {
		return "./config.toml"
	} else if FileExists("./config/config.toml") {
		return "./config/config.toml"
	} else if FileExists(os.ExpandEnv("$HOME/.saturn/config.toml")) {
		return os.ExpandEnv("$HOME/.saturn/config.toml")
	} else if FileExists("/etc/saturn/config.toml") {
		return "/etc/saturn/config.toml"
	}
	return ""
}

// NewConfig will load a Config object using the given TOML file.
func NewConfig(configFile string) (conf *Config, err error) {
	defer func() {
		if r := recover(); r != nil {
			switch x := r.(type) {
			case string:
				err = errors.New(x)
			case error:
				err = x
			default:
				err = errors.New("Unknown panic")
			}
		}
	}()

	if configFile == "" {
		configFile = "config.toml"
	}

	f, err := os.Open(configFile)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	buf, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	var con Config
	if err := toml.Unmarshal(buf, &con); err != nil {
		return nil, err
	}

	if len(con.Hosts) == 0 {
		return nil, errors.New("No hosts defined")
	}

	con.sourceFile = configFile

	con.HostsMap = make(map[string]*ConfigHost)
	for _, host := range con.Hosts {
		if _, exists := con.HostsMap[host.Name]; exists {
			return nil, fmt.Errorf("Host %s duplicated in configuration", host.Name)
		}
		host.appConfig = &con
		con.HostsMap[host.Name] = host
	}

	return setConfigDefaults(&con)
}

func setConfigDefaults(c *Config) (*Config, error) {
	// Anything not set here implies its zero value is the default

	c.Core.TempDir = setStringOrDefault(c.Core.TempDir, "./tmp")
	c.Core.ModuleDir = setStringOrDefault(c.Core.ModuleDir, "./modules")
	c.Core.RemoteBaseDir = setStringOrDefault(c.Core.RemoteBaseDir, "$HOME")
	c.SSH.Username = setStringOrDefault(c.SSH.Username, os.Getenv("USER"))
	c.SSH.Timeout = setStringOrDefault(c.SSH.Timeout, "20s")
	if _, err := time.ParseDuration(c.SSH.Timeout); err == nil {
		c.SSH.Timeout = "20s"
	}
	return c, nil
}

// Given string s, if it is empty, return v else return s.
func setStringOrDefault(s, v string) string {
	if s == "" {
		return v
	}
	return s
}

// Given integer s, if it is 0, return v else return s.
func setIntOrDefault(s, v int) int {
	if s == 0 {
		return v
	}
	return s
}

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