<commit_before>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package common

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

	"github.com/uber/cherami-thrift/.generated/go/controller"

	"github.com/uber-common/bark"
	"github.com/uber/tchannel-go/thrift"
)

var (
	// ErrLoadReportThrottled is the error returned by LoadReporter when it runs out of tokens to serve the request
	ErrLoadReportThrottled = errors.New("Too many load reports from the host.")
)

type (
	// LoadReporterSource is implemented by any component reporting load to controller
	LoadReporterSource interface {
		Report(reporter LoadReporter)
	}

	// LoadReporterDaemon is used for periodically reporting load to controller
	LoadReporterDaemon interface {
		Daemon
	}

	// LoadReporterDaemonFactory is used to create a daemon task for reporting load to controller
	LoadReporterDaemonFactory interface {
		CreateReporter(interval time.Duration, source LoadReporterSource, logger bark.Logger) LoadReporterDaemon
	}

	// LoadReporter is used by each component interested in reporting load to Controller
	LoadReporter interface {
		ReportHostMetric(metrics controller.NodeMetrics) error
		ReportDestinationMetric(destinationUUID string, metrics controller.DestinationMetrics) error
		ReportDestinationExtentMetric(destinationUUID string, extentUUID string, metrics controller.DestinationExtentMetrics) error
		ReportConsumerGroupMetric(destinationUUID string, consumerGroupUUID string, metrics controller.ConsumerGroupMetrics) error
		ReportConsumerGroupExtentMetric(destinationUUID string, consumerGroupUUID string, extentUUID string, metrics controller.ConsumerGroupExtentMetrics) error
		ReportStoreExtentMetric(extentUUID string, metrics controller.StoreExtentMetrics) error
	}

	loadReporterDaemonFactoryImpl struct {
		reporter      LoadReporter
		tickerFactory TickerSourceFactory
		logger        bark.Logger
	}

	loadReporterDaemonImpl struct {
		interval      time.Duration
		tickerFactory TickerSourceFactory
		source        LoadReporterSource
		reporter      LoadReporter
		started       int32
		shutdownCh    chan struct{}
		shutdownWG    sync.WaitGroup
		logger        bark.Logger
	}

	loadReporterImpl struct {
		hostUUID         string
		sku              controller.SKU
		role             controller.Role
		requestThrottler TokenBucket
		clientFactory    ClientFactory
		logger           bark.Logger
	}
)

// ReportExtentMetric rate limiting: throttle requests beyond 1000 every second
const (
	reportLoadMetricThrottleRequests = 1000
	reportLoadMetricThriftTimeout    = 2 * time.Second
)

// NewLoadReporterDaemonFactory is used to create a factory for creating LoadReporterDaemon
func NewLoadReporterDaemonFactory(hostUUID string, sku controller.SKU, role controller.Role,
	clientFactory ClientFactory, tickerFactory TickerSourceFactory, tokenBucketFactory TokenBucketFactory, logger bark.Logger) LoadReporterDaemonFactory {
	r := &loadReporterImpl{
		hostUUID:         hostUUID,
		sku:              sku,
		role:             role,
		clientFactory:    clientFactory,
		requestThrottler: tokenBucketFactory.CreateTokenBucket(reportLoadMetricThrottleRequests, NewRealTimeSource()),
		logger:           logger,
	}

	return &loadReporterDaemonFactoryImpl{
		reporter:      r,
		tickerFactory: tickerFactory,
		logger:        logger,
	}
}

// CreateReporter creates a daemon task for reporting load metric
func (f *loadReporterDaemonFactoryImpl) CreateReporter(interval time.Duration, source LoadReporterSource, logger bark.Logger) LoadReporterDaemon {
	return &loadReporterDaemonImpl{
		interval:      interval,
		tickerFactory: f.tickerFactory,
		source:        source,
		reporter:      f.reporter,
		shutdownCh:    make(chan struct{}),
		logger:        logger,
	}
}

// Start is used to start the daemon task to start reporting metric at regular interval
func (d *loadReporterDaemonImpl) Start() {
	if !atomic.CompareAndSwapInt32(&d.started, 0, 1) {
		return
	}

	d.shutdownWG.Add(1)
	go d.reporterPump()

	d.logger.Info("Load reporter started.")
}

// Stop is used to stop the daemon task from reporting metric at regular interval
func (d *loadReporterDaemonImpl) Stop() {
	if !atomic.CompareAndSwapInt32(&d.started, 1, 0) {
		return
	}

	close(d.shutdownCh)
	if success := AwaitWaitGroup(&d.shutdownWG, time.Minute); !success {
		d.logger.Warn("Timed out waiting to shutdown load reporter.")
	}

	d.logger.Info("Load reporter stopped.")
}

func (d *loadReporterDaemonImpl) reporterPump() {
	defer d.shutdownWG.Done()
	reportTicker := d.tickerFactory.CreateTicker(d.interval)
	defer reportTicker.Stop()

	for {
		select {
		case <-reportTicker.Ticks():
			d.source.Report(d.reporter)
		case <-d.shutdownCh:
			d.logger.Info("Load reporter pump shutting down.")
			return
		}
	}
}

// ReportHostMetric is the API exposed by LoadReporter for reporting host load to controller
func (d *loadReporterImpl) ReportHostMetric(metrics controller.NodeMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportNodeMetricRequest()
	request.HostId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.Role = RolePtr(d.role)
	request.Metrics = NodeMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportNodeMetric(ctx, request)
}

// ReportDestinationMetric is the API exposed by LoadReporter for reporting destination load to controller
func (d *loadReporterImpl) ReportDestinationMetric(destinationUUID string, metrics controller.DestinationMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportDestinationMetricRequest()
	request.HostId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.DestinationUUID = StringPtr(destinationUUID)
	request.Metrics = DestinationMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportDestinationMetric(ctx, request)
}

// ReportDestinationExtentMetric is the API exposed by LoadReporter for reporting destination extent load to controller
func (d *loadReporterImpl) ReportDestinationExtentMetric(destinationUUID string, extentUUID string, metrics controller.DestinationExtentMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportDestinationExtentMetricRequest()
	request.HostId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.DestinationUUID = StringPtr(destinationUUID)
	request.ExtentUUID = StringPtr(extentUUID)
	request.Metrics = DestinationExtentMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportDestinationExtentMetric(ctx, request)
}

// ReportConsumerGroupMetric is the API exposed by LoadReporter for reporting consumer group load to controller
func (d *loadReporterImpl) ReportConsumerGroupMetric(destinationUUID string, consumerGroupUUID string, metrics controller.ConsumerGroupMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportConsumerGroupMetricRequest()
	request.HostId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.DestinationUUID = StringPtr(destinationUUID)
	request.ConsumerGroupUUID = StringPtr(consumerGroupUUID)
	request.Metrics = ConsumerGroupMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportConsumerGroupMetric(ctx, request)
}

// ReportConsumerGroupExtentMetric is the API exposed by LoadReporter for reporting consumer group extent load to controller
func (d *loadReporterImpl) ReportConsumerGroupExtentMetric(destinationUUID string, consumerGroupUUID string, extentUUID string, metrics controller.ConsumerGroupExtentMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportConsumerGroupExtentMetricRequest()
	request.HostId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.DestinationUUID = StringPtr(destinationUUID)
	request.ConsumerGroupUUID = StringPtr(consumerGroupUUID)
	request.ExtentUUID = StringPtr(extentUUID)
	request.Metrics = ConsumerGroupExtentMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportConsumerGroupExtentMetric(ctx, request)
}

// ReportStoreExtentMetric is the API exposed by LoadReporter for reporting store extent load to controller
func (d *loadReporterImpl) ReportStoreExtentMetric(extentUUID string, metrics controller.StoreExtentMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportStoreExtentMetricRequest()
	request.StoreId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.ExtentUUID = StringPtr(extentUUID)
	request.Metrics = StoreExtentMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportStoreExtentMetric(ctx, request)
}

func (d *loadReporterImpl) getControllerClient() (controller.TChanController, error) {
	if ok, _ := d.requestThrottler.TryConsume(1); !ok {
		return nil, ErrLoadReportThrottled
	}

	return d.clientFactory.GetControllerClient()
}
<commit_msg>Send atleast one report per reporter session (#70)<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package common

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

	"github.com/uber/cherami-thrift/.generated/go/controller"

	"github.com/uber-common/bark"
	"github.com/uber/tchannel-go/thrift"
)

var (
	// ErrLoadReportThrottled is the error returned by LoadReporter when it runs out of tokens to serve the request
	ErrLoadReportThrottled = errors.New("Too many load reports from the host.")
)

type (
	// LoadReporterSource is implemented by any component reporting load to controller
	LoadReporterSource interface {
		Report(reporter LoadReporter)
	}

	// LoadReporterDaemon is used for periodically reporting load to controller
	LoadReporterDaemon interface {
		Daemon
	}

	// LoadReporterDaemonFactory is used to create a daemon task for reporting load to controller
	LoadReporterDaemonFactory interface {
		CreateReporter(interval time.Duration, source LoadReporterSource, logger bark.Logger) LoadReporterDaemon
	}

	// LoadReporter is used by each component interested in reporting load to Controller
	LoadReporter interface {
		ReportHostMetric(metrics controller.NodeMetrics) error
		ReportDestinationMetric(destinationUUID string, metrics controller.DestinationMetrics) error
		ReportDestinationExtentMetric(destinationUUID string, extentUUID string, metrics controller.DestinationExtentMetrics) error
		ReportConsumerGroupMetric(destinationUUID string, consumerGroupUUID string, metrics controller.ConsumerGroupMetrics) error
		ReportConsumerGroupExtentMetric(destinationUUID string, consumerGroupUUID string, extentUUID string, metrics controller.ConsumerGroupExtentMetrics) error
		ReportStoreExtentMetric(extentUUID string, metrics controller.StoreExtentMetrics) error
	}

	loadReporterDaemonFactoryImpl struct {
		reporter      LoadReporter
		tickerFactory TickerSourceFactory
		logger        bark.Logger
	}

	loadReporterDaemonImpl struct {
		interval      time.Duration
		tickerFactory TickerSourceFactory
		source        LoadReporterSource
		reporter      LoadReporter
		started       int32
		reported      bool
		shutdownCh    chan struct{}
		shutdownWG    sync.WaitGroup
		logger        bark.Logger
	}

	loadReporterImpl struct {
		hostUUID         string
		sku              controller.SKU
		role             controller.Role
		requestThrottler TokenBucket
		clientFactory    ClientFactory
		logger           bark.Logger
	}
)

// ReportExtentMetric rate limiting: throttle requests beyond 1000 every second
const (
	reportLoadMetricThrottleRequests = 1000
	reportLoadMetricThriftTimeout    = 2 * time.Second
)

// NewLoadReporterDaemonFactory is used to create a factory for creating LoadReporterDaemon
func NewLoadReporterDaemonFactory(hostUUID string, sku controller.SKU, role controller.Role,
	clientFactory ClientFactory, tickerFactory TickerSourceFactory, tokenBucketFactory TokenBucketFactory, logger bark.Logger) LoadReporterDaemonFactory {
	r := &loadReporterImpl{
		hostUUID:         hostUUID,
		sku:              sku,
		role:             role,
		clientFactory:    clientFactory,
		requestThrottler: tokenBucketFactory.CreateTokenBucket(reportLoadMetricThrottleRequests, NewRealTimeSource()),
		logger:           logger,
	}

	return &loadReporterDaemonFactoryImpl{
		reporter:      r,
		tickerFactory: tickerFactory,
		logger:        logger,
	}
}

// CreateReporter creates a daemon task for reporting load metric
func (f *loadReporterDaemonFactoryImpl) CreateReporter(interval time.Duration, source LoadReporterSource, logger bark.Logger) LoadReporterDaemon {
	return &loadReporterDaemonImpl{
		interval:      interval,
		tickerFactory: f.tickerFactory,
		source:        source,
		reporter:      f.reporter,
		shutdownCh:    make(chan struct{}),
		logger:        logger,
	}
}

// Start is used to start the daemon task to start reporting metric at regular interval
func (d *loadReporterDaemonImpl) Start() {
	if !atomic.CompareAndSwapInt32(&d.started, 0, 1) {
		return
	}

	d.shutdownWG.Add(1)
	go d.reporterPump()

	d.logger.Info("Load reporter started.")
}

// Stop is used to stop the daemon task from reporting metric at regular interval
func (d *loadReporterDaemonImpl) Stop() {
	if !atomic.CompareAndSwapInt32(&d.started, 1, 0) {
		return
	}

	close(d.shutdownCh)
	if success := AwaitWaitGroup(&d.shutdownWG, time.Minute); !success {
		d.logger.Warn("Timed out waiting to shutdown load reporter.")
	}

	d.logger.Info("Load reporter stopped.")
}

func (d *loadReporterDaemonImpl) reporterPump() {

	defer d.shutdownWG.Done()

	reportTicker := d.tickerFactory.CreateTicker(d.interval)
	defer reportTicker.Stop()

	for {
		select {
		case <-reportTicker.Ticks():

			d.source.Report(d.reporter)
			d.reported = true

		case <-d.shutdownCh:

			// ensure we report at least once
			if !d.reported {
				d.source.Report(d.reporter)
				d.reported = true
			}

			d.logger.Info("Load reporter pump shutting down.")
			return
		}
	}
}

// ReportHostMetric is the API exposed by LoadReporter for reporting host load to controller
func (d *loadReporterImpl) ReportHostMetric(metrics controller.NodeMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportNodeMetricRequest()
	request.HostId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.Role = RolePtr(d.role)
	request.Metrics = NodeMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportNodeMetric(ctx, request)
}

// ReportDestinationMetric is the API exposed by LoadReporter for reporting destination load to controller
func (d *loadReporterImpl) ReportDestinationMetric(destinationUUID string, metrics controller.DestinationMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportDestinationMetricRequest()
	request.HostId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.DestinationUUID = StringPtr(destinationUUID)
	request.Metrics = DestinationMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportDestinationMetric(ctx, request)
}

// ReportDestinationExtentMetric is the API exposed by LoadReporter for reporting destination extent load to controller
func (d *loadReporterImpl) ReportDestinationExtentMetric(destinationUUID string, extentUUID string, metrics controller.DestinationExtentMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportDestinationExtentMetricRequest()
	request.HostId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.DestinationUUID = StringPtr(destinationUUID)
	request.ExtentUUID = StringPtr(extentUUID)
	request.Metrics = DestinationExtentMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportDestinationExtentMetric(ctx, request)
}

// ReportConsumerGroupMetric is the API exposed by LoadReporter for reporting consumer group load to controller
func (d *loadReporterImpl) ReportConsumerGroupMetric(destinationUUID string, consumerGroupUUID string, metrics controller.ConsumerGroupMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportConsumerGroupMetricRequest()
	request.HostId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.DestinationUUID = StringPtr(destinationUUID)
	request.ConsumerGroupUUID = StringPtr(consumerGroupUUID)
	request.Metrics = ConsumerGroupMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportConsumerGroupMetric(ctx, request)
}

// ReportConsumerGroupExtentMetric is the API exposed by LoadReporter for reporting consumer group extent load to controller
func (d *loadReporterImpl) ReportConsumerGroupExtentMetric(destinationUUID string, consumerGroupUUID string, extentUUID string, metrics controller.ConsumerGroupExtentMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportConsumerGroupExtentMetricRequest()
	request.HostId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.DestinationUUID = StringPtr(destinationUUID)
	request.ConsumerGroupUUID = StringPtr(consumerGroupUUID)
	request.ExtentUUID = StringPtr(extentUUID)
	request.Metrics = ConsumerGroupExtentMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportConsumerGroupExtentMetric(ctx, request)
}

// ReportStoreExtentMetric is the API exposed by LoadReporter for reporting store extent load to controller
func (d *loadReporterImpl) ReportStoreExtentMetric(extentUUID string, metrics controller.StoreExtentMetrics) error {
	controllerClient, err := d.getControllerClient()
	if err != nil {
		return err
	}

	request := controller.NewReportStoreExtentMetricRequest()
	request.StoreId = StringPtr(d.hostUUID)
	request.Sku = SKUPtr(d.sku)
	request.ExtentUUID = StringPtr(extentUUID)
	request.Metrics = StoreExtentMetricsPtr(metrics)
	request.Timestamp = Int64Ptr(time.Now().UnixNano())

	ctx, cancel := thrift.NewContext(reportLoadMetricThriftTimeout)
	defer cancel()

	return controllerClient.ReportStoreExtentMetric(ctx, request)
}

func (d *loadReporterImpl) getControllerClient() (controller.TChanController, error) {
	if ok, _ := d.requestThrottler.TryConsume(1); !ok {
		return nil, ErrLoadReportThrottled
	}

	return d.clientFactory.GetControllerClient()
}
<|endoftext|>


<commit_before>package chat

import (
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/keybase/client/go/chat/s3"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/chat1"
	"golang.org/x/net/context"
)

const MB = 1024 * 1024

func TestSignEncrypter(t *testing.T) {
	e := NewSignEncrypter()
	el := e.EncryptedLen(100)
	if el != 180 {
		t.Errorf("enc len: %d, expected 180", el)
	}

	el = e.EncryptedLen(50 * 1024 * 1024)
	if el != 52432880 {
		t.Errorf("enc len: %d, expected 52432880", el)
	}

	pt := "plain text"
	er, err := e.Encrypt(strings.NewReader(pt))
	if err != nil {
		t.Fatal(err)
	}
	ct, err := ioutil.ReadAll(er)
	if err != nil {
		t.Fatal(err)
	}

	if string(ct) == pt {
		t.Fatal("Encrypt did not change plaintext")
	}

	d := NewSignDecrypter()
	dr := d.Decrypt(bytes.NewReader(ct), e.EncryptKey(), e.VerifyKey())
	ptOut, err := ioutil.ReadAll(dr)
	if err != nil {
		t.Fatal(err)
	}
	if string(ptOut) != pt {
		t.Errorf("decrypted ciphertext doesn't match plaintext: %q, expected %q", ptOut, pt)
	}

	// reuse e to do another Encrypt, make sure keys change:
	firstEncKey := e.EncryptKey()
	firstVerifyKey := e.VerifyKey()

	er2, err := e.Encrypt(strings.NewReader(pt))
	if err != nil {
		t.Fatal(err)
	}
	ct2, err := ioutil.ReadAll(er2)
	if err != nil {
		t.Fatal(err)
	}

	if string(ct2) == pt {
		t.Fatal("Encrypt did not change plaintext")
	}
	if bytes.Equal(ct, ct2) {
		t.Fatal("second Encrypt result same as first")
	}
	if bytes.Equal(firstEncKey, e.EncryptKey()) {
		t.Fatal("first enc key reused")
	}
	if bytes.Equal(firstVerifyKey, e.VerifyKey()) {
		t.Fatal("first verify key reused")
	}

	dr2 := d.Decrypt(bytes.NewReader(ct2), e.EncryptKey(), e.VerifyKey())
	ptOut2, err := ioutil.ReadAll(dr2)
	if err != nil {
		t.Fatal(err)
	}
	if string(ptOut2) != pt {
		t.Errorf("decrypted ciphertext doesn't match plaintext: %q, expected %q", ptOut2, pt)
	}
}

func makeTestStore(t *testing.T) *AttachmentStore {
	s := NewAttachmentStore(logger.NewTestLogger(t))

	// use in-memory implementation of s3 interface
	s.s3c = &s3.Mem{}

	return s
}

func randBytes(t *testing.T, n int) []byte {
	buf := make([]byte, n)
	if _, err := rand.Read(buf); err != nil {
		t.Fatal(err)
	}
	return buf
}

func randString(t *testing.T, nbytes int) string {
	return hex.EncodeToString(randBytes(t, nbytes))
}

type ptsigner struct{}

func (p *ptsigner) Sign(payload []byte) ([]byte, error) {
	s := sha256.Sum256(payload)
	return s[:], nil
}

func makeUploadTask(t *testing.T, size int) (plaintext []byte, task *UploadTask) {
	plaintext = randBytes(t, size)
	task = &UploadTask{
		S3Params: chat1.S3Params{
			Bucket:    "upload-test",
			ObjectKey: randString(t, 8),
		},
		LocalSrc: chat1.LocalSource{
			Filename: randString(t, 8),
			Size:     size,
		},
		Plaintext:      bytes.NewReader(plaintext),
		PlaintextHash:  randBytes(t, 16),
		S3Signer:       &ptsigner{},
		ConversationID: randBytes(t, 16),
	}
	return plaintext, task
}

func TestUploadAssetSmall(t *testing.T) {
	s := makeTestStore(t)
	ctx := context.Background()
	plaintext, task := makeUploadTask(t, 1*MB)
	_ = plaintext
	a, err := s.UploadAsset(ctx, task)
	if err != nil {
		t.Fatal(err)
	}

	var buf bytes.Buffer
	if err = s.DownloadAsset(ctx, task.S3Params, a, &buf, task.S3Signer, nil); err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(plaintext, buf.Bytes()) {
		t.Errorf("downloaded asset did not match uploaded asset")
	}
}

func TestUploadAssetLarge(t *testing.T) {
	s := makeTestStore(t)
	ctx := context.Background()
	plaintext, task := makeUploadTask(t, 12*MB)
	_ = plaintext
	a, err := s.UploadAsset(ctx, task)
	if err != nil {
		t.Fatal(err)
	}

	var buf bytes.Buffer
	if err = s.DownloadAsset(ctx, task.S3Params, a, &buf, task.S3Signer, nil); err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(plaintext, buf.Bytes()) {
		t.Errorf("downloaded asset did not match uploaded asset")
	}
}
<commit_msg>Test resume<commit_after>package chat

import (
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/keybase/client/go/chat/s3"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/chat1"
	"golang.org/x/net/context"
)

const MB = 1024 * 1024

func TestSignEncrypter(t *testing.T) {
	e := NewSignEncrypter()
	el := e.EncryptedLen(100)
	if el != 180 {
		t.Errorf("enc len: %d, expected 180", el)
	}

	el = e.EncryptedLen(50 * 1024 * 1024)
	if el != 52432880 {
		t.Errorf("enc len: %d, expected 52432880", el)
	}

	pt := "plain text"
	er, err := e.Encrypt(strings.NewReader(pt))
	if err != nil {
		t.Fatal(err)
	}
	ct, err := ioutil.ReadAll(er)
	if err != nil {
		t.Fatal(err)
	}

	if string(ct) == pt {
		t.Fatal("Encrypt did not change plaintext")
	}

	d := NewSignDecrypter()
	dr := d.Decrypt(bytes.NewReader(ct), e.EncryptKey(), e.VerifyKey())
	ptOut, err := ioutil.ReadAll(dr)
	if err != nil {
		t.Fatal(err)
	}
	if string(ptOut) != pt {
		t.Errorf("decrypted ciphertext doesn't match plaintext: %q, expected %q", ptOut, pt)
	}

	// reuse e to do another Encrypt, make sure keys change:
	firstEncKey := e.EncryptKey()
	firstVerifyKey := e.VerifyKey()

	er2, err := e.Encrypt(strings.NewReader(pt))
	if err != nil {
		t.Fatal(err)
	}
	ct2, err := ioutil.ReadAll(er2)
	if err != nil {
		t.Fatal(err)
	}

	if string(ct2) == pt {
		t.Fatal("Encrypt did not change plaintext")
	}
	if bytes.Equal(ct, ct2) {
		t.Fatal("second Encrypt result same as first")
	}
	if bytes.Equal(firstEncKey, e.EncryptKey()) {
		t.Fatal("first enc key reused")
	}
	if bytes.Equal(firstVerifyKey, e.VerifyKey()) {
		t.Fatal("first verify key reused")
	}

	dr2 := d.Decrypt(bytes.NewReader(ct2), e.EncryptKey(), e.VerifyKey())
	ptOut2, err := ioutil.ReadAll(dr2)
	if err != nil {
		t.Fatal(err)
	}
	if string(ptOut2) != pt {
		t.Errorf("decrypted ciphertext doesn't match plaintext: %q, expected %q", ptOut2, pt)
	}
}

func makeTestStore(t *testing.T) *AttachmentStore {
	s := NewAttachmentStore(logger.NewTestLogger(t))

	// use in-memory implementation of s3 interface
	s.s3c = &s3.Mem{}

	return s
}

func randBytes(t *testing.T, n int) []byte {
	buf := make([]byte, n)
	if _, err := rand.Read(buf); err != nil {
		t.Fatal(err)
	}
	return buf
}

func randString(t *testing.T, nbytes int) string {
	return hex.EncodeToString(randBytes(t, nbytes))
}

type ptsigner struct{}

func (p *ptsigner) Sign(payload []byte) ([]byte, error) {
	s := sha256.Sum256(payload)
	return s[:], nil
}

func makeUploadTask(t *testing.T, size int) (plaintext []byte, task *UploadTask) {
	plaintext = randBytes(t, size)
	task = &UploadTask{
		S3Params: chat1.S3Params{
			Bucket:    "upload-test",
			ObjectKey: randString(t, 8),
		},
		LocalSrc: chat1.LocalSource{
			Filename: randString(t, 8),
			Size:     size,
		},
		Plaintext:      bytes.NewReader(plaintext),
		PlaintextHash:  randBytes(t, 16),
		S3Signer:       &ptsigner{},
		ConversationID: randBytes(t, 16),
	}
	return plaintext, task
}

func TestUploadAssetSmall(t *testing.T) {
	s := makeTestStore(t)
	ctx := context.Background()
	plaintext, task := makeUploadTask(t, 1*MB)
	a, err := s.UploadAsset(ctx, task)
	if err != nil {
		t.Fatal(err)
	}

	var buf bytes.Buffer
	if err = s.DownloadAsset(ctx, task.S3Params, a, &buf, task.S3Signer, nil); err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(plaintext, buf.Bytes()) {
		t.Errorf("downloaded asset did not match uploaded asset")
	}
}

func TestUploadAssetLarge(t *testing.T) {
	s := makeTestStore(t)
	ctx := context.Background()
	plaintext, task := makeUploadTask(t, 12*MB)
	a, err := s.UploadAsset(ctx, task)
	if err != nil {
		t.Fatal(err)
	}

	var buf bytes.Buffer
	if err = s.DownloadAsset(ctx, task.S3Params, a, &buf, task.S3Signer, nil); err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(plaintext, buf.Bytes()) {
		t.Errorf("downloaded asset did not match uploaded asset")
	}
}

type pwcancel struct {
	cancel func()
	after  int
}

func (p *pwcancel) report(bytesCompleted, bytesTotal int) {
	if bytesCompleted > p.after {
		p.cancel()
	}
}

func TestUploadAssetResume(t *testing.T) {
	s := makeTestStore(t)
	ctx, cancel := context.WithCancel(context.Background())
	plaintext, task := makeUploadTask(t, 12*MB)
	pw := &pwcancel{
		cancel: cancel,
		after:  7 * MB,
	}
	task.Progress = pw.report
	_, err := s.UploadAsset(ctx, task)
	if err == nil {
		t.Fatal("expected upload to be canceled")
	}

	ctx = context.Background()
	task.Progress = nil
	a, err := s.UploadAsset(ctx, task)
	if err != nil {
		t.Fatalf("expected second UploadAsset call to work, got: %s", err)
	}

	var buf bytes.Buffer
	if err = s.DownloadAsset(ctx, task.S3Params, a, &buf, task.S3Signer, nil); err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(plaintext, buf.Bytes()) {
		t.Errorf("downloaded asset did not match uploaded asset")
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package service

import (
	"encoding/json"
	"fmt"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/gregor"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/teams"
)

const teamHandlerName = "teamHandler"

type teamHandler struct {
	libkb.Contextified
}

var _ libkb.GregorInBandMessageHandler = (*teamHandler)(nil)

func newTeamHandler(g *libkb.GlobalContext) *teamHandler {
	return &teamHandler{
		Contextified: libkb.NewContextified(g),
	}
}

func (r *teamHandler) Create(ctx context.Context, cli gregor1.IncomingInterface, category string, item gregor.Item) (bool, error) {
	switch category {
	case "team.clkr":
		return true, r.rotateTeam(ctx, cli, item)
	case "team.sbs":
		return true, r.sharingBeforeSignup(ctx, cli, item)
	case "team.openreq":
		return true, r.openTeamAccessRequest(ctx, item)
	case "team.change":
		return true, r.changeTeam(ctx, cli, item, keybase1.TeamChangeSet{})
	case "team.rename":
		return true, r.changeTeam(ctx, cli, item, keybase1.TeamChangeSet{Renamed: true})
	case "team.delete":
		return true, r.deleteTeam(ctx, cli, item)
	case "team.exit":
		return true, r.exitTeam(ctx, cli, item)
	default:
		return false, fmt.Errorf("unknown teamHandler category: %q", category)
	}
}

func (r *teamHandler) rotateTeam(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item) error {
	r.G().Log.Debug("team.clkr received")
	var msg keybase1.TeamCLKRMsg
	if err := json.Unmarshal(item.Body().Bytes(), &msg); err != nil {
		r.G().Log.Debug("error unmarshaling team.clkr item: %s", err)
		return err
	}
	r.G().Log.Debug("team.clkr unmarshaled: %+v", msg)

	if err := teams.HandleRotateRequest(ctx, r.G(), msg.TeamID, keybase1.PerTeamKeyGeneration(msg.Generation)); err != nil {
		return err
	}

	r.G().Log.Debug("dismissing team.clkr item since rotate succeeded")
	return r.G().GregorDismisser.DismissItem(cli, item.Metadata().MsgID())
}

func (r *teamHandler) changeTeam(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item, changes keybase1.TeamChangeSet) error {
	var rows []keybase1.TeamChangeRow
	if err := json.Unmarshal(item.Body().Bytes(), &rows); err != nil {
		r.G().Log.Debug("error unmarshaling team.(change|rename) item: %s", err)
		return err
	}
	r.G().Log.Debug("team.(change|rename) unmarshaled: %+v", rows)

	return teams.HandleChangeNotification(ctx, r.G(), rows, changes)
}

func (r *teamHandler) deleteTeam(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item) error {
	var rows []keybase1.TeamChangeRow
	if err := json.Unmarshal(item.Body().Bytes(), &rows); err != nil {
		r.G().Log.Debug("error unmarshaling team.(change|rename) item: %s", err)
		return err
	}
	r.G().Log.Debug("team.delete unmarshaled: %+v", rows)
	return teams.HandleDeleteNotification(ctx, r.G(), rows)
}

func (r *teamHandler) exitTeam(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item) error {
	var rows []keybase1.TeamExitRow
	if err := json.Unmarshal(item.Body().Bytes(), &rows); err != nil {
		r.G().Log.Debug("error unmarshaling team.exit item: %s", err)
		return err
	}
	r.G().Log.Debug("team.exit unmarshaled: %+v", rows)
	err := teams.HandleExitNotification(ctx, r.G(), rows)
	if err != nil {
		return err
	}

	r.G().Log.Debug("dismissing team.exit: %v", item.Metadata().MsgID().String())
	return r.G().GregorDismisser.DismissItem(cli, item.Metadata().MsgID())
}

func (r *teamHandler) sharingBeforeSignup(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item) error {
	r.G().Log.Debug("team.sbs received")
	var msg keybase1.TeamSBSMsg
	if err := json.Unmarshal(item.Body().Bytes(), &msg); err != nil {
		r.G().Log.Debug("error unmarshaling team.sbs item: %s", err)
		return err
	}
	r.G().Log.Debug("team.sbs unmarshaled: %+v", msg)

	if err := teams.HandleSBSRequest(ctx, r.G(), msg); err != nil {
		return err
	}

	r.G().Log.Debug("dismissing team.sbs item since it succeeded")
	return r.G().GregorDismisser.DismissItem(cli, item.Metadata().MsgID())
}

func (r *teamHandler) openTeamAccessRequest(ctx context.Context, item gregor.Item) error {
	r.G().Log.Debug("team.openreq received")
	var msg keybase1.TeamOpenReqMsg
	if err := json.Unmarshal(item.Body().Bytes(), &msg); err != nil {
		r.G().Log.Debug("error unmarshaling team.openreq item: %s", err)
		return err
	}
	r.G().Log.Debug("team.openreq unmarshaled: %+v", msg)

	if err := teams.HandleOpenTeamAccessRequest(ctx, r.G(), msg); err != nil {
		return err
	}

	r.G().Log.Debug("dismissing team.openreq item since it succeeded")
	return r.G().GregorDismisser.DismissItem(item.Metadata().MsgID())
}

func (r *teamHandler) Dismiss(ctx context.Context, cli gregor1.IncomingInterface, category string, item gregor.Item) (bool, error) {
	return false, nil
}

func (r *teamHandler) IsAlive() bool {
	return true
}

func (r *teamHandler) Name() string {
	return teamHandlerName
}
<commit_msg>Fix gregor call after rebase<commit_after>// Copyright 2017 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package service

import (
	"encoding/json"
	"fmt"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/gregor"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/teams"
)

const teamHandlerName = "teamHandler"

type teamHandler struct {
	libkb.Contextified
}

var _ libkb.GregorInBandMessageHandler = (*teamHandler)(nil)

func newTeamHandler(g *libkb.GlobalContext) *teamHandler {
	return &teamHandler{
		Contextified: libkb.NewContextified(g),
	}
}

func (r *teamHandler) Create(ctx context.Context, cli gregor1.IncomingInterface, category string, item gregor.Item) (bool, error) {
	switch category {
	case "team.clkr":
		return true, r.rotateTeam(ctx, cli, item)
	case "team.sbs":
		return true, r.sharingBeforeSignup(ctx, cli, item)
	case "team.openreq":
		return true, r.openTeamAccessRequest(ctx, cli, item)
	case "team.change":
		return true, r.changeTeam(ctx, cli, item, keybase1.TeamChangeSet{})
	case "team.rename":
		return true, r.changeTeam(ctx, cli, item, keybase1.TeamChangeSet{Renamed: true})
	case "team.delete":
		return true, r.deleteTeam(ctx, cli, item)
	case "team.exit":
		return true, r.exitTeam(ctx, cli, item)
	default:
		return false, fmt.Errorf("unknown teamHandler category: %q", category)
	}
}

func (r *teamHandler) rotateTeam(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item) error {
	r.G().Log.Debug("team.clkr received")
	var msg keybase1.TeamCLKRMsg
	if err := json.Unmarshal(item.Body().Bytes(), &msg); err != nil {
		r.G().Log.Debug("error unmarshaling team.clkr item: %s", err)
		return err
	}
	r.G().Log.Debug("team.clkr unmarshaled: %+v", msg)

	if err := teams.HandleRotateRequest(ctx, r.G(), msg.TeamID, keybase1.PerTeamKeyGeneration(msg.Generation)); err != nil {
		return err
	}

	r.G().Log.Debug("dismissing team.clkr item since rotate succeeded")
	return r.G().GregorDismisser.DismissItem(cli, item.Metadata().MsgID())
}

func (r *teamHandler) changeTeam(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item, changes keybase1.TeamChangeSet) error {
	var rows []keybase1.TeamChangeRow
	if err := json.Unmarshal(item.Body().Bytes(), &rows); err != nil {
		r.G().Log.Debug("error unmarshaling team.(change|rename) item: %s", err)
		return err
	}
	r.G().Log.Debug("team.(change|rename) unmarshaled: %+v", rows)

	return teams.HandleChangeNotification(ctx, r.G(), rows, changes)
}

func (r *teamHandler) deleteTeam(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item) error {
	var rows []keybase1.TeamChangeRow
	if err := json.Unmarshal(item.Body().Bytes(), &rows); err != nil {
		r.G().Log.Debug("error unmarshaling team.(change|rename) item: %s", err)
		return err
	}
	r.G().Log.Debug("team.delete unmarshaled: %+v", rows)
	return teams.HandleDeleteNotification(ctx, r.G(), rows)
}

func (r *teamHandler) exitTeam(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item) error {
	var rows []keybase1.TeamExitRow
	if err := json.Unmarshal(item.Body().Bytes(), &rows); err != nil {
		r.G().Log.Debug("error unmarshaling team.exit item: %s", err)
		return err
	}
	r.G().Log.Debug("team.exit unmarshaled: %+v", rows)
	err := teams.HandleExitNotification(ctx, r.G(), rows)
	if err != nil {
		return err
	}

	r.G().Log.Debug("dismissing team.exit: %v", item.Metadata().MsgID().String())
	return r.G().GregorDismisser.DismissItem(cli, item.Metadata().MsgID())
}

func (r *teamHandler) sharingBeforeSignup(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item) error {
	r.G().Log.Debug("team.sbs received")
	var msg keybase1.TeamSBSMsg
	if err := json.Unmarshal(item.Body().Bytes(), &msg); err != nil {
		r.G().Log.Debug("error unmarshaling team.sbs item: %s", err)
		return err
	}
	r.G().Log.Debug("team.sbs unmarshaled: %+v", msg)

	if err := teams.HandleSBSRequest(ctx, r.G(), msg); err != nil {
		return err
	}

	r.G().Log.Debug("dismissing team.sbs item since it succeeded")
	return r.G().GregorDismisser.DismissItem(cli, item.Metadata().MsgID())
}

func (r *teamHandler) openTeamAccessRequest(ctx context.Context, cli gregor1.IncomingInterface, item gregor.Item) error {
	r.G().Log.Debug("team.openreq received")
	var msg keybase1.TeamOpenReqMsg
	if err := json.Unmarshal(item.Body().Bytes(), &msg); err != nil {
		r.G().Log.Debug("error unmarshaling team.openreq item: %s", err)
		return err
	}
	r.G().Log.Debug("team.openreq unmarshaled: %+v", msg)

	if err := teams.HandleOpenTeamAccessRequest(ctx, r.G(), msg); err != nil {
		return err
	}

	r.G().Log.Debug("dismissing team.openreq item since it succeeded")
	return r.G().GregorDismisser.DismissItem(cli, item.Metadata().MsgID())
}

func (r *teamHandler) Dismiss(ctx context.Context, cli gregor1.IncomingInterface, category string, item gregor.Item) (bool, error) {
	return false, nil
}

func (r *teamHandler) IsAlive() bool {
	return true
}

func (r *teamHandler) Name() string {
	return teamHandlerName
}
<|endoftext|>


<commit_before>package virt

import (
	"bufio"
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"
)

type Package struct {
	Fields []string
	Values map[string]string
}

var DPKG_INFO_EXTENSIONS = []string{"conffiles", "list", "md5sums", "postinst", "postrm", "preinst", "prerm", "shlibs", "symbols", "templates"}

func (pkg *Package) ID() string {
	name := strings.TrimSpace(pkg.Values["Package"])
	multiArch := strings.TrimSpace(pkg.Values["Multi-Arch"])
	if multiArch == "same" {
		arch := strings.TrimSpace(pkg.Values["Architecture"])
		return name + ":" + arch
	}
	return name
}

func (pkg *Package) InfoFile(extension string) string {
	return fmt.Sprintf("/var/lib/dpkg/info/%s.%s", pkg.ID(), extension)
}

func (vm *VM) MergeDpkgDatabase() {
	dpkgStatusFile := vm.OverlayFile("/var/lib/dpkg/status")
	upperPackages, err := ReadDpkgStatus(dpkgStatusFile)
	if err != nil {
		if !os.IsNotExist(err) {
			panic(err)
		}
		return // no file in upper, no need to merge
	}

	lowerPackages, err := ReadDpkgStatus(vm.LowerdirFile("/var/lib/dpkg/status"))
	if err != nil {
		panic(err)
	}

	// merge packages from lower to upper
	for name, lowerPkg := range lowerPackages {
		upperPkg, found := upperPackages[name]
		if found && upperPkg.Values["Version"] != lowerPkg.Values["Version"] {
			conffiles := make(map[string]string)

			lines := strings.Split(upperPkg.Values["Conffiles"], "\n")
			for _, line := range lines {
				parts := strings.Split(strings.TrimSpace(line), " ")
				if len(parts) == 2 {
					conffiles[parts[0]] = parts[1]
				}
			}

			// delete package files from overlay
			listFile := upperPkg.InfoFile("list")
			list, err := ioutil.ReadFile(vm.OverlayFile(listFile))
			if err != nil {
				list, err = ioutil.ReadFile(vm.LowerdirFile(listFile))
				if err != nil {
					panic(err)
				}
			}
			files := strings.Split(string(list), "\n")
			for _, file := range files {
				overlayFile := vm.OverlayFile(file)
				originalHash, found := conffiles[file]
				if found {
					hash := md5.New()
					f, err := os.Open(overlayFile)
					if err != nil {
						if os.IsNotExist(err) {
							continue // file not found in overlay
						}
						panic(err)
					}
					io.Copy(hash, f)
					f.Close()
					if string(hash.Sum(nil)) != originalHash {
						continue // config file was changed, do not delete from overlay
					}
				}
				os.Remove(overlayFile)
			}

			// delete informations from overlay
			for _, ext := range DPKG_INFO_EXTENSIONS {
				os.Remove(vm.OverlayFile(upperPkg.InfoFile(ext)))
			}
		}
		upperPackages[name] = lowerPkg
	}

	// delete packages that were removed in lower
	for name, upperPkg := range upperPackages {
		if _, found := lowerPackages[name]; found {
			continue // still in lower
		}
		_, err := os.Stat(vm.OverlayFile(upperPkg.InfoFile("list")))
		if err == nil {
			continue // files in overlay
		}
		delete(upperPackages, name)
	}

	if err = WriteDpkgStatus(upperPackages, dpkgStatusFile); err != nil {
		panic(err)
	}
}

func ReadDpkgStatus(fileName string) (map[string]*Package, error) {
	f, err := os.Open(fileName)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	r := bufio.NewReader(f)

	packages := make(map[string]*Package)
	for !atEndOfFile(r) {
		pkg := &Package{make([]string, 0), make(map[string]string)}
		for !tryReadByte(r, '\n') {
			field := readUntil(r, ':')
			value := ""
			if tryReadByte(r, '\n') {
				value += "\n"
			}
			for tryReadByte(r, ' ') {
				line := readUntil(r, '\n')
				value += " " + line + "\n"
			}
			pkg.Fields = append(pkg.Fields, field)
			pkg.Values[field] = value
		}
		packages[pkg.ID()] = pkg
	}
	return packages, nil
}

func WriteDpkgStatus(packages map[string]*Package, fileName string) error {
	f, err := os.Create(fileName)
	if err != nil {
		return err
	}
	defer f.Close()

	for _, pkg := range packages {
		for _, field := range pkg.Fields {
			f.Write([]byte(field + ":" + pkg.Values[field]))
		}
		f.Write([]byte{'\n'})
	}

	return nil
}
<commit_msg>os kite: Added (quite complicated) dpkg version parsing and comparing algorithm to prevent packages from downgrade caused by lowerdir merge.<commit_after>package virt

import (
	"bufio"
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"unicode"
)

type Package struct {
	Fields []string
	Values map[string]string
}

type DpkgVersion struct {
	Epoch    int
	Version  string
	Revision string
}

func ParseDpkgVerion(s string) *DpkgVersion {
	var version DpkgVersion

	epochParts := strings.SplitN(s, ":", 2)
	if len(epochParts) == 2 {
		version.Epoch, _ = strconv.Atoi(epochParts[0])
	}

	revisionParts := strings.SplitN(epochParts[len(epochParts)-1], "-", 2)
	if len(revisionParts) == 2 {
		version.Revision = revisionParts[1]
	}

	version.Version = revisionParts[0]

	return &version
}

// translated from dpkg source code
func dpkgOrder(r rune) int {
	if unicode.IsDigit(r) {
		return 0
	}
	if unicode.IsLetter(r) {
		return int(r)
	}
	if r == '~' {
		return -1
	}
	if r == 0 {
		return 0
	}
	return int(r) + 256
}

// translated from dpkg source code
func dpkgCompare(a string, b string) int {
	c := func(s string, i int) rune {
		if i >= len(s) {
			return 0
		}
		return rune(s[i])
	}

	ai := 0
	bi := 0
	for c(a, ai) != 0 || c(b, bi) != 0 {
		for (c(a, ai) != 0 && !unicode.IsDigit(c(a, ai))) || (c(b, bi) != 0 && !unicode.IsDigit(c(b, bi))) {
			ac := dpkgOrder(c(a, ai))
			bc := dpkgOrder(c(b, bi))

			if ac != bc {
				return ac - bc
			}

			ai += 1
			bi += 1
		}

		for c(a, ai) == '0' {
			ai += 1
		}
		for c(b, bi) == '0' {
			bi += 1
		}

		firstDiff := 0
		for unicode.IsDigit(c(a, ai)) && unicode.IsDigit(c(b, bi)) {
			if firstDiff == 0 {
				firstDiff = int(c(a, ai) - c(b, bi))
			}
			ai += 1
			bi += 1
		}
		if unicode.IsDigit(c(a, ai)) {
			return 1
		}
		if unicode.IsDigit(c(b, bi)) {
			return -1
		}
		if firstDiff != 0 {
			return firstDiff
		}
	}

	return 0
}

func (a *DpkgVersion) Compare(b *DpkgVersion) int {
	if a.Epoch > b.Epoch {
		return 1
	}
	if a.Epoch < b.Epoch {
		return -1
	}

	if r := dpkgCompare(a.Version, b.Version); r != 0 {
		return r
	}

	return dpkgCompare(a.Revision, b.Revision)
}

var DPKG_INFO_EXTENSIONS = []string{"conffiles", "list", "md5sums", "postinst", "postrm", "preinst", "prerm", "shlibs", "symbols", "templates"}

func (pkg *Package) ID() string {
	name := strings.TrimSpace(pkg.Values["Package"])
	multiArch := strings.TrimSpace(pkg.Values["Multi-Arch"])
	if multiArch == "same" {
		arch := strings.TrimSpace(pkg.Values["Architecture"])
		return name + ":" + arch
	}
	return name
}

func (pkg *Package) InfoFile(extension string) string {
	return fmt.Sprintf("/var/lib/dpkg/info/%s.%s", pkg.ID(), extension)
}

func (vm *VM) MergeDpkgDatabase() {
	dpkgStatusFile := vm.OverlayFile("/var/lib/dpkg/status")
	upperPackages, err := ReadDpkgStatus(dpkgStatusFile)
	if err != nil {
		if !os.IsNotExist(err) {
			panic(err)
		}
		return // no file in upper, no need to merge
	}

	lowerPackages, err := ReadDpkgStatus(vm.LowerdirFile("/var/lib/dpkg/status"))
	if err != nil {
		panic(err)
	}

	// merge packages from lower to upper
	for name, lowerPkg := range lowerPackages {
		upperPkg, found := upperPackages[name]
		if found && ParseDpkgVerion(upperPkg.Values["Version"]).Compare(ParseDpkgVerion(lowerPkg.Values["Version"])) < 0 {
			conffiles := make(map[string]string)

			lines := strings.Split(upperPkg.Values["Conffiles"], "\n")
			for _, line := range lines {
				parts := strings.Split(strings.TrimSpace(line), " ")
				if len(parts) == 2 {
					conffiles[parts[0]] = parts[1]
				}
			}

			// delete package files from overlay
			listFile := upperPkg.InfoFile("list")
			list, err := ioutil.ReadFile(vm.OverlayFile(listFile))
			if err != nil {
				list, err = ioutil.ReadFile(vm.LowerdirFile(listFile))
				if err != nil {
					panic(err)
				}
			}
			files := strings.Split(string(list), "\n")
			for _, file := range files {
				overlayFile := vm.OverlayFile(file)
				originalHash, found := conffiles[file]
				if found {
					hash := md5.New()
					f, err := os.Open(overlayFile)
					if err != nil {
						if os.IsNotExist(err) {
							continue // file not found in overlay
						}
						panic(err)
					}
					io.Copy(hash, f)
					f.Close()
					if string(hash.Sum(nil)) != originalHash {
						continue // config file was changed, do not delete from overlay
					}
				}
				os.Remove(overlayFile)
			}

			// delete informations from overlay
			for _, ext := range DPKG_INFO_EXTENSIONS {
				os.Remove(vm.OverlayFile(upperPkg.InfoFile(ext)))
			}
		}
		upperPackages[name] = lowerPkg
	}

	// delete packages that were removed in lower
	for name, upperPkg := range upperPackages {
		if _, found := lowerPackages[name]; found {
			continue // still in lower
		}
		_, err := os.Stat(vm.OverlayFile(upperPkg.InfoFile("list")))
		if err == nil {
			continue // files in overlay
		}
		delete(upperPackages, name)
	}

	if err = WriteDpkgStatus(upperPackages, dpkgStatusFile); err != nil {
		panic(err)
	}
}

func ReadDpkgStatus(fileName string) (map[string]*Package, error) {
	f, err := os.Open(fileName)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	r := bufio.NewReader(f)

	packages := make(map[string]*Package)
	for !atEndOfFile(r) {
		pkg := &Package{make([]string, 0), make(map[string]string)}
		for !tryReadByte(r, '\n') {
			field := readUntil(r, ':')
			value := ""
			if tryReadByte(r, '\n') {
				value += "\n"
			}
			for tryReadByte(r, ' ') {
				line := readUntil(r, '\n')
				value += " " + line + "\n"
			}
			pkg.Fields = append(pkg.Fields, field)
			pkg.Values[field] = value
		}
		packages[pkg.ID()] = pkg
	}
	return packages, nil
}

func WriteDpkgStatus(packages map[string]*Package, fileName string) error {
	f, err := os.Create(fileName)
	if err != nil {
		return err
	}
	defer f.Close()

	for _, pkg := range packages {
		for _, field := range pkg.Fields {
			f.Write([]byte(field + ":" + pkg.Values[field]))
		}
		f.Write([]byte{'\n'})
	}

	return nil
}
<|endoftext|>


<commit_before>// ----------------------------------------------------------------------------
//
//     ***     AUTO GENERATED CODE    ***    AUTO GENERATED CODE     ***
//
// ----------------------------------------------------------------------------
//
//     This file is automatically generated by Magic Modules and manual
//     changes will be clobbered when the file is regenerated.
//
//     Please read more about how to change this file in
//     .github/CONTRIBUTING.md.
//
// ----------------------------------------------------------------------------

package google

import (
	"bytes"
	"fmt"
	"reflect"
	"sort"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceComputeFirewallRuleHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%s-", m["protocol"].(string)))

	// We need to make sure to sort the strings below so that we always
	// generate the same hash code no matter what is in the set.
	if v, ok := m["ports"]; ok && v != nil {
		s := convertStringArr(v.([]interface{}))
		sort.Strings(s)

		for _, v := range s {
			buf.WriteString(fmt.Sprintf("%s-", v))
		}
	}

	return hashcode.String(buf.String())
}

func GetComputeFirewallCaiObject(d TerraformResourceData, config *Config) (Asset, error) {
	name, err := assetName(d, config, "//compute.googleapis.com/projects/{{project}}/global/firewalls/{{name}}")
	if err != nil {
		return Asset{}, err
	}
	if obj, err := GetComputeFirewallApiObject(d, config); err == nil {
		return Asset{
			Name: name,
			Type: "compute.googleapis.com/Firewall",
			Resource: &AssetResource{
				Version:              "v1",
				DiscoveryDocumentURI: "https://www.googleapis.com/discovery/v1/apis/compute/v1/rest",
				DiscoveryName:        "Firewall",
				Data:                 obj,
			},
		}, nil
	} else {
		return Asset{}, err
	}
}

func GetComputeFirewallApiObject(d TerraformResourceData, config *Config) (map[string]interface{}, error) {
	obj := make(map[string]interface{})
	allowedProp, err := expandComputeFirewallAllow(d.Get("allow"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("allow"); !isEmptyValue(reflect.ValueOf(allowedProp)) && (ok || !reflect.DeepEqual(v, allowedProp)) {
		obj["allowed"] = allowedProp
	}
	deniedProp, err := expandComputeFirewallDeny(d.Get("deny"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("deny"); !isEmptyValue(reflect.ValueOf(deniedProp)) && (ok || !reflect.DeepEqual(v, deniedProp)) {
		obj["denied"] = deniedProp
	}
	descriptionProp, err := expandComputeFirewallDescription(d.Get("description"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("description"); !isEmptyValue(reflect.ValueOf(descriptionProp)) && (ok || !reflect.DeepEqual(v, descriptionProp)) {
		obj["description"] = descriptionProp
	}
	destinationRangesProp, err := expandComputeFirewallDestinationRanges(d.Get("destination_ranges"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("destination_ranges"); !isEmptyValue(reflect.ValueOf(destinationRangesProp)) && (ok || !reflect.DeepEqual(v, destinationRangesProp)) {
		obj["destinationRanges"] = destinationRangesProp
	}
	directionProp, err := expandComputeFirewallDirection(d.Get("direction"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("direction"); !isEmptyValue(reflect.ValueOf(directionProp)) && (ok || !reflect.DeepEqual(v, directionProp)) {
		obj["direction"] = directionProp
	}
	disabledProp, err := expandComputeFirewallDisabled(d.Get("disabled"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("disabled"); ok || !reflect.DeepEqual(v, disabledProp) {
		obj["disabled"] = disabledProp
	}
	nameProp, err := expandComputeFirewallName(d.Get("name"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("name"); !isEmptyValue(reflect.ValueOf(nameProp)) && (ok || !reflect.DeepEqual(v, nameProp)) {
		obj["name"] = nameProp
	}
	networkProp, err := expandComputeFirewallNetwork(d.Get("network"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("network"); !isEmptyValue(reflect.ValueOf(networkProp)) && (ok || !reflect.DeepEqual(v, networkProp)) {
		obj["network"] = networkProp
	}
	priorityProp, err := expandComputeFirewallPriority(d.Get("priority"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("priority"); !isEmptyValue(reflect.ValueOf(priorityProp)) && (ok || !reflect.DeepEqual(v, priorityProp)) {
		obj["priority"] = priorityProp
	}
	sourceRangesProp, err := expandComputeFirewallSourceRanges(d.Get("source_ranges"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("source_ranges"); !isEmptyValue(reflect.ValueOf(sourceRangesProp)) && (ok || !reflect.DeepEqual(v, sourceRangesProp)) {
		obj["sourceRanges"] = sourceRangesProp
	}
	sourceServiceAccountsProp, err := expandComputeFirewallSourceServiceAccounts(d.Get("source_service_accounts"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("source_service_accounts"); !isEmptyValue(reflect.ValueOf(sourceServiceAccountsProp)) && (ok || !reflect.DeepEqual(v, sourceServiceAccountsProp)) {
		obj["sourceServiceAccounts"] = sourceServiceAccountsProp
	}
	sourceTagsProp, err := expandComputeFirewallSourceTags(d.Get("source_tags"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("source_tags"); !isEmptyValue(reflect.ValueOf(sourceTagsProp)) && (ok || !reflect.DeepEqual(v, sourceTagsProp)) {
		obj["sourceTags"] = sourceTagsProp
	}
	targetServiceAccountsProp, err := expandComputeFirewallTargetServiceAccounts(d.Get("target_service_accounts"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("target_service_accounts"); !isEmptyValue(reflect.ValueOf(targetServiceAccountsProp)) && (ok || !reflect.DeepEqual(v, targetServiceAccountsProp)) {
		obj["targetServiceAccounts"] = targetServiceAccountsProp
	}
	targetTagsProp, err := expandComputeFirewallTargetTags(d.Get("target_tags"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("target_tags"); !isEmptyValue(reflect.ValueOf(targetTagsProp)) && (ok || !reflect.DeepEqual(v, targetTagsProp)) {
		obj["targetTags"] = targetTagsProp
	}

	return obj, nil
}

func expandComputeFirewallAllow(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	l := v.([]interface{})
	req := make([]interface{}, 0, len(l))
	for _, raw := range l {
		if raw == nil {
			continue
		}
		original := raw.(map[string]interface{})
		transformed := make(map[string]interface{})

		transformedProtocol, err := expandComputeFirewallAllowProtocol(original["protocol"], d, config)
		if err != nil {
			return nil, err
		} else if val := reflect.ValueOf(transformedProtocol); val.IsValid() && !isEmptyValue(val) {
			transformed["IPProtocol"] = transformedProtocol
		}

		transformedPorts, err := expandComputeFirewallAllowPorts(original["ports"], d, config)
		if err != nil {
			return nil, err
		} else if val := reflect.ValueOf(transformedPorts); val.IsValid() && !isEmptyValue(val) {
			transformed["ports"] = transformedPorts
		}

		req = append(req, transformed)
	}
	return req, nil
}

func expandComputeFirewallAllowProtocol(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallAllowPorts(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallDeny(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	l := v.([]interface{})
	req := make([]interface{}, 0, len(l))
	for _, raw := range l {
		if raw == nil {
			continue
		}
		original := raw.(map[string]interface{})
		transformed := make(map[string]interface{})

		transformedProtocol, err := expandComputeFirewallDenyProtocol(original["protocol"], d, config)
		if err != nil {
			return nil, err
		} else if val := reflect.ValueOf(transformedProtocol); val.IsValid() && !isEmptyValue(val) {
			transformed["IPProtocol"] = transformedProtocol
		}

		transformedPorts, err := expandComputeFirewallDenyPorts(original["ports"], d, config)
		if err != nil {
			return nil, err
		} else if val := reflect.ValueOf(transformedPorts); val.IsValid() && !isEmptyValue(val) {
			transformed["ports"] = transformedPorts
		}

		req = append(req, transformed)
	}
	return req, nil
}

func expandComputeFirewallDenyProtocol(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallDenyPorts(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallDescription(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallDestinationRanges(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}

func expandComputeFirewallDirection(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallDisabled(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallName(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallNetwork(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	f, err := parseGlobalFieldValue("networks", v.(string), "project", d, config, true)
	if err != nil {
		return nil, fmt.Errorf("Invalid value for network: %s", err)
	}
	return f.RelativeLink(), nil
}

func expandComputeFirewallPriority(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallSourceRanges(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}

func expandComputeFirewallSourceServiceAccounts(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}

func expandComputeFirewallSourceTags(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}

func expandComputeFirewallTargetServiceAccounts(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}

func expandComputeFirewallTargetTags(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}
<commit_msg>allow setting firewall priority to 0 (#86)<commit_after>// ----------------------------------------------------------------------------
//
//     ***     AUTO GENERATED CODE    ***    AUTO GENERATED CODE     ***
//
// ----------------------------------------------------------------------------
//
//     This file is automatically generated by Magic Modules and manual
//     changes will be clobbered when the file is regenerated.
//
//     Please read more about how to change this file in
//     .github/CONTRIBUTING.md.
//
// ----------------------------------------------------------------------------

package google

import (
	"bytes"
	"fmt"
	"reflect"
	"sort"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceComputeFirewallRuleHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%s-", m["protocol"].(string)))

	// We need to make sure to sort the strings below so that we always
	// generate the same hash code no matter what is in the set.
	if v, ok := m["ports"]; ok && v != nil {
		s := convertStringArr(v.([]interface{}))
		sort.Strings(s)

		for _, v := range s {
			buf.WriteString(fmt.Sprintf("%s-", v))
		}
	}

	return hashcode.String(buf.String())
}

func GetComputeFirewallCaiObject(d TerraformResourceData, config *Config) (Asset, error) {
	name, err := assetName(d, config, "//compute.googleapis.com/projects/{{project}}/global/firewalls/{{name}}")
	if err != nil {
		return Asset{}, err
	}
	if obj, err := GetComputeFirewallApiObject(d, config); err == nil {
		return Asset{
			Name: name,
			Type: "compute.googleapis.com/Firewall",
			Resource: &AssetResource{
				Version:              "v1",
				DiscoveryDocumentURI: "https://www.googleapis.com/discovery/v1/apis/compute/v1/rest",
				DiscoveryName:        "Firewall",
				Data:                 obj,
			},
		}, nil
	} else {
		return Asset{}, err
	}
}

func GetComputeFirewallApiObject(d TerraformResourceData, config *Config) (map[string]interface{}, error) {
	obj := make(map[string]interface{})
	allowedProp, err := expandComputeFirewallAllow(d.Get("allow"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("allow"); !isEmptyValue(reflect.ValueOf(allowedProp)) && (ok || !reflect.DeepEqual(v, allowedProp)) {
		obj["allowed"] = allowedProp
	}
	deniedProp, err := expandComputeFirewallDeny(d.Get("deny"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("deny"); !isEmptyValue(reflect.ValueOf(deniedProp)) && (ok || !reflect.DeepEqual(v, deniedProp)) {
		obj["denied"] = deniedProp
	}
	descriptionProp, err := expandComputeFirewallDescription(d.Get("description"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("description"); !isEmptyValue(reflect.ValueOf(descriptionProp)) && (ok || !reflect.DeepEqual(v, descriptionProp)) {
		obj["description"] = descriptionProp
	}
	destinationRangesProp, err := expandComputeFirewallDestinationRanges(d.Get("destination_ranges"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("destination_ranges"); !isEmptyValue(reflect.ValueOf(destinationRangesProp)) && (ok || !reflect.DeepEqual(v, destinationRangesProp)) {
		obj["destinationRanges"] = destinationRangesProp
	}
	directionProp, err := expandComputeFirewallDirection(d.Get("direction"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("direction"); !isEmptyValue(reflect.ValueOf(directionProp)) && (ok || !reflect.DeepEqual(v, directionProp)) {
		obj["direction"] = directionProp
	}
	disabledProp, err := expandComputeFirewallDisabled(d.Get("disabled"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("disabled"); ok || !reflect.DeepEqual(v, disabledProp) {
		obj["disabled"] = disabledProp
	}
	nameProp, err := expandComputeFirewallName(d.Get("name"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("name"); !isEmptyValue(reflect.ValueOf(nameProp)) && (ok || !reflect.DeepEqual(v, nameProp)) {
		obj["name"] = nameProp
	}
	networkProp, err := expandComputeFirewallNetwork(d.Get("network"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("network"); !isEmptyValue(reflect.ValueOf(networkProp)) && (ok || !reflect.DeepEqual(v, networkProp)) {
		obj["network"] = networkProp
	}
	priorityProp, err := expandComputeFirewallPriority(d.Get("priority"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("priority"); ok || !reflect.DeepEqual(v, priorityProp) {
		obj["priority"] = priorityProp
	}
	sourceRangesProp, err := expandComputeFirewallSourceRanges(d.Get("source_ranges"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("source_ranges"); !isEmptyValue(reflect.ValueOf(sourceRangesProp)) && (ok || !reflect.DeepEqual(v, sourceRangesProp)) {
		obj["sourceRanges"] = sourceRangesProp
	}
	sourceServiceAccountsProp, err := expandComputeFirewallSourceServiceAccounts(d.Get("source_service_accounts"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("source_service_accounts"); !isEmptyValue(reflect.ValueOf(sourceServiceAccountsProp)) && (ok || !reflect.DeepEqual(v, sourceServiceAccountsProp)) {
		obj["sourceServiceAccounts"] = sourceServiceAccountsProp
	}
	sourceTagsProp, err := expandComputeFirewallSourceTags(d.Get("source_tags"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("source_tags"); !isEmptyValue(reflect.ValueOf(sourceTagsProp)) && (ok || !reflect.DeepEqual(v, sourceTagsProp)) {
		obj["sourceTags"] = sourceTagsProp
	}
	targetServiceAccountsProp, err := expandComputeFirewallTargetServiceAccounts(d.Get("target_service_accounts"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("target_service_accounts"); !isEmptyValue(reflect.ValueOf(targetServiceAccountsProp)) && (ok || !reflect.DeepEqual(v, targetServiceAccountsProp)) {
		obj["targetServiceAccounts"] = targetServiceAccountsProp
	}
	targetTagsProp, err := expandComputeFirewallTargetTags(d.Get("target_tags"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("target_tags"); !isEmptyValue(reflect.ValueOf(targetTagsProp)) && (ok || !reflect.DeepEqual(v, targetTagsProp)) {
		obj["targetTags"] = targetTagsProp
	}

	return obj, nil
}

func expandComputeFirewallAllow(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	l := v.([]interface{})
	req := make([]interface{}, 0, len(l))
	for _, raw := range l {
		if raw == nil {
			continue
		}
		original := raw.(map[string]interface{})
		transformed := make(map[string]interface{})

		transformedProtocol, err := expandComputeFirewallAllowProtocol(original["protocol"], d, config)
		if err != nil {
			return nil, err
		} else if val := reflect.ValueOf(transformedProtocol); val.IsValid() && !isEmptyValue(val) {
			transformed["IPProtocol"] = transformedProtocol
		}

		transformedPorts, err := expandComputeFirewallAllowPorts(original["ports"], d, config)
		if err != nil {
			return nil, err
		} else if val := reflect.ValueOf(transformedPorts); val.IsValid() && !isEmptyValue(val) {
			transformed["ports"] = transformedPorts
		}

		req = append(req, transformed)
	}
	return req, nil
}

func expandComputeFirewallAllowProtocol(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallAllowPorts(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallDeny(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	l := v.([]interface{})
	req := make([]interface{}, 0, len(l))
	for _, raw := range l {
		if raw == nil {
			continue
		}
		original := raw.(map[string]interface{})
		transformed := make(map[string]interface{})

		transformedProtocol, err := expandComputeFirewallDenyProtocol(original["protocol"], d, config)
		if err != nil {
			return nil, err
		} else if val := reflect.ValueOf(transformedProtocol); val.IsValid() && !isEmptyValue(val) {
			transformed["IPProtocol"] = transformedProtocol
		}

		transformedPorts, err := expandComputeFirewallDenyPorts(original["ports"], d, config)
		if err != nil {
			return nil, err
		} else if val := reflect.ValueOf(transformedPorts); val.IsValid() && !isEmptyValue(val) {
			transformed["ports"] = transformedPorts
		}

		req = append(req, transformed)
	}
	return req, nil
}

func expandComputeFirewallDenyProtocol(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallDenyPorts(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallDescription(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallDestinationRanges(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}

func expandComputeFirewallDirection(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallDisabled(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallName(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallNetwork(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	f, err := parseGlobalFieldValue("networks", v.(string), "project", d, config, true)
	if err != nil {
		return nil, fmt.Errorf("Invalid value for network: %s", err)
	}
	return f.RelativeLink(), nil
}

func expandComputeFirewallPriority(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeFirewallSourceRanges(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}

func expandComputeFirewallSourceServiceAccounts(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}

func expandComputeFirewallSourceTags(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}

func expandComputeFirewallTargetServiceAccounts(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}

func expandComputeFirewallTargetTags(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	v = v.(*schema.Set).List()
	return v, nil
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2014-2015 VMware, Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package esxcli

import (
	"context"
	"flag"
	"fmt"
	"io"
	"sort"
	"strings"
	"text/tabwriter"

	"github.com/vmware/govmomi/govc/cli"
	"github.com/vmware/govmomi/govc/flags"
)

type esxcli struct {
	*flags.HostSystemFlag

	hints bool
}

func init() {
	cli.Register("host.esxcli", &esxcli{})
}

func (cmd *esxcli) Usage() string {
	return "COMMAND [ARG]..."
}

func (cmd *esxcli) Register(ctx context.Context, f *flag.FlagSet) {
	cmd.HostSystemFlag, ctx = flags.NewHostSystemFlag(ctx)
	cmd.HostSystemFlag.Register(ctx, f)

	f.BoolVar(&cmd.hints, "hints", true, "Use command info hints when formatting output")
}

func (cmd *esxcli) Description() string {
	return `Invoke esxcli command on HOST.

Output is rendered in table form when possible, unless disabled with '-hints=false'.

Examples:
  govc host.esxcli network ip connection list
  govc host.esxcli system settings advanced set -o /Net/GuestIPHack -i 1
  govc host.esxcli network firewall ruleset set -r remoteSerialPort -e true
  govc host.esxcli network firewall set -e false
  govc host.esxcli hardware platform get`
}

func (cmd *esxcli) Process(ctx context.Context) error {
	if err := cmd.HostSystemFlag.Process(ctx); err != nil {
		return err
	}
	return nil
}

func (cmd *esxcli) Run(ctx context.Context, f *flag.FlagSet) error {
	c, err := cmd.Client()
	if err != nil {
		return err
	}

	host, err := cmd.HostSystem()
	if err != nil {
		return err
	}

	e, err := NewExecutor(c, host)
	if err != nil {
		return err
	}

	res, err := e.Run(f.Args())
	if err != nil {
		return err
	}

	if len(res.Values) == 0 {
		if res.String != "" {
			fmt.Print(res.String)
			if !strings.HasSuffix(res.String, "\n") {
				fmt.Println()
			}
		}
		return nil
	}

	return cmd.WriteResult(&result{res, cmd})
}

type result struct {
	*Response
	cmd *esxcli
}

func (r *result) Write(w io.Writer) error {
	var formatType string
	if r.cmd.hints {
		formatType = r.Info.Hints.Formatter()
	}

	switch formatType {
	case "table":
		r.cmd.formatTable(w, r.Response)
	default:
		r.cmd.formatSimple(w, r.Response)
	}

	return nil
}

func (cmd *esxcli) formatSimple(w io.Writer, res *Response) {
	var keys []string
	for key := range res.Values[0] {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	tw := tabwriter.NewWriter(w, 2, 0, 2, ' ', 0)

	for i, rv := range res.Values {
		if i > 0 {
			fmt.Fprintln(tw)
			_ = tw.Flush()
		}
		for _, key := range keys {
			fmt.Fprintf(tw, "%s:\t%s\n", key, strings.Join(rv[key], ", "))
		}
	}

	_ = tw.Flush()
}

func (cmd *esxcli) formatTable(w io.Writer, res *Response) {
	fields := res.Info.Hints.Fields()
	if len(fields) == 0 {
		cmd.formatSimple(w, res)
		return
	}
	tw := tabwriter.NewWriter(w, len(fields), 0, 2, ' ', 0)

	var hr []string
	for _, name := range fields {
		hr = append(hr, strings.Repeat("-", len(name)))
	}

	fmt.Fprintln(tw, strings.Join(fields, "\t"))
	fmt.Fprintln(tw, strings.Join(hr, "\t"))

	for _, vals := range res.Values {
		var row []string

		for _, name := range fields {
			key := strings.Replace(name, " ", "", -1)
			if val, ok := vals[key]; ok {
				row = append(row, strings.Join(val, ", "))
			} else {
				row = append(row, "")
			}
		}

		fmt.Fprintln(tw, strings.Join(row, "\t"))
	}

	_ = tw.Flush()
}
<commit_msg>govc: fix host.esxcli runtime error occurred when no arguments specified<commit_after>/*
Copyright (c) 2014-2015 VMware, Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package esxcli

import (
	"context"
	"flag"
	"fmt"
	"io"
	"sort"
	"strings"
	"text/tabwriter"

	"github.com/vmware/govmomi/govc/cli"
	"github.com/vmware/govmomi/govc/flags"
)

type esxcli struct {
	*flags.HostSystemFlag

	hints bool
}

func init() {
	cli.Register("host.esxcli", &esxcli{})
}

func (cmd *esxcli) Usage() string {
	return "COMMAND [ARG]..."
}

func (cmd *esxcli) Register(ctx context.Context, f *flag.FlagSet) {
	cmd.HostSystemFlag, ctx = flags.NewHostSystemFlag(ctx)
	cmd.HostSystemFlag.Register(ctx, f)

	f.BoolVar(&cmd.hints, "hints", true, "Use command info hints when formatting output")
}

func (cmd *esxcli) Description() string {
	return `Invoke esxcli command on HOST.

Output is rendered in table form when possible, unless disabled with '-hints=false'.

Examples:
  govc host.esxcli network ip connection list
  govc host.esxcli system settings advanced set -o /Net/GuestIPHack -i 1
  govc host.esxcli network firewall ruleset set -r remoteSerialPort -e true
  govc host.esxcli network firewall set -e false
  govc host.esxcli hardware platform get`
}

func (cmd *esxcli) Process(ctx context.Context) error {
	if err := cmd.HostSystemFlag.Process(ctx); err != nil {
		return err
	}
	return nil
}

func (cmd *esxcli) Run(ctx context.Context, f *flag.FlagSet) error {
	if f.NArg() == 0 {
		return flag.ErrHelp
	}

	c, err := cmd.Client()
	if err != nil {
		return err
	}

	host, err := cmd.HostSystem()
	if err != nil {
		return err
	}

	e, err := NewExecutor(c, host)
	if err != nil {
		return err
	}

	res, err := e.Run(f.Args())
	if err != nil {
		return err
	}

	if len(res.Values) == 0 {
		if res.String != "" {
			fmt.Print(res.String)
			if !strings.HasSuffix(res.String, "\n") {
				fmt.Println()
			}
		}
		return nil
	}

	return cmd.WriteResult(&result{res, cmd})
}

type result struct {
	*Response
	cmd *esxcli
}

func (r *result) Write(w io.Writer) error {
	var formatType string
	if r.cmd.hints {
		formatType = r.Info.Hints.Formatter()
	}

	switch formatType {
	case "table":
		r.cmd.formatTable(w, r.Response)
	default:
		r.cmd.formatSimple(w, r.Response)
	}

	return nil
}

func (cmd *esxcli) formatSimple(w io.Writer, res *Response) {
	var keys []string
	for key := range res.Values[0] {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	tw := tabwriter.NewWriter(w, 2, 0, 2, ' ', 0)

	for i, rv := range res.Values {
		if i > 0 {
			fmt.Fprintln(tw)
			_ = tw.Flush()
		}
		for _, key := range keys {
			fmt.Fprintf(tw, "%s:\t%s\n", key, strings.Join(rv[key], ", "))
		}
	}

	_ = tw.Flush()
}

func (cmd *esxcli) formatTable(w io.Writer, res *Response) {
	fields := res.Info.Hints.Fields()
	if len(fields) == 0 {
		cmd.formatSimple(w, res)
		return
	}
	tw := tabwriter.NewWriter(w, len(fields), 0, 2, ' ', 0)

	var hr []string
	for _, name := range fields {
		hr = append(hr, strings.Repeat("-", len(name)))
	}

	fmt.Fprintln(tw, strings.Join(fields, "\t"))
	fmt.Fprintln(tw, strings.Join(hr, "\t"))

	for _, vals := range res.Values {
		var row []string

		for _, name := range fields {
			key := strings.Replace(name, " ", "", -1)
			if val, ok := vals[key]; ok {
				row = append(row, strings.Join(val, ", "))
			} else {
				row = append(row, "")
			}
		}

		fmt.Fprintln(tw, strings.Join(row, "\t"))
	}

	_ = tw.Flush()
}
<|endoftext|>


<commit_before>package signal

import (
	"context"
	"time"
)

type ActivityUpdater interface {
	Update()
}

type ActivityTimer struct {
	updated chan bool
	timeout chan time.Duration
	ctx     context.Context
	cancel  context.CancelFunc
}

func (t *ActivityTimer) Update() {
	select {
	case t.updated <- true:
	default:
	}
}

func (t *ActivityTimer) SetTimeout(timeout time.Duration) {
	t.timeout <- timeout
}

func (t *ActivityTimer) run() {
	ticker := time.NewTicker(<-t.timeout)
	defer func() {
		ticker.Stop()
	}()

	for {
		select {
		case <-ticker.C:
		case <-t.ctx.Done():
			return
		case timeout := <-t.timeout:
			if timeout == 0 {
				t.cancel()
				return
			}

			ticker.Stop()
			ticker = time.NewTicker(timeout)
		}

		select {
		case <-t.updated:
		// Updated keep waiting.
		default:
			t.cancel()
			return
		}
	}
}

func CancelAfterInactivity(ctx context.Context, cancel context.CancelFunc, timeout time.Duration) *ActivityTimer {
	timer := &ActivityTimer{
		ctx:     ctx,
		cancel:  cancel,
		timeout: make(chan time.Duration, 1),
		updated: make(chan bool, 1),
	}
	timer.timeout <- timeout
	go timer.run()
	return timer
}
<commit_msg>move ctx into function<commit_after>package signal

import (
	"context"
	"time"
)

type ActivityUpdater interface {
	Update()
}

type ActivityTimer struct {
	updated chan bool
	timeout chan time.Duration
}

func (t *ActivityTimer) Update() {
	select {
	case t.updated <- true:
	default:
	}
}

func (t *ActivityTimer) SetTimeout(timeout time.Duration) {
	t.timeout <- timeout
}

func (t *ActivityTimer) run(ctx context.Context, cancel context.CancelFunc) {
	ticker := time.NewTicker(<-t.timeout)
	defer func() {
		ticker.Stop()
	}()

	for {
		select {
		case <-ticker.C:
		case <-ctx.Done():
			return
		case timeout := <-t.timeout:
			if timeout == 0 {
				cancel()
				return
			}

			ticker.Stop()
			ticker = time.NewTicker(timeout)
			continue
		}

		select {
		case <-t.updated:
		// Updated keep waiting.
		default:
			cancel()
			return
		}
	}
}

func CancelAfterInactivity(ctx context.Context, cancel context.CancelFunc, timeout time.Duration) *ActivityTimer {
	timer := &ActivityTimer{
		timeout: make(chan time.Duration, 1),
		updated: make(chan bool, 1),
	}
	timer.timeout <- timeout
	go timer.run(ctx, cancel)
	return timer
}
<|endoftext|>


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

package compilergraph

import (
	"fmt"

	"github.com/google/cayley"
	"github.com/google/cayley/graph"
	"github.com/google/cayley/graph/path"
)

var _ = fmt.Printf

// GraphQuery is a type which wraps a Cayley Path and provides nice accessors for querying
// the graph layer.
type GraphQuery struct {
	path  *path.Path  // The wrapped Cayley Path.
	layer *GraphLayer // The layer under which this query was created.
}

// graphNodeIterator represents an iterator over a GraphQuery, with each call to Next()
// updating the iterator with a node ID and a map of values found off of the specified predicates.
type graphNodeIterator struct {
	layer      *GraphLayer    // The parent graph layer.
	iterator   graph.Iterator // The wrapped Cayley Iterator.
	predicates []string       // The set of predicates to retrieve.

	Node   GraphNode         // The current node (if any).
	Values map[string]string // The current predicate values (if any).
}

// StartQuery returns a new query starting at the nodes with the given names.
func (gl *GraphLayer) StartQuery(nodeNames ...string) *GraphQuery {
	return &GraphQuery{
		path:  cayley.StartPath(gl.cayleyStore, nodeNames...),
		layer: gl,
	}
}

// AllNodesQuery returns a query starting from all the nodes in the layer.
func (gl *GraphLayer) AllNodesQuery() *GraphQuery {
	return &GraphQuery{
		path:  cayley.StartPath(gl.cayleyStore, gl.id).In(nodeMemberPredicate),
		layer: gl,
	}
}

// FindNodesOfKind returns a new query starting at the nodes who have the given kind in this layer.
func (gl *GraphLayer) FindNodesOfKind(kinds ...TaggedValue) *GraphQuery {
	return gl.FindNodesWithTaggedType(gl.nodeKindPredicate, kinds...)
}

// FindNodesWithTaggedType returns a new query starting at the nodes who are linked to tagged values
// (of the given name) by the given predicate.
//
// For example:
//
// `FindNodesWithTaggedType("parser-ast-node-type", NodeType.Class, NodeType.Interface)`
// would return all classes and interfaces.
func (gl *GraphLayer) FindNodesWithTaggedType(predicate string, values ...TaggedValue) *GraphQuery {
	var nodeNames []string
	for _, value := range values {
		nodeNames = append(nodeNames, gl.getTaggedKey(value))
	}

	return gl.StartQuery(nodeNames...).In(predicate)
}

// IsKind updates thi Query to represent only those nodes that are of the given kind.
func (gq *GraphQuery) IsKind(nodeKinds ...TaggedValue) *GraphQuery {
	return gq.HasTagged(gq.layer.nodeKindPredicate, nodeKinds...)
}

// In updates this Query to represent the nodes that are adjacent to the
// current nodes, via the given inbound predicate.
func (gq *GraphQuery) In(via ...string) *GraphQuery {
	adjustedVia := gq.getAdjustedPredicates(via...)

	return &GraphQuery{
		path:  gq.path.In(adjustedVia...),
		layer: gq.layer,
	}
}

// Out updates this Query to represent the nodes that are adjacent to the
// current nodes, via the given outbound predicate.
func (gq *GraphQuery) Out(via ...string) *GraphQuery {
	adjustedVia := gq.getAdjustedPredicates(via...)
	return &GraphQuery{
		path:  gq.path.Out(adjustedVia...),
		layer: gq.layer,
	}
}

// HasTagged filters this Query to represent the nodes that have some linkage
// to some known node.
func (gq *GraphQuery) HasTagged(via string, taggedValues ...TaggedValue) *GraphQuery {
	var values []string = make([]string, len(taggedValues))
	for index, taggedValue := range taggedValues {
		values[index] = gq.layer.getTaggedKey(taggedValue)
	}

	return gq.Has(via, values...)
}

// Has filters this Query to represent the nodes that have some linkage
// to some known node.
func (gq *GraphQuery) Has(via string, nodes ...string) *GraphQuery {
	adjustedVia := gq.getAdjustedPredicates(via)[0]
	return &GraphQuery{
		path:  gq.path.Has(adjustedVia, nodes...),
		layer: gq.layer,
	}
}

func (gq *GraphQuery) getAdjustedPredicates(predicates ...string) []interface{} {
	adjusted := make([]interface{}, 0, len(predicates))

	for _, predicate := range predicates {
		fullPredicate := gq.layer.prefix + "-" + predicate
		adjusted = append(adjusted, fullPredicate)
	}
	return adjusted
}

// GetValue executes the query and returns the name of the node found, as a value.
func (gq *GraphQuery) GetValue() (string, bool) {
	it := gq.path.BuildIterator()
	result := cayley.RawNext(it)
	if !result {
		return "", false
	}

	return gq.layer.cayleyStore.NameOf(it.Result()), true
}

// GetNode executes the query and returns the single node found or panics.
func (gq *GraphQuery) GetNode() GraphNode {
	node, found := gq.TryGetNode()
	if !found {
		panic(fmt.Sprintf("Could not return node for query: %v", gq))
	}
	return node
}

// TryGetNode executes the query and returns the single node found or false. If there is
// more than a single node as a result of the query, the first node is returned.
func (gq *GraphQuery) TryGetNode() (GraphNode, bool) {
	it := gq.BuildNodeIterator()
	if !it.Next() {
		return GraphNode{}, false
	}

	return it.Node, true
}

// BuildNodeIterator returns an iterator for retrieving the results of the query, with
// each result being a struct representing the node and the values found outgoing at the
// given predicates.
func (gq *GraphQuery) BuildNodeIterator(predicates ...string) *graphNodeIterator {
	var updatedPath *path.Path = gq.path

	// Save the predicates the user requested.
	for _, predicate := range predicates {
		fullPredicate := gq.layer.prefix + "-" + predicate
		updatedPath = updatedPath.Save(fullPredicate, fullPredicate)
	}

	// Save the predicate for the kind of the node as well.
	fullKindPredicate := gq.layer.prefix + "-" + gq.layer.nodeKindPredicate
	updatedPath = updatedPath.Save(fullKindPredicate, fullKindPredicate)

	it := updatedPath.BuildIterator()

	// TODO(jschorr): Uncomment and use this once fixed.
	// oit, _ := it.Optimize()

	return &graphNodeIterator{
		layer:      gq.layer,
		iterator:   it,
		predicates: predicates,
	}
}

// Returns the tagged value at the given predicate on the current node.
func (gni *graphNodeIterator) GetTaggedValue(predicate string, example TaggedValue) interface{} {
	strValue := gni.Values[predicate]
	return gni.layer.parseTaggedKey(strValue, example)
}

// Next move the iterator forward.
func (gni *graphNodeIterator) Next() bool {
	result := cayley.RawNext(gni.iterator)
	if !result {
		return false
	}

	tags := make(map[string]graph.Value)
	gni.iterator.TagResults(tags)

	// Copy the tags over, making sure to update the predicates to reflect
	// the current layer.
	updatedTags := make(map[string]string)
	for _, predicate := range gni.predicates {
		fullPredicate := gni.layer.prefix + "-" + predicate
		updatedTags[predicate] = gni.layer.cayleyStore.NameOf(tags[fullPredicate])
	}

	// Load the kind of the node.
	fullKindPredicate := gni.layer.prefix + "-" + gni.layer.nodeKindPredicate
	kindString := gni.layer.cayleyStore.NameOf(tags[fullKindPredicate])

	node := GraphNode{
		NodeId: GraphNodeId(gni.layer.cayleyStore.NameOf(gni.iterator.Result())),
		Kind:   gni.layer.parseTaggedKey(kindString, gni.layer.nodeKindEnum).(TaggedValue),
		layer:  gni.layer,
	}

	gni.Node = node
	gni.Values = updatedTags
	return true
}
<commit_msg>Enable Cayley optimizations<commit_after>// Copyright 2015 The Serulian Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package compilergraph

import (
	"fmt"

	"github.com/google/cayley"
	"github.com/google/cayley/graph"
	"github.com/google/cayley/graph/path"
)

var _ = fmt.Printf

// GraphQuery is a type which wraps a Cayley Path and provides nice accessors for querying
// the graph layer.
type GraphQuery struct {
	path  *path.Path  // The wrapped Cayley Path.
	layer *GraphLayer // The layer under which this query was created.
}

// graphNodeIterator represents an iterator over a GraphQuery, with each call to Next()
// updating the iterator with a node ID and a map of values found off of the specified predicates.
type graphNodeIterator struct {
	layer      *GraphLayer    // The parent graph layer.
	iterator   graph.Iterator // The wrapped Cayley Iterator.
	predicates []string       // The set of predicates to retrieve.

	Node   GraphNode         // The current node (if any).
	Values map[string]string // The current predicate values (if any).
}

// StartQuery returns a new query starting at the nodes with the given names.
func (gl *GraphLayer) StartQuery(nodeNames ...string) *GraphQuery {
	return &GraphQuery{
		path:  cayley.StartPath(gl.cayleyStore, nodeNames...),
		layer: gl,
	}
}

// AllNodesQuery returns a query starting from all the nodes in the layer.
func (gl *GraphLayer) AllNodesQuery() *GraphQuery {
	return &GraphQuery{
		path:  cayley.StartPath(gl.cayleyStore, gl.id).In(nodeMemberPredicate),
		layer: gl,
	}
}

// FindNodesOfKind returns a new query starting at the nodes who have the given kind in this layer.
func (gl *GraphLayer) FindNodesOfKind(kinds ...TaggedValue) *GraphQuery {
	return gl.FindNodesWithTaggedType(gl.nodeKindPredicate, kinds...)
}

// FindNodesWithTaggedType returns a new query starting at the nodes who are linked to tagged values
// (of the given name) by the given predicate.
//
// For example:
//
// `FindNodesWithTaggedType("parser-ast-node-type", NodeType.Class, NodeType.Interface)`
// would return all classes and interfaces.
func (gl *GraphLayer) FindNodesWithTaggedType(predicate string, values ...TaggedValue) *GraphQuery {
	var nodeNames []string
	for _, value := range values {
		nodeNames = append(nodeNames, gl.getTaggedKey(value))
	}

	return gl.StartQuery(nodeNames...).In(predicate)
}

// IsKind updates thi Query to represent only those nodes that are of the given kind.
func (gq *GraphQuery) IsKind(nodeKinds ...TaggedValue) *GraphQuery {
	return gq.HasTagged(gq.layer.nodeKindPredicate, nodeKinds...)
}

// In updates this Query to represent the nodes that are adjacent to the
// current nodes, via the given inbound predicate.
func (gq *GraphQuery) In(via ...string) *GraphQuery {
	adjustedVia := gq.getAdjustedPredicates(via...)

	return &GraphQuery{
		path:  gq.path.In(adjustedVia...),
		layer: gq.layer,
	}
}

// Out updates this Query to represent the nodes that are adjacent to the
// current nodes, via the given outbound predicate.
func (gq *GraphQuery) Out(via ...string) *GraphQuery {
	adjustedVia := gq.getAdjustedPredicates(via...)
	return &GraphQuery{
		path:  gq.path.Out(adjustedVia...),
		layer: gq.layer,
	}
}

// HasTagged filters this Query to represent the nodes that have some linkage
// to some known node.
func (gq *GraphQuery) HasTagged(via string, taggedValues ...TaggedValue) *GraphQuery {
	var values []string = make([]string, len(taggedValues))
	for index, taggedValue := range taggedValues {
		values[index] = gq.layer.getTaggedKey(taggedValue)
	}

	return gq.Has(via, values...)
}

// Has filters this Query to represent the nodes that have some linkage
// to some known node.
func (gq *GraphQuery) Has(via string, nodes ...string) *GraphQuery {
	adjustedVia := gq.getAdjustedPredicates(via)[0]
	return &GraphQuery{
		path:  gq.path.Has(adjustedVia, nodes...),
		layer: gq.layer,
	}
}

func (gq *GraphQuery) getAdjustedPredicates(predicates ...string) []interface{} {
	adjusted := make([]interface{}, 0, len(predicates))

	for _, predicate := range predicates {
		fullPredicate := gq.layer.prefix + "-" + predicate
		adjusted = append(adjusted, fullPredicate)
	}
	return adjusted
}

// GetValue executes the query and returns the name of the node found, as a value.
func (gq *GraphQuery) GetValue() (string, bool) {
	it := gq.path.BuildIterator()
	result := cayley.RawNext(it)
	if !result {
		return "", false
	}

	return gq.layer.cayleyStore.NameOf(it.Result()), true
}

// GetNode executes the query and returns the single node found or panics.
func (gq *GraphQuery) GetNode() GraphNode {
	node, found := gq.TryGetNode()
	if !found {
		panic(fmt.Sprintf("Could not return node for query: %v", gq))
	}
	return node
}

// TryGetNode executes the query and returns the single node found or false. If there is
// more than a single node as a result of the query, the first node is returned.
func (gq *GraphQuery) TryGetNode() (GraphNode, bool) {
	it := gq.BuildNodeIterator()
	if !it.Next() {
		return GraphNode{}, false
	}

	return it.Node, true
}

// BuildNodeIterator returns an iterator for retrieving the results of the query, with
// each result being a struct representing the node and the values found outgoing at the
// given predicates.
func (gq *GraphQuery) BuildNodeIterator(predicates ...string) *graphNodeIterator {
	var updatedPath *path.Path = gq.path

	// Save the predicates the user requested.
	for _, predicate := range predicates {
		fullPredicate := gq.layer.prefix + "-" + predicate
		updatedPath = updatedPath.Save(fullPredicate, fullPredicate)
	}

	// Save the predicate for the kind of the node as well.
	fullKindPredicate := gq.layer.prefix + "-" + gq.layer.nodeKindPredicate
	updatedPath = updatedPath.Save(fullKindPredicate, fullKindPredicate)

	it := updatedPath.BuildIterator()
	oit, _ := it.Optimize()

	return &graphNodeIterator{
		layer:      gq.layer,
		iterator:   oit,
		predicates: predicates,
	}
}

// Returns the tagged value at the given predicate on the current node.
func (gni *graphNodeIterator) GetTaggedValue(predicate string, example TaggedValue) interface{} {
	strValue := gni.Values[predicate]
	return gni.layer.parseTaggedKey(strValue, example)
}

// Next move the iterator forward.
func (gni *graphNodeIterator) Next() bool {
	result := cayley.RawNext(gni.iterator)
	if !result {
		return false
	}

	tags := make(map[string]graph.Value)
	gni.iterator.TagResults(tags)

	// Copy the tags over, making sure to update the predicates to reflect
	// the current layer.
	updatedTags := make(map[string]string)
	for _, predicate := range gni.predicates {
		fullPredicate := gni.layer.prefix + "-" + predicate
		updatedTags[predicate] = gni.layer.cayleyStore.NameOf(tags[fullPredicate])
	}

	// Load the kind of the node.
	fullKindPredicate := gni.layer.prefix + "-" + gni.layer.nodeKindPredicate
	kindString := gni.layer.cayleyStore.NameOf(tags[fullKindPredicate])

	node := GraphNode{
		NodeId: GraphNodeId(gni.layer.cayleyStore.NameOf(gni.iterator.Result())),
		Kind:   gni.layer.parseTaggedKey(kindString, gni.layer.nodeKindEnum).(TaggedValue),
		layer:  gni.layer,
	}

	gni.Node = node
	gni.Values = updatedTags
	return true
}
<|endoftext|>