<commit_before>package commands

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

	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/git/odb"
	"github.com/git-lfs/git-lfs/tools"
	"github.com/git-lfs/git-lfs/tools/humanize"
	"github.com/spf13/cobra"
)

var (
	// migrateInfoTopN is a flag given to the git-lfs-migrate(1) subcommand
	// 'info' which specifies how many info entries to show by default.
	migrateInfoTopN int

	// migrateInfoAboveFmt is a flag given to the git-lfs-migrate(1)
	// subcommand 'info' specifying a human-readable string threshold of
	// filesize before entries are counted.
	migrateInfoAboveFmt string
	// migrateInfoAbove is the number of bytes parsed from the above
	// migrateInfoAboveFmt flag.
	migrateInfoAbove uint64
)

func migrateInfoCommand(cmd *cobra.Command, args []string) {
	exts := make(map[string]*MigrateInfoEntry)

	above, err := humanize.ParseBytes(migrateInfoAboveFmt)
	if err != nil {
		ExitWithError(errors.Wrap(err, "cannot parse --above=<n>"))
	}

	migrateInfoAbove = above

	migrate(cmd, args, func(path string, b *odb.Blob) (*odb.Blob, error) {
		ext := fmt.Sprintf("*%s", filepath.Ext(path))

		if len(ext) > 1 {
			entry := exts[ext]
			if entry == nil {
				entry = &MigrateInfoEntry{Qualifier: ext}
			}

			entry.Total++
			entry.BytesTotal += b.Size

			if b.Size > int64(migrateInfoAbove) {
				entry.TotalAbove++
				entry.BytesAbove += b.Size
			}

			exts[ext] = entry
		}

		return b, nil
	})

	entries := EntriesBySize(MapToEntries(exts))
	sort.Sort(sort.Reverse(entries))

	migrateInfoTopN = tools.ClampInt(migrateInfoTopN, len(entries), 0)

	entries = entries[:tools.MaxInt(0, migrateInfoTopN)]

	entries.Print(os.Stderr)
}

// MigrateInfoEntry represents a tuple of filetype to bytes and entry count
// above and below a threshold.
type MigrateInfoEntry struct {
	// Qualifier is the filepath's extension.
	Qualifier string

	// BytesAbove is total size of all files above a given threshold.
	BytesAbove int64
	// TotalAbove is the count of all files above a given size threshold.
	TotalAbove int64
	// BytesTotal is the number of bytes of all files
	BytesTotal int64
	// Total is the count of all files.
	Total int64
}

// MapToEntries creates a set of `*MigrateInfoEntry`'s for a given map of
// filepath extensions to file size in bytes.
func MapToEntries(exts map[string]*MigrateInfoEntry) []*MigrateInfoEntry {
	entries := make([]*MigrateInfoEntry, 0, len(exts))
	for _, entry := range exts {
		entries = append(entries, entry)
	}

	return entries
}

// EntriesBySize is an implementation of sort.Interface that sorts a set of
// `*MigrateInfoEntry`'s
type EntriesBySize []*MigrateInfoEntry

// Len returns the total length of the set of `*MigrateInfoEntry`'s.
func (e EntriesBySize) Len() int { return len(e) }

// Less returns the whether or not the MigrateInfoEntry given at `i` takes up
// less total size than the MigrateInfoEntry given at `j`.
func (e EntriesBySize) Less(i, j int) bool { return e[i].BytesAbove < e[j].BytesAbove }

// Swap swaps the entries given at i, j.
func (e EntriesBySize) Swap(i, j int) { e[i], e[j] = e[j], e[i] }

// Print formats the `*MigrateInfoEntry`'s in the set and prints them to the
// given io.Writer, "to", returning "n" the number of bytes written, and any
// error, if one occurred.
func (e EntriesBySize) Print(to io.Writer) (int, error) {
	extensions := make([]string, 0, len(e))
	for _, entry := range e {
		extensions = append(extensions, entry.Qualifier)
	}
	extensions = tools.Ljust(extensions)

	files := make([]string, 0, len(e))
	for _, entry := range e {
		bytes := humanize.FormatBytes(uint64(entry.BytesAbove))
		above := entry.TotalAbove
		total := entry.Total

		file := fmt.Sprintf("%s, %d/%d files(s)",
			bytes, above, total)

		files = append(files, file)
	}
	files = tools.Rjust(files)

	percentages := make([]string, 0, len(e))
	for _, entry := range e {
		percentAbove := 100 * (float64(entry.TotalAbove) / float64(entry.Total))

		percentage := fmt.Sprintf("%.0f%%", percentAbove)

		percentages = append(percentages, percentage)
	}
	percentages = tools.Rjust(percentages)

	output := make([]string, 0, len(e))
	for i := 0; i < len(e); i++ {
		extension := extensions[i]
		fileCount := files[i]
		percentage := percentages[i]

		line := strings.Join([]string{extension, fileCount, percentage}, "\t")

		output = append(output, line)
	}

	header := fmt.Sprintf("Files above %s:", humanize.FormatBytes(migrateInfoAbove))
	output = append([]string{header}, output...)

	return fmt.Fprintln(to, strings.Join(output, "\n"))
}
<commit_msg>commands/command_migrate_info: build cols all at once<commit_after>package commands

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

	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/git/odb"
	"github.com/git-lfs/git-lfs/tools"
	"github.com/git-lfs/git-lfs/tools/humanize"
	"github.com/spf13/cobra"
)

var (
	// migrateInfoTopN is a flag given to the git-lfs-migrate(1) subcommand
	// 'info' which specifies how many info entries to show by default.
	migrateInfoTopN int

	// migrateInfoAboveFmt is a flag given to the git-lfs-migrate(1)
	// subcommand 'info' specifying a human-readable string threshold of
	// filesize before entries are counted.
	migrateInfoAboveFmt string
	// migrateInfoAbove is the number of bytes parsed from the above
	// migrateInfoAboveFmt flag.
	migrateInfoAbove uint64
)

func migrateInfoCommand(cmd *cobra.Command, args []string) {
	exts := make(map[string]*MigrateInfoEntry)

	above, err := humanize.ParseBytes(migrateInfoAboveFmt)
	if err != nil {
		ExitWithError(errors.Wrap(err, "cannot parse --above=<n>"))
	}

	migrateInfoAbove = above

	migrate(cmd, args, func(path string, b *odb.Blob) (*odb.Blob, error) {
		ext := fmt.Sprintf("*%s", filepath.Ext(path))

		if len(ext) > 1 {
			entry := exts[ext]
			if entry == nil {
				entry = &MigrateInfoEntry{Qualifier: ext}
			}

			entry.Total++
			entry.BytesTotal += b.Size

			if b.Size > int64(migrateInfoAbove) {
				entry.TotalAbove++
				entry.BytesAbove += b.Size
			}

			exts[ext] = entry
		}

		return b, nil
	})

	entries := EntriesBySize(MapToEntries(exts))
	sort.Sort(sort.Reverse(entries))

	migrateInfoTopN = tools.ClampInt(migrateInfoTopN, len(entries), 0)

	entries = entries[:tools.MaxInt(0, migrateInfoTopN)]

	entries.Print(os.Stderr)
}

// MigrateInfoEntry represents a tuple of filetype to bytes and entry count
// above and below a threshold.
type MigrateInfoEntry struct {
	// Qualifier is the filepath's extension.
	Qualifier string

	// BytesAbove is total size of all files above a given threshold.
	BytesAbove int64
	// TotalAbove is the count of all files above a given size threshold.
	TotalAbove int64
	// BytesTotal is the number of bytes of all files
	BytesTotal int64
	// Total is the count of all files.
	Total int64
}

// MapToEntries creates a set of `*MigrateInfoEntry`'s for a given map of
// filepath extensions to file size in bytes.
func MapToEntries(exts map[string]*MigrateInfoEntry) []*MigrateInfoEntry {
	entries := make([]*MigrateInfoEntry, 0, len(exts))
	for _, entry := range exts {
		entries = append(entries, entry)
	}

	return entries
}

// EntriesBySize is an implementation of sort.Interface that sorts a set of
// `*MigrateInfoEntry`'s
type EntriesBySize []*MigrateInfoEntry

// Len returns the total length of the set of `*MigrateInfoEntry`'s.
func (e EntriesBySize) Len() int { return len(e) }

// Less returns the whether or not the MigrateInfoEntry given at `i` takes up
// less total size than the MigrateInfoEntry given at `j`.
func (e EntriesBySize) Less(i, j int) bool { return e[i].BytesAbove < e[j].BytesAbove }

// Swap swaps the entries given at i, j.
func (e EntriesBySize) Swap(i, j int) { e[i], e[j] = e[j], e[i] }

// Print formats the `*MigrateInfoEntry`'s in the set and prints them to the
// given io.Writer, "to", returning "n" the number of bytes written, and any
// error, if one occurred.
func (e EntriesBySize) Print(to io.Writer) (int, error) {
	extensions := make([]string, 0, len(e))
	files := make([]string, 0, len(e))
	percentages := make([]string, 0, len(e))

	for _, entry := range e {
		bytes := humanize.FormatBytes(uint64(entry.BytesAbove))
		above := entry.TotalAbove
		total := entry.Total
		percentAbove := 100 * (float64(above) / float64(total))

		file := fmt.Sprintf("%s, %d/%d files(s)",
			bytes, above, total)

		percentage := fmt.Sprintf("%.0f%%", percentAbove)

		extensions = append(extensions, entry.Qualifier)
		files = append(files, file)
		percentages = append(percentages, percentage)
	}

	extensions = tools.Ljust(extensions)
	files = tools.Rjust(files)
	percentages = tools.Rjust(percentages)

	output := make([]string, 0, len(e))
	for i := 0; i < len(e); i++ {
		extension := extensions[i]
		fileCount := files[i]
		percentage := percentages[i]

		line := strings.Join([]string{extension, fileCount, percentage}, "\t")

		output = append(output, line)
	}

	header := fmt.Sprintf("Files above %s:", humanize.FormatBytes(migrateInfoAbove))
	output = append([]string{header}, output...)

	return fmt.Fprintln(to, strings.Join(output, "\n"))
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"strconv"
	"time"

	"github.com/pachyderm/pachyderm/src/client/limit"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
	"golang.org/x/sync/errgroup"
)

const (
	// MB is a megabyte
	MB = 1024 * 1024
)

var (
	// Flags

	// Number of objects.
	// (bryce) change to int64 later.
	numObjects int
	// Size of the objects.
	// (bryce) change to int64 later.
	objectSize int
	// Maximum concurrent writes.
	concurrency int
)

func init() {
	flag.IntVar(&numObjects, "num-objects", 10, "number of objects")
	flag.IntVar(&objectSize, "object-size", MB, "size of the objects")
	flag.IntVar(&concurrency, "concurrency", 5, "maximum concurrent writes")
}

// PrintFlags just prints the flag values, set above, to stdout. Useful for
// comparing benchmark runs
func PrintFlags() {
	fmt.Printf("num-objects: %v\n", numObjects)
	fmt.Printf("object-size: %v\n", objectSize)
	fmt.Printf("concurrency: %v\n", concurrency)
}

var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

// RandSeq generates a random sequence of data (n is number of bytes)
func RandSeq(n int) []byte {
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return []byte(string(b))
}

func main() {
	flag.Parse()
	PrintFlags()
	// Setup object to write and client.
	data := RandSeq(objectSize)
	c, err := obj.NewClientFromEnv("/pach")
	if err != nil {
		log.Fatalf("Error creating client (%v)", err)
	}
	// Start timing load test.
	var start = time.Now()
	defer func() {
		fmt.Printf("Benchmark complete. Total time: %.3f\n", time.Now().Sub(start).Seconds())
	}()
	// Start writing objects.
	eg, ctx := errgroup.WithContext(context.Background())
	limiter := limit.New(concurrency)
	for i := 0; i < numObjects; i++ {
		i := i
		limiter.Acquire()
		eg.Go(func() error {
			defer limiter.Release()
			w, err := c.Writer(ctx, strconv.Itoa(i))
			if err != nil {
				return fmt.Errorf("Error creating writer for object %v (%v)", i, err)
			}
			r := bytes.NewReader(data)
			if _, err := io.Copy(w, r); err != nil {
				return fmt.Errorf("Error writing to object %v (%v)", i, err)
			}
			if err := w.Close(); err != nil {
				return fmt.Errorf("Error closing object %v (%v)", i, err)
			}
			return nil
		})
	}
	if err := eg.Wait(); err != nil {
		log.Fatalf(err.Error())
	}
}
<commit_msg>Upgrade obj load test<commit_after>package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"strconv"
	"time"

	"github.com/pachyderm/pachyderm/src/client/limit"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
	"golang.org/x/sync/errgroup"
)

const (
	// MB is a megabyte
	MB     = 1024 * 1024
	prefix = "/pach"
)

var (
	// Flags
	// Number of objects.
	// (bryce) change to int64 later.
	numObjects int
	// Size of the objects.
	// (bryce) change to int64 later.
	objectSize int
	// Maximum concurrent writes.
	concurrency int
)

func init() {
	flag.IntVar(&numObjects, "num-objects", 10, "number of objects")
	flag.IntVar(&objectSize, "object-size", MB, "size of the objects")
	flag.IntVar(&concurrency, "concurrency", 5, "maximum concurrent writes")
}

// PrintFlags just prints the flag values, set above, to stdout. Useful for
// comparing benchmark runs
func PrintFlags() {
	fmt.Printf("num-objects: %v\n", numObjects)
	fmt.Printf("object-size: %v\n", objectSize)
	fmt.Printf("concurrency: %v\n", concurrency)
}

var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

// RandSeq generates a random sequence of data (n is number of bytes)
func RandSeq(n int) []byte {
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return []byte(string(b))
}

func main() {
	flag.Parse()
	PrintFlags()
	// Setup client.
	c, err := obj.NewClientFromEnv(prefix)
	if err != nil {
		log.Fatalf("Error creating client (%v)", err)
	}
	// (bryce) It might make sense to clean up the bucket here before running the tests.
	// Run basic test.
	start := time.Now()
	fmt.Printf("Basic test started.\n")
	if err := basicTest(c); err != nil {
		log.Fatalf("Basic test error: %v", err)
	}
	fmt.Printf("Basic test completed. Total time: %.3f\n", time.Now().Sub(start).Seconds())
	// Run load test.
	start = time.Now()
	fmt.Printf("Load test started.\n")
	if err := loadTest(c); err != nil {
		log.Fatalf("Load test error: %v", err)
	}
	fmt.Printf("Load test completed. Total time: %.3f\n", time.Now().Sub(start).Seconds())
}

func basicTest(c obj.Client) error {
	ctx := context.Background()
	name := "0"
	// Confirm that an existence check and deletion for a non-existent object works correctly.
	if c.Exists(ctx, name) {
		return fmt.Errorf("existence check returns true when the object should not exist")
	}
	if err := c.Delete(ctx, name); err != nil {
		return fmt.Errorf("deletion errored on non-existent object (%v)", err)
	}
	if err := walk(ctx, c, 0, nil); err != nil {
		return err
	}
	numObjects := 5
	basicObjectSize := 1024
	data := RandSeq(basicObjectSize)
	// Write then read objects.
	for i := 0; i < numObjects; i++ {
		name := strconv.Itoa(i)
		if err := writeObject(ctx, c, name, data); err != nil {
			return err
		}
		if err := readTest(ctx, c, name, 0, 0, data); err != nil {
			return err
		}
	}
	// Confirm range reads work correctly
	offset, size := basicObjectSize/2, 0
	if err := readTest(ctx, c, name, offset, size, data[offset:]); err != nil {
		return err
	}
	offset, size = basicObjectSize/2, basicObjectSize/4
	if err := readTest(ctx, c, name, offset, size, data[offset:offset+size]); err != nil {
		return err
	}
	// Walk the objects and for each check the existence and delete it.
	if err := walk(ctx, c, 5, func(name string) error {
		if !c.Exists(ctx, name) {
			return fmt.Errorf("existence check returns false when the object should exist")
		}
		return c.Delete(ctx, name)
	}); err != nil {
		return err
	}
	// Confirm that no objects exist after deletion.
	return walk(ctx, c, 0, nil)
}

func walk(ctx context.Context, c obj.Client, expected int, f func(string) error) error {
	objCount := 0
	if err := c.Walk(ctx, "", func(name string) error {
		objCount++
		if f != nil {
			return f(name)
		}
		return nil
	}); err != nil {
		return err
	}
	if objCount != expected {
		return fmt.Errorf("walk should have returned %v objects, not %v", expected, objCount)
	}
	return nil
}

func writeObject(ctx context.Context, c obj.Client, name string, data []byte) (retErr error) {
	w, err := c.Writer(ctx, name)
	if err != nil {
		return err
	}
	defer func() {
		if err := w.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	r := bytes.NewReader(data)
	_, err = io.Copy(w, r)
	return err
}

func readObject(ctx context.Context, c obj.Client, name string, offset, size int, buf []byte) (retErr error) {
	r, err := c.Reader(ctx, name, uint64(offset), uint64(size))
	if err != nil {
		return err
	}
	defer func() {
		if err := r.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	_, err = io.ReadFull(r, buf)
	return err
}

func readTest(ctx context.Context, c obj.Client, name string, offset, size int, expected []byte) error {
	buf := make([]byte, len(expected))
	if err := readObject(ctx, c, name, offset, size, buf); err != nil {
		return err
	}
	if bytes.Compare(expected, buf) != 0 {
		return fmt.Errorf("range read for object %v incorrect (offset: %v, size: %v)", name, offset, size)
	}
	return nil
}

func loadTest(c obj.Client) error {
	limiter := limit.New(concurrency)
	eg, ctx := errgroup.WithContext(context.Background())
	data := RandSeq(objectSize)
	bufPool := grpcutil.NewBufPool(objectSize)
	for i := 0; i < numObjects; i++ {
		i := i
		limiter.Acquire()
		eg.Go(func() error {
			defer limiter.Release()
			name := strconv.Itoa(i)
			if err := writeObject(ctx, c, name, data); err != nil {
				return err
			}
			buf := bufPool.GetBuffer()
			defer bufPool.PutBuffer(buf)
			if err := readObject(ctx, c, name, 0, 0, buf); err != nil {
				return err
			}
			if bytes.Compare(data, buf) != 0 {
				return fmt.Errorf("data writen does not equal data read for object %v", i)
			}
			return c.Delete(ctx, name)
		})
	}
	return eg.Wait()
}
<|endoftext|>


<commit_before><commit_msg>"Jeff: Give yourself to the Dark Side. It is the only way you can save your friends."<commit_after><|endoftext|>


<commit_before>package worker

import (
	"context"
	"encoding/json"
	"fmt"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
	"github.com/stretchr/testify/mock"
)

type mockedSqsClient struct {
	Config   *aws.Config
	Response sqs.ReceiveMessageOutput
	sqsiface.SQSAPI
}

func (c *mockedSqsClient) GetQueueUrl(urlInput *sqs.GetQueueUrlInput) (*sqs.GetQueueUrlOutput, error) {
	url := fmt.Sprintf("https://sqs.%v.amazonaws.com/123456789/%v", c.Config.Region, urlInput.QueueName)

	return &sqs.GetQueueUrlOutput{QueueUrl: &url}, nil
}

func (c *mockedSqsClient) ReceiveMessage(*sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error) {
	return &c.Response, nil
}

func (c *mockedSqsClient) DeleteMessage(*sqs.DeleteMessageInput) (*sqs.DeleteMessageOutput, error) {
	c.Response = sqs.ReceiveMessageOutput{}

	return &sqs.DeleteMessageOutput{}, nil
}

type mockedHandler struct {
	mock.Mock
}

type sqsEvent struct {
	Foo string `json:"foo"`
	Qux string `json:"qux"`
}

func (mh *mockedHandler) HandleMessage(foo string, qux string) {
	mh.Called(foo, qux)
}

func TestStart(t *testing.T) {
	region := "eu-west-1"
	awsConfig := &aws.Config{Region: &region}
	workerConfig := &Config{QueueName: "my-sqs-queue"}
	client := setupMockedSqsClient(awsConfig)
	worker := New(client, workerConfig)

	ctx, cancel := contextAndCancel()
	defer cancel()

	handler := new(mockedHandler)
	handlerFunc := HandlerFunc(func(msg *sqs.Message) (err error) {
		event := &sqsEvent{}

		json.Unmarshal([]byte(aws.StringValue(msg.Body)), event)

		handler.HandleMessage(event.Foo, event.Qux)

		return
	})

	t.Run("when worker successfully receives a message", func(t *testing.T) {
		handler.On("HandleMessage", "bar", "baz").Return().Once()
		worker.Start(ctx, handlerFunc)

		handler.AssertExpectations(t)
	})
}

func contextAndCancel() (context.Context, context.CancelFunc) {
	delay := time.Now().Add(1 * time.Millisecond)

	return context.WithDeadline(context.Background(), delay)
}

func setupMockedSqsClient(awsConfig *aws.Config) sqsiface.SQSAPI {
	sqsMessage := &sqs.Message{Body: aws.String(`{ "foo": "bar", "qux": "baz" }`)}
	sqsResponse := sqs.ReceiveMessageOutput{
		Messages: []*sqs.Message{sqsMessage},
	}

	return &mockedSqsClient{Response: sqsResponse, Config: awsConfig}
}
<commit_msg>Add tests for verifying proper sqs configuration<commit_after>package worker

import (
	"context"
	"encoding/json"
	"fmt"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
)

type mockedSqsClient struct {
	Config   *aws.Config
	Response sqs.ReceiveMessageOutput
	sqsiface.SQSAPI
	mock.Mock
}

func (c *mockedSqsClient) GetQueueUrl(urlInput *sqs.GetQueueUrlInput) (*sqs.GetQueueUrlOutput, error) {
	url := fmt.Sprintf("https://sqs.%v.amazonaws.com/123456789/%v", *c.Config.Region, *urlInput.QueueName)

	return &sqs.GetQueueUrlOutput{QueueUrl: &url}, nil
}

func (c *mockedSqsClient) ReceiveMessage(input *sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error) {
	c.Called(input)

	return &c.Response, nil
}

func (c *mockedSqsClient) DeleteMessage(input *sqs.DeleteMessageInput) (*sqs.DeleteMessageOutput, error) {
	c.Called(input)
	c.Response = sqs.ReceiveMessageOutput{}

	return &sqs.DeleteMessageOutput{}, nil
}

type mockedHandler struct {
	mock.Mock
}

func (mh *mockedHandler) HandleMessage(foo string, qux string) {
	mh.Called(foo, qux)
}

type sqsEvent struct {
	Foo string `json:"foo"`
	Qux string `json:"qux"`
}

func TestStart(t *testing.T) {
	region := "eu-west-1"
	awsConfig := &aws.Config{Region: &region}
	workerConfig := &Config{QueueName: "my-sqs-queue"}
	client := setupMockedSqsClient(awsConfig)
	worker := New(client, workerConfig)

	ctx, cancel := contextAndCancel()
	defer cancel()

	handler := new(mockedHandler)
	handlerFunc := HandlerFunc(func(msg *sqs.Message) (err error) {
		event := &sqsEvent{}

		json.Unmarshal([]byte(aws.StringValue(msg.Body)), event)

		handler.HandleMessage(event.Foo, event.Qux)

		return
	})

	t.Run("the worker has correct configuration", func(t *testing.T) {
		assert.Equal(t, worker.Config.QueueName, "my-sqs-queue", "QueueName has been set properly")
		assert.Equal(t, worker.Config.QueueURL, "https://sqs.eu-west-1.amazonaws.com/123456789/my-sqs-queue", "QueueURL has been set properly")
		assert.Equal(t, worker.Config.MaxNumberOfMessage, int64(10), "MaxNumberOfMessage has been set properly")
		assert.Equal(t, worker.Config.WaitTimeSecond, int64(20), "WaitTimeSecond has been set properly")
	})

	t.Run("the worker successfully processes a message", func(t *testing.T) {
		setupClientSpies(client)
		handler.On("HandleMessage", "bar", "baz").Return().Once()
		worker.Start(ctx, handlerFunc)

		client.AssertExpectations(t)
		handler.AssertExpectations(t)
	})
}

func contextAndCancel() (context.Context, context.CancelFunc) {
	delay := time.Now().Add(1 * time.Millisecond)

	return context.WithDeadline(context.Background(), delay)
}

func setupMockedSqsClient(awsConfig *aws.Config) *mockedSqsClient {
	sqsMessage := &sqs.Message{Body: aws.String(`{ "foo": "bar", "qux": "baz" }`)}
	sqsResponse := sqs.ReceiveMessageOutput{
		Messages: []*sqs.Message{sqsMessage},
	}

	return &mockedSqsClient{Response: sqsResponse, Config: awsConfig}
}

func setupClientSpies(client *mockedSqsClient) {
	url := aws.String("https://sqs.eu-west-1.amazonaws.com/123456789/my-sqs-queue")
	receiveInput := &sqs.ReceiveMessageInput{
		QueueUrl:            url,
		MaxNumberOfMessages: aws.Int64(10),
		AttributeNames: []*string{
			aws.String("All"),
		},
		WaitTimeSeconds: aws.Int64(10),
	}
	client.On("ReceiveMessage", receiveInput).Return()

	deleteInput := &sqs.DeleteMessageInput{
		QueueUrl: url,
	}
	client.On("DeleteMessage", deleteInput).Return()
}
<|endoftext|>


<commit_before>package docker

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"github.com/hashicorp/packer/packer"
)

// Windows containers are a special beast in Docker; you can't use docker cp
// to move files between the container and host.

// This communicator works around that limitation by reusing all possible
// methods and fields of the normal Docker Communicator, but we overwrite the
// Upload, Download, and UploadDir methods to utilize a mounted directory and
// native powershell commands rather than relying on docker cp.

type WindowsContainerCommunicator struct {
	Communicator
}

// Upload uses docker exec to copy the file from the host to the container
func (c *WindowsContainerCommunicator) Upload(dst string, src io.Reader, fi *os.FileInfo) error {
	// Create a temporary file to store the upload
	tempfile, err := ioutil.TempFile(c.HostDir, "upload")
	if err != nil {
		return err
	}
	defer os.Remove(tempfile.Name())

	// Copy the contents to the temporary file
	_, err = io.Copy(tempfile, src)
	if err != nil {
		return err
	}
	if fi != nil {
		tempfile.Chmod((*fi).Mode())
	}
	tempfile.Close()

	// Copy the file into place by copying the temporary file we put
	// into the shared folder into the proper location in the container
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("Copy-Item -Path %s/%s -Destination %s", c.ContainerDir,
			filepath.Base(tempfile.Name()), dst),
	}

	if err := c.Start(cmd); err != nil {
		return err
	}

	// Wait for the copy to complete
	cmd.Wait()
	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Upload failed with non-zero exit status: %d", cmd.ExitStatus)
	}

	return nil
}

func (c *WindowsContainerCommunicator) UploadDir(dst string, src string, exclude []string) error {
	// Create the temporary directory that will store the contents of "src"
	// for copying into the container.
	td, err := ioutil.TempDir(c.HostDir, "dirupload")
	if err != nil {
		return err
	}
	defer os.RemoveAll(td)

	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		relpath, err := filepath.Rel(src, path)
		if err != nil {
			return err
		}
		hostpath := filepath.Join(td, relpath)

		// If it is a directory, just create it
		if info.IsDir() {
			return os.MkdirAll(hostpath, info.Mode())
		}

		if info.Mode()&os.ModeSymlink == os.ModeSymlink {
			dest, err := os.Readlink(path)

			if err != nil {
				return err
			}

			return os.Symlink(dest, hostpath)
		}

		// It is a file, copy it over, including mode.
		src, err := os.Open(path)
		if err != nil {
			return err
		}
		defer src.Close()

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

		if _, err := io.Copy(dst, src); err != nil {
			return err
		}

		return nil
	}

	// Copy the entire directory tree to the temporary directory
	if err := filepath.Walk(src, walkFn); err != nil {
		return err
	}

	// Determine the destination directory
	containerSrc := filepath.Join(c.ContainerDir, filepath.Base(td))
	containerDst := dst
	if src[len(src)-1] != '/' {
		containerDst = filepath.Join(dst, filepath.Base(src))
	}

	// Make the directory, then copy into it
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("Copy-Item %s -Destination %s -Recurse",
			containerSrc, containerDst),
	}
	if err := c.Start(cmd); err != nil {
		return err
	}

	// Wait for the copy to complete
	cmd.Wait()
	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Upload failed with non-zero exit status: %d", cmd.ExitStatus)
	}

	return nil
}

// Download pulls a file out of a container using `docker cp`. We have a source
// path and want to write to an io.Writer
func (c *WindowsContainerCommunicator) Download(src string, dst io.Writer) error {
	log.Printf("Downloading file from container: %s:%s", c.ContainerID, src)
	// Copy file onto temp file on mounted volume inside container
	var stdout, stderr bytes.Buffer
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("Copy-Item -Path %s -Destination %s/%s", src, c.ContainerDir,
			filepath.Base(src)),
		Stdout: &stdout,
		Stderr: &stderr,
	}
	if err := c.Start(cmd); err != nil {
		return err
	}

	// Wait for the copy to complete
	cmd.Wait()

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Failed to copy file to shared drive: %s, %s, %d", stderr.String(), stdout.String(), cmd.ExitStatus)
	}

	// Read that copied file into a new file opened on host machine
	fsrc, err := os.Open(filepath.Join(c.HostDir, filepath.Base(src)))
	if err != nil {
		return err
	}
	defer fsrc.Close()
	defer os.Remove(fsrc.Name())

	_, err = io.Copy(dst, fsrc)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Update windows_container_communicator.go<commit_after>package docker

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"github.com/hashicorp/packer/packer"
)

// Windows containers are a special beast in Docker; you can't use docker cp
// to move files between the container and host.

// This communicator works around that limitation by reusing all possible
// methods and fields of the normal Docker Communicator, but we overwrite the
// Upload, Download, and UploadDir methods to utilize a mounted directory and
// native powershell commands rather than relying on docker cp.

type WindowsContainerCommunicator struct {
	Communicator
}

// Upload uses docker exec to copy the file from the host to the container
func (c *WindowsContainerCommunicator) Upload(dst string, src io.Reader, fi *os.FileInfo) error {
	// Create a temporary file to store the upload
	tempfile, err := ioutil.TempFile(c.HostDir, "upload")
	if err != nil {
		return err
	}
	defer os.Remove(tempfile.Name())

	// Copy the contents to the temporary file
	_, err = io.Copy(tempfile, src)
	if err != nil {
		return err
	}
	if fi != nil {
		tempfile.Chmod((*fi).Mode())
	}
	tempfile.Close()

	// Copy the file into place by copying the temporary file we put
	// into the shared folder into the proper location in the container
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("Copy-Item -Path %s/%s -Destination %s", c.ContainerDir,
			filepath.Base(tempfile.Name()), dst),
	}
	ctx := context.TODO()
	if err := c.Start(ctx, cmd); err != nil {
		return err
	}

	// Wait for the copy to complete
	cmd.Wait()
	if cmd.ExitStatus() != 0 {
		return fmt.Errorf("Upload failed with non-zero exit status: %d", cmd.ExitStatus())
	}

	return nil
}

func (c *WindowsContainerCommunicator) UploadDir(dst string, src string, exclude []string) error {
	// Create the temporary directory that will store the contents of "src"
	// for copying into the container.
	td, err := ioutil.TempDir(c.HostDir, "dirupload")
	if err != nil {
		return err
	}
	defer os.RemoveAll(td)

	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		relpath, err := filepath.Rel(src, path)
		if err != nil {
			return err
		}
		hostpath := filepath.Join(td, relpath)

		// If it is a directory, just create it
		if info.IsDir() {
			return os.MkdirAll(hostpath, info.Mode())
		}

		if info.Mode()&os.ModeSymlink == os.ModeSymlink {
			dest, err := os.Readlink(path)

			if err != nil {
				return err
			}

			return os.Symlink(dest, hostpath)
		}

		// It is a file, copy it over, including mode.
		src, err := os.Open(path)
		if err != nil {
			return err
		}
		defer src.Close()

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

		if _, err := io.Copy(dst, src); err != nil {
			return err
		}

		return nil
	}

	// Copy the entire directory tree to the temporary directory
	if err := filepath.Walk(src, walkFn); err != nil {
		return err
	}

	// Determine the destination directory
	containerSrc := filepath.Join(c.ContainerDir, filepath.Base(td))
	containerDst := dst
	if src[len(src)-1] != '/' {
		containerDst = filepath.Join(dst, filepath.Base(src))
	}

	// Make the directory, then copy into it
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("Copy-Item %s -Destination %s -Recurse",
			containerSrc, containerDst),
	}
	ctx := context.TODO()
	if err := c.Start(ctx, cmd); err != nil {
		return err
	}

	// Wait for the copy to complete
	cmd.Wait()
	if cmd.ExitStatus() != 0 {
		return fmt.Errorf("Upload failed with non-zero exit status: %d", cmd.ExitStatus())
	}

	return nil
}

// Download pulls a file out of a container using `docker cp`. We have a source
// path and want to write to an io.Writer
func (c *WindowsContainerCommunicator) Download(src string, dst io.Writer) error {
	log.Printf("Downloading file from container: %s:%s", c.ContainerID, src)
	// Copy file onto temp file on mounted volume inside container
	var stdout, stderr bytes.Buffer
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("Copy-Item -Path %s -Destination %s/%s", src, c.ContainerDir,
			filepath.Base(src)),
		Stdout: &stdout,
		Stderr: &stderr,
	}
	ctx := context.TODO()
	if err := c.Start(ctx, cmd); err != nil {
		return err
	}

	// Wait for the copy to complete
	cmd.Wait()

	if cmd.ExitStatus() != 0 {
		return fmt.Errorf("Failed to copy file to shared drive: %s, %s, %d", stderr.String(), stdout.String(), cmd.ExitStatus())
	}

	// Read that copied file into a new file opened on host machine
	fsrc, err := os.Open(filepath.Join(c.HostDir, filepath.Base(src)))
	if err != nil {
		return err
	}
	defer fsrc.Close()
	defer os.Remove(fsrc.Name())

	_, err = io.Copy(dst, fsrc)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package core

import (
	"errors"
	"fmt"
	"math/big"
	"sort"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/state"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/event"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"gopkg.in/fatih/set.v0"
)

var (
	ErrInvalidSender      = errors.New("Invalid sender")
	ErrNonce              = errors.New("Nonce too low")
	ErrBalance            = errors.New("Insufficient balance")
	ErrNonExistentAccount = errors.New("Account does not exist")
	ErrInsufficientFunds  = errors.New("Insufficient funds for gas * price + value")
	ErrIntrinsicGas       = errors.New("Intrinsic gas too low")
	ErrGasLimit           = errors.New("Exceeds block gas limit")
)

const txPoolQueueSize = 50

type TxPoolHook chan *types.Transaction
type TxMsg struct{ Tx *types.Transaction }

type stateFn func() *state.StateDB

const (
	minGasPrice = 1000000
)

type TxProcessor interface {
	ProcessTransaction(tx *types.Transaction)
}

// The tx pool a thread safe transaction pool handler. In order to
// guarantee a non blocking pool we use a queue channel which can be
// independently read without needing access to the actual pool.
type TxPool struct {
	mu sync.RWMutex
	// Queueing channel for reading and writing incoming
	// transactions to
	queueChan chan *types.Transaction
	// Quiting channel
	quit chan bool
	// The state function which will allow us to do some pre checkes
	currentState stateFn
	// The current gas limit function callback
	gasLimit func() *big.Int
	// The actual pool
	txs           map[common.Hash]*types.Transaction
	invalidHashes *set.Set

	queue map[common.Address]types.Transactions

	subscribers []chan TxMsg

	eventMux *event.TypeMux
}

func NewTxPool(eventMux *event.TypeMux, currentStateFn stateFn, gasLimitFn func() *big.Int) *TxPool {
	txPool := &TxPool{
		txs:           make(map[common.Hash]*types.Transaction),
		queue:         make(map[common.Address]types.Transactions),
		queueChan:     make(chan *types.Transaction, txPoolQueueSize),
		quit:          make(chan bool),
		eventMux:      eventMux,
		invalidHashes: set.New(),
		currentState:  currentStateFn,
		gasLimit:      gasLimitFn,
	}
	return txPool
}

func (pool *TxPool) Start() {
	// Queue timer will tick so we can attempt to move items from the queue to the
	// main transaction pool.
	queueTimer := time.NewTicker(300 * time.Millisecond)
	// Removal timer will tick and attempt to remove bad transactions (account.nonce>tx.nonce)
	removalTimer := time.NewTicker(1 * time.Second)
done:
	for {
		select {
		case <-queueTimer.C:
			pool.checkQueue()
		case <-removalTimer.C:
			pool.validatePool()
		case <-pool.quit:
			break done
		}
	}
}

func (pool *TxPool) ValidateTransaction(tx *types.Transaction) error {
	// Validate sender
	var (
		from common.Address
		err  error
	)

	if from, err = tx.From(); err != nil {
		return ErrInvalidSender
	}

	// Validate curve param
	v, _, _ := tx.Curve()
	if v > 28 || v < 27 {
		return fmt.Errorf("tx.v != (28 || 27) => %v", v)
	}

	if !pool.currentState().HasAccount(from) {
		return ErrNonExistentAccount
	}

	if pool.gasLimit().Cmp(tx.GasLimit) < 0 {
		return ErrGasLimit
	}

	total := new(big.Int).Mul(tx.Price, tx.GasLimit)
	total.Add(total, tx.Value())
	if pool.currentState().GetBalance(from).Cmp(total) < 0 {
		return ErrInsufficientFunds
	}

	if tx.GasLimit.Cmp(IntrinsicGas(tx)) < 0 {
		return ErrIntrinsicGas
	}

	if pool.currentState().GetNonce(from) > tx.Nonce() {
		return ErrNonce
	}

	return nil
}

func (self *TxPool) add(tx *types.Transaction) error {
	hash := tx.Hash()

	/* XXX I'm unsure about this. This is extremely dangerous and may result
	 in total black listing of certain transactions
	if self.invalidHashes.Has(hash) {
		return fmt.Errorf("Invalid transaction (%x)", hash[:4])
	}
	*/
	if self.txs[hash] != nil {
		return fmt.Errorf("Known transaction (%x)", hash[:4])
	}
	err := self.ValidateTransaction(tx)
	if err != nil {
		return err
	}

	self.queueTx(tx)

	var toname string
	if to := tx.To(); to != nil {
		toname = common.Bytes2Hex(to[:4])
	} else {
		toname = "[NEW_CONTRACT]"
	}
	// we can ignore the error here because From is
	// verified in ValidateTransaction.
	f, _ := tx.From()
	from := common.Bytes2Hex(f[:4])

	if glog.V(logger.Debug) {
		glog.Infof("(t) %x => %s (%v) %x\n", from, toname, tx.Value, tx.Hash())
	}

	return nil
}

func (self *TxPool) Size() int {
	return len(self.txs)
}

func (self *TxPool) Add(tx *types.Transaction) error {
	self.mu.Lock()
	defer self.mu.Unlock()

	return self.add(tx)
}

func (self *TxPool) AddTransactions(txs []*types.Transaction) {
	self.mu.Lock()
	defer self.mu.Unlock()

	for _, tx := range txs {
		if err := self.add(tx); err != nil {
			glog.V(logger.Debug).Infoln("tx error:", err)
		} else {
			h := tx.Hash()
			glog.V(logger.Debug).Infof("tx %x\n", h[:4])
		}
	}
}

func (self *TxPool) GetTransactions() (txs types.Transactions) {
	self.mu.RLock()
	defer self.mu.RUnlock()

	txs = make(types.Transactions, self.Size())
	i := 0
	for _, tx := range self.txs {
		txs[i] = tx
		i++
	}

	return
}

func (self *TxPool) GetQueuedTransactions() types.Transactions {
	self.mu.RLock()
	defer self.mu.RUnlock()

	var txs types.Transactions
	for _, ts := range self.queue {
		txs = append(txs, ts...)
	}

	return txs
}

func (self *TxPool) RemoveTransactions(txs types.Transactions) {
	self.mu.Lock()
	defer self.mu.Unlock()

	for _, tx := range txs {
		delete(self.txs, tx.Hash())
	}
}

func (pool *TxPool) Flush() {
	pool.txs = make(map[common.Hash]*types.Transaction)
}

func (pool *TxPool) Stop() {
	pool.Flush()
	close(pool.quit)

	glog.V(logger.Info).Infoln("TX Pool stopped")
}

func (self *TxPool) queueTx(tx *types.Transaction) {
	from, _ := tx.From()
	self.queue[from] = append(self.queue[from], tx)
}

func (pool *TxPool) addTx(tx *types.Transaction) {
	if _, ok := pool.txs[tx.Hash()]; !ok {
		pool.txs[tx.Hash()] = tx
		// Notify the subscribers. This event is posted in a goroutine
		// because it's possible that somewhere during the post "Remove transaction"
		// gets called which will then wait for the global tx pool lock and deadlock.
		go pool.eventMux.Post(TxPreEvent{tx})
	}
}

// check queue will attempt to insert
func (pool *TxPool) checkQueue() {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	statedb := pool.currentState()
	for address, txs := range pool.queue {
		sort.Sort(types.TxByNonce{txs})

		var (
			nonce = statedb.GetNonce(address)
			start int
		)
		// Clean up the transactions first and determine the start of the nonces
		for _, tx := range txs {
			if tx.Nonce() >= nonce {
				break
			}
			start++
		}
		pool.queue[address] = txs[start:]

		// expected nonce
		enonce := nonce
		for _, tx := range pool.queue[address] {
			// If the expected nonce does not match up with the next one
			// (i.e. a nonce gap), we stop the loop
			if enonce != tx.Nonce() {
				break
			}
			enonce++

			pool.addTx(tx)
		}
		// delete the entire queue entry if it's empty. There's no need to keep it
		if len(pool.queue[address]) == 0 {
			delete(pool.queue, address)
		}
	}
}

func (pool *TxPool) removeTx(hash common.Hash) {
	// delete from pending pool
	delete(pool.txs, hash)

	// delete from queue
out:
	for address, txs := range pool.queue {
		for i, tx := range txs {
			if tx.Hash() == hash {
				if len(txs) == 1 {
					// if only one tx, remove entire address entry
					delete(pool.queue, address)
				} else {
					pool.queue[address][len(txs)-1], pool.queue[address] = nil, append(txs[:i], txs[i+1:]...)
				}
				break out
			}
		}
	}
}

func (pool *TxPool) validatePool() {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	for hash, tx := range pool.txs {
		if err := pool.ValidateTransaction(tx); err != nil {
			if glog.V(logger.Info) {
				glog.Infof("removed tx (%x) from pool: %v\n", hash[:4], err)
			}

			pool.removeTx(hash)
		}
	}
}
<commit_msg>core: use removeTx instead of delete<commit_after>package core

import (
	"errors"
	"fmt"
	"math/big"
	"sort"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/state"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/event"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"gopkg.in/fatih/set.v0"
)

var (
	ErrInvalidSender      = errors.New("Invalid sender")
	ErrNonce              = errors.New("Nonce too low")
	ErrBalance            = errors.New("Insufficient balance")
	ErrNonExistentAccount = errors.New("Account does not exist")
	ErrInsufficientFunds  = errors.New("Insufficient funds for gas * price + value")
	ErrIntrinsicGas       = errors.New("Intrinsic gas too low")
	ErrGasLimit           = errors.New("Exceeds block gas limit")
)

const txPoolQueueSize = 50

type TxPoolHook chan *types.Transaction
type TxMsg struct{ Tx *types.Transaction }

type stateFn func() *state.StateDB

const (
	minGasPrice = 1000000
)

type TxProcessor interface {
	ProcessTransaction(tx *types.Transaction)
}

// The tx pool a thread safe transaction pool handler. In order to
// guarantee a non blocking pool we use a queue channel which can be
// independently read without needing access to the actual pool.
type TxPool struct {
	mu sync.RWMutex
	// Queueing channel for reading and writing incoming
	// transactions to
	queueChan chan *types.Transaction
	// Quiting channel
	quit chan bool
	// The state function which will allow us to do some pre checkes
	currentState stateFn
	// The current gas limit function callback
	gasLimit func() *big.Int
	// The actual pool
	txs           map[common.Hash]*types.Transaction
	invalidHashes *set.Set

	queue map[common.Address]types.Transactions

	subscribers []chan TxMsg

	eventMux *event.TypeMux
}

func NewTxPool(eventMux *event.TypeMux, currentStateFn stateFn, gasLimitFn func() *big.Int) *TxPool {
	txPool := &TxPool{
		txs:           make(map[common.Hash]*types.Transaction),
		queue:         make(map[common.Address]types.Transactions),
		queueChan:     make(chan *types.Transaction, txPoolQueueSize),
		quit:          make(chan bool),
		eventMux:      eventMux,
		invalidHashes: set.New(),
		currentState:  currentStateFn,
		gasLimit:      gasLimitFn,
	}
	return txPool
}

func (pool *TxPool) Start() {
	// Queue timer will tick so we can attempt to move items from the queue to the
	// main transaction pool.
	queueTimer := time.NewTicker(300 * time.Millisecond)
	// Removal timer will tick and attempt to remove bad transactions (account.nonce>tx.nonce)
	removalTimer := time.NewTicker(1 * time.Second)
done:
	for {
		select {
		case <-queueTimer.C:
			pool.checkQueue()
		case <-removalTimer.C:
			pool.validatePool()
		case <-pool.quit:
			break done
		}
	}
}

func (pool *TxPool) ValidateTransaction(tx *types.Transaction) error {
	// Validate sender
	var (
		from common.Address
		err  error
	)

	if from, err = tx.From(); err != nil {
		return ErrInvalidSender
	}

	// Validate curve param
	v, _, _ := tx.Curve()
	if v > 28 || v < 27 {
		return fmt.Errorf("tx.v != (28 || 27) => %v", v)
	}

	if !pool.currentState().HasAccount(from) {
		return ErrNonExistentAccount
	}

	if pool.gasLimit().Cmp(tx.GasLimit) < 0 {
		return ErrGasLimit
	}

	total := new(big.Int).Mul(tx.Price, tx.GasLimit)
	total.Add(total, tx.Value())
	if pool.currentState().GetBalance(from).Cmp(total) < 0 {
		return ErrInsufficientFunds
	}

	if tx.GasLimit.Cmp(IntrinsicGas(tx)) < 0 {
		return ErrIntrinsicGas
	}

	if pool.currentState().GetNonce(from) > tx.Nonce() {
		return ErrNonce
	}

	return nil
}

func (self *TxPool) add(tx *types.Transaction) error {
	hash := tx.Hash()

	/* XXX I'm unsure about this. This is extremely dangerous and may result
	 in total black listing of certain transactions
	if self.invalidHashes.Has(hash) {
		return fmt.Errorf("Invalid transaction (%x)", hash[:4])
	}
	*/
	if self.txs[hash] != nil {
		return fmt.Errorf("Known transaction (%x)", hash[:4])
	}
	err := self.ValidateTransaction(tx)
	if err != nil {
		return err
	}

	self.queueTx(tx)

	var toname string
	if to := tx.To(); to != nil {
		toname = common.Bytes2Hex(to[:4])
	} else {
		toname = "[NEW_CONTRACT]"
	}
	// we can ignore the error here because From is
	// verified in ValidateTransaction.
	f, _ := tx.From()
	from := common.Bytes2Hex(f[:4])

	if glog.V(logger.Debug) {
		glog.Infof("(t) %x => %s (%v) %x\n", from, toname, tx.Value, tx.Hash())
	}

	return nil
}

func (self *TxPool) Size() int {
	return len(self.txs)
}

func (self *TxPool) Add(tx *types.Transaction) error {
	self.mu.Lock()
	defer self.mu.Unlock()

	return self.add(tx)
}

func (self *TxPool) AddTransactions(txs []*types.Transaction) {
	self.mu.Lock()
	defer self.mu.Unlock()

	for _, tx := range txs {
		if err := self.add(tx); err != nil {
			glog.V(logger.Debug).Infoln("tx error:", err)
		} else {
			h := tx.Hash()
			glog.V(logger.Debug).Infof("tx %x\n", h[:4])
		}
	}
}

func (self *TxPool) GetTransactions() (txs types.Transactions) {
	self.mu.RLock()
	defer self.mu.RUnlock()

	txs = make(types.Transactions, self.Size())
	i := 0
	for _, tx := range self.txs {
		txs[i] = tx
		i++
	}

	return
}

func (self *TxPool) GetQueuedTransactions() types.Transactions {
	self.mu.RLock()
	defer self.mu.RUnlock()

	var txs types.Transactions
	for _, ts := range self.queue {
		txs = append(txs, ts...)
	}

	return txs
}

func (self *TxPool) RemoveTransactions(txs types.Transactions) {
	self.mu.Lock()
	defer self.mu.Unlock()

	for _, tx := range txs {
		self.removeTx(tx.Hash())
	}
}

func (pool *TxPool) Flush() {
	pool.txs = make(map[common.Hash]*types.Transaction)
}

func (pool *TxPool) Stop() {
	pool.Flush()
	close(pool.quit)

	glog.V(logger.Info).Infoln("TX Pool stopped")
}

func (self *TxPool) queueTx(tx *types.Transaction) {
	from, _ := tx.From()
	self.queue[from] = append(self.queue[from], tx)
}

func (pool *TxPool) addTx(tx *types.Transaction) {
	if _, ok := pool.txs[tx.Hash()]; !ok {
		pool.txs[tx.Hash()] = tx
		// Notify the subscribers. This event is posted in a goroutine
		// because it's possible that somewhere during the post "Remove transaction"
		// gets called which will then wait for the global tx pool lock and deadlock.
		go pool.eventMux.Post(TxPreEvent{tx})
	}
}

// check queue will attempt to insert
func (pool *TxPool) checkQueue() {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	statedb := pool.currentState()
	for address, txs := range pool.queue {
		sort.Sort(types.TxByNonce{txs})

		var (
			nonce = statedb.GetNonce(address)
			start int
		)
		// Clean up the transactions first and determine the start of the nonces
		for _, tx := range txs {
			if tx.Nonce() >= nonce {
				break
			}
			start++
		}
		pool.queue[address] = txs[start:]

		// expected nonce
		enonce := nonce
		for _, tx := range pool.queue[address] {
			// If the expected nonce does not match up with the next one
			// (i.e. a nonce gap), we stop the loop
			if enonce != tx.Nonce() {
				break
			}
			enonce++

			pool.addTx(tx)
		}
		// delete the entire queue entry if it's empty. There's no need to keep it
		if len(pool.queue[address]) == 0 {
			delete(pool.queue, address)
		}
	}
}

func (pool *TxPool) removeTx(hash common.Hash) {
	// delete from pending pool
	delete(pool.txs, hash)

	// delete from queue
out:
	for address, txs := range pool.queue {
		for i, tx := range txs {
			if tx.Hash() == hash {
				if len(txs) == 1 {
					// if only one tx, remove entire address entry
					delete(pool.queue, address)
				} else {
					pool.queue[address][len(txs)-1], pool.queue[address] = nil, append(txs[:i], txs[i+1:]...)
				}
				break out
			}
		}
	}
}

func (pool *TxPool) validatePool() {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	for hash, tx := range pool.txs {
		if err := pool.ValidateTransaction(tx); err != nil {
			if glog.V(logger.Info) {
				glog.Infof("removed tx (%x) from pool: %v\n", hash[:4], err)
			}

			pool.removeTx(hash)
		}
	}
}
<|endoftext|>


<commit_before>package cassandra

import (
	"fmt"
	bulkQuerygen "github.com/influxdata/influxdb-comparisons/bulk_query_gen"
	"math/rand"
	"time"
)

// CassandraIot produces Cassandra-specific queries for all the devops query types.
type CassandraIot struct {
	KeyspaceName string
	AllInterval  bulkQuerygen.TimeInterval
}

// NewCassandraIot makes an CassandraIot object ready to generate Queries.
func newCassandraIotCommon(dbConfig bulkQuerygen.DatabaseConfig, start, end time.Time) bulkQuerygen.QueryGenerator {
	if !start.Before(end) {
		panic("bad time order")
	}

	return &CassandraIot{
		KeyspaceName: dbConfig["database-name"],
		AllInterval:  bulkQuerygen.NewTimeInterval(start, end),
	}
}

// Dispatch fulfills the QueryGenerator interface.
func (d *CassandraIot) Dispatch(i, scaleVar int) bulkQuerygen.Query {
	q := NewCassandraQuery() // from pool
	bulkQuerygen.IotDispatchAll(d, i, q, scaleVar)
	return q
}

func (d *CassandraIot) AverageTemperatureDayByHourOneHome(q bulkQuerygen.Query, scaleVar int) {
	d.averageTemperatureDayByHourNHomes(q.(*CassandraQuery), scaleVar, 1, time.Hour)
}

// averageTemperatureHourByMinuteNHomes populates a Query with a query that looks like:
// SELECT avg(temperature) from air_condition_room where (home_id = '$HHOME_ID_1' or ... or hostname = '$HOSTNAME_N') and time >= '$HOUR_START' and time < '$HOUR_END' group by time(1h)
func (d *CassandraIot) averageTemperatureDayByHourNHomes(qi bulkQuerygen.Query, scaleVar, nHomes int, timeRange time.Duration) {
	interval := d.AllInterval.RandWindow(timeRange)
	nn := rand.Perm(scaleVar)[:nHomes]

	tagSets := [][]string{}
	tagSet := []string{}
	for _, n := range nn {
		hostname := fmt.Sprintf("host_%d", n)
		tag := fmt.Sprintf("hostname=%s", hostname)
		tagSet = append(tagSet, tag)
	}
	tagSets = append(tagSets, tagSet)

	humanLabel := fmt.Sprintf("Cassandra average temperature, rand %4d homes, rand %s by 1h", nHomes, timeRange)
	q := qi.(*CassandraQuery)
	q.HumanLabel = []byte(humanLabel)
	q.HumanDescription = []byte(fmt.Sprintf("%s: %s", humanLabel, interval.StartString()))

	q.AggregationType = []byte("avg")
	q.MeasurementName = []byte("air_condition_room")
	q.FieldName = []byte("temperature")

	q.TimeStart = interval.Start
	q.TimeEnd = interval.End
	q.GroupByDuration = time.Hour

	q.TagSets = tagSets
}
<commit_msg>Fixed tag names for CQL IoT queries<commit_after>package cassandra

import (
	"fmt"
	bulkDataGenIot "github.com/influxdata/influxdb-comparisons/bulk_data_gen/iot"
	bulkQuerygen "github.com/influxdata/influxdb-comparisons/bulk_query_gen"
	"math/rand"
	"time"
)

// CassandraIot produces Cassandra-specific queries for all the devops query types.
type CassandraIot struct {
	KeyspaceName string
	AllInterval  bulkQuerygen.TimeInterval
}

// NewCassandraIot makes an CassandraIot object ready to generate Queries.
func newCassandraIotCommon(dbConfig bulkQuerygen.DatabaseConfig, start, end time.Time) bulkQuerygen.QueryGenerator {
	if !start.Before(end) {
		panic("bad time order")
	}

	return &CassandraIot{
		KeyspaceName: dbConfig["database-name"],
		AllInterval:  bulkQuerygen.NewTimeInterval(start, end),
	}
}

// Dispatch fulfills the QueryGenerator interface.
func (d *CassandraIot) Dispatch(i, scaleVar int) bulkQuerygen.Query {
	q := NewCassandraQuery() // from pool
	bulkQuerygen.IotDispatchAll(d, i, q, scaleVar)
	return q
}

func (d *CassandraIot) AverageTemperatureDayByHourOneHome(q bulkQuerygen.Query, scaleVar int) {
	d.averageTemperatureDayByHourNHomes(q.(*CassandraQuery), scaleVar, 1, time.Hour)
}

// averageTemperatureHourByMinuteNHomes populates a Query with a query that looks like:
// SELECT avg(temperature) from air_condition_room where (home_id = '$HHOME_ID_1' or ... or hostname = '$HOSTNAME_N') and time >= '$HOUR_START' and time < '$HOUR_END' group by time(1h)
func (d *CassandraIot) averageTemperatureDayByHourNHomes(qi bulkQuerygen.Query, scaleVar, nHomes int, timeRange time.Duration) {
	interval := d.AllInterval.RandWindow(timeRange)
	nn := rand.Perm(scaleVar)[:nHomes]

	tagSets := [][]string{}
	tagSet := []string{}
	for _, n := range nn {
		home := fmt.Sprintf(bulkDataGenIot.SmartHomeIdFormat, n)
		tag := fmt.Sprintf("home_id = %s", home)
		tagSet = append(tagSet, tag)
	}
	tagSets = append(tagSets, tagSet)

	humanLabel := fmt.Sprintf("Cassandra average temperature, rand %4d homes, rand %s by 1h", nHomes, timeRange)
	q := qi.(*CassandraQuery)
	q.HumanLabel = []byte(humanLabel)
	q.HumanDescription = []byte(fmt.Sprintf("%s: %s", humanLabel, interval.StartString()))

	q.AggregationType = []byte("avg")
	q.MeasurementName = []byte("air_condition_room")
	q.FieldName = []byte("temperature")

	q.TimeStart = interval.Start
	q.TimeEnd = interval.End
	q.GroupByDuration = time.Hour

	q.TagSets = tagSets
}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
  "io/ioutil"
  "github.com/hyperledger/fabric/protos/utils"
)

func main(){
  fmt.Println("Use codes to verify fabric")
  blockFile := "/root/go/src/github.com/hyperledger/fabric/examples/e2e_cli/channel-artifacts/genesis.block"
  data, err := ioutil.ReadFile(blockFile)
	if err != nil {
		fmt.Errorf("Could not read block %s", inspectBlock)
  }
  block, err := utils.UnmarshalBlock(data)
	if err != nil {
		fmt.Errorf("error unmarshaling to block: %s", err)
  }
  fmt.Println(utils.GetChainIDFromBlock(block))
  fmt.Println(" get what we want")
}
<commit_msg>Update main.go<commit_after>package main

import (
  "fmt"
  "io/ioutil"
  "github.com/hyperledger/fabric/protos/utils"
)

func main(){
  fmt.Println("Use codes to verify fabric")
  blockFile := "/root/go/src/github.com/hyperledger/fabric/examples/e2e_cli/channel-artifacts/genesis.block"
  data, err := ioutil.ReadFile(blockFile)
	if err != nil {
		fmt.Errorf("Could not read block %s", blockFile)
  }
  block, err := utils.UnmarshalBlock(data)
	if err != nil {
		fmt.Errorf("error unmarshaling to block: %s", err)
  }
  fmt.Println(utils.GetChainIDFromBlock(block))
  fmt.Println(" get what we want")
}
<|endoftext|>


<commit_before>package clusterdeploy

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"reflect"
	"sync"
	"time"

	"github.com/rancher/norman/types"
	util "github.com/rancher/rancher/pkg/cluster"
	"github.com/rancher/rancher/pkg/clustermanager"
	"github.com/rancher/rancher/pkg/features"
	"github.com/rancher/rancher/pkg/image"
	"github.com/rancher/rancher/pkg/kubectl"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/rancher/pkg/systemaccount"
	"github.com/rancher/rancher/pkg/systemtemplate"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/rancher/types/user"
	"github.com/sirupsen/logrus"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
)

const (
	AgentForceDeployAnn = "io.cattle.agent.force.deploy"
	nodeImage           = "nodeImage"
	clusterImage        = "clusterImage"
)

var (
	agentImagesMutex sync.RWMutex
	agentImages      = map[string]map[string]string{
		nodeImage:    map[string]string{},
		clusterImage: map[string]string{},
	}
)

func Register(ctx context.Context, management *config.ManagementContext, clusterManager *clustermanager.Manager) {
	c := &clusterDeploy{
		systemAccountManager: systemaccount.NewManager(management),
		userManager:          management.UserManager,
		clusters:             management.Management.Clusters(""),
		nodeLister:           management.Management.Nodes("").Controller().Lister(),
		clusterManager:       clusterManager,
	}

	management.Management.Clusters("").AddHandler(ctx, "cluster-deploy", c.sync)
}

type clusterDeploy struct {
	systemAccountManager *systemaccount.Manager
	userManager          user.Manager
	clusters             v3.ClusterInterface
	clusterManager       *clustermanager.Manager
	nodeLister           v3.NodeLister
}

func (cd *clusterDeploy) sync(key string, cluster *v3.Cluster) (runtime.Object, error) {
	var (
		err, updateErr error
	)

	if cluster == nil || cluster.DeletionTimestamp != nil {
		// remove the system account user created for this cluster
		if err := cd.systemAccountManager.RemoveSystemAccount(key); err != nil {
			return nil, err
		}
		return nil, nil
	}

	original := cluster
	cluster = original.DeepCopy()

	if cluster.Status.Driver == v3.ClusterDriverRKE {
		if cluster.Spec.LocalClusterAuthEndpoint.Enabled {
			cluster.Spec.RancherKubernetesEngineConfig.Authentication.Strategy = "x509|webhook"
		} else {
			cluster.Spec.RancherKubernetesEngineConfig.Authentication.Strategy = "x509"
		}
	}

	err = cd.doSync(cluster)
	if cluster != nil && !reflect.DeepEqual(cluster, original) {
		_, updateErr = cd.clusters.Update(cluster)
	}

	if err != nil {
		return nil, err
	}
	return nil, updateErr
}

func (cd *clusterDeploy) doSync(cluster *v3.Cluster) error {
	if !v3.ClusterConditionProvisioned.IsTrue(cluster) {
		return nil
	}

	nodes, err := cd.nodeLister.List(cluster.Name, labels.Everything())
	if err != nil {
		return err
	}
	if len(nodes) == 0 {
		return nil
	}

	_, err = v3.ClusterConditionSystemAccountCreated.DoUntilTrue(cluster, func() (runtime.Object, error) {
		return cluster, cd.systemAccountManager.CreateSystemAccount(cluster)
	})
	if err != nil {
		return err
	}

	if cluster.Status.AgentImage != "" && !agentImagesCached(cluster.Name) {
		if err := cd.cacheAgentImages(cluster.Name); err != nil {
			return err
		}
	}

	err = cd.deployAgent(cluster)
	if err != nil {
		return err
	}

	return cd.setNetworkPolicyAnn(cluster)
}

// agentFeaturesChanged will treat a missing key as false. This means we only detect changes
// when we set a feature to true so we can't reliably set a feature to false that is enabled by default.
// This behavior makes adding new def false features not cause the agent to redeploy.
func agentFeaturesChanged(desired, actual map[string]bool) bool {
	for k, v := range desired {
		if actual[k] != v {
			return true
		}
	}

	for k, v := range actual {
		if desired[k] != v {
			return true
		}
	}

	return false
}

func redeployAgent(cluster *v3.Cluster, desiredAgent, desiredAuth string, desiredFeatures map[string]bool) bool {
	if !v3.ClusterConditionAgentDeployed.IsTrue(cluster) {
		return true
	}

	forceDeploy := cluster.Annotations[AgentForceDeployAnn] == "true"
	imageChange := cluster.Status.AgentImage != desiredAgent || cluster.Status.AuthImage != desiredAuth
	agentFeaturesChanged := agentFeaturesChanged(desiredFeatures, cluster.Status.AgentFeatures)
	repoChange := false
	if cluster.Spec.RancherKubernetesEngineConfig != nil {
		if cluster.Status.AppliedSpec.RancherKubernetesEngineConfig != nil {
			desiredRepo := util.GetPrivateRepo(cluster)
			var appliedRepo *v3.PrivateRegistry
			if len(cluster.Status.AppliedSpec.RancherKubernetesEngineConfig.PrivateRegistries) > 0 {
				appliedRepo = &cluster.Status.AppliedSpec.RancherKubernetesEngineConfig.PrivateRegistries[0]
			}
			if desiredRepo != nil && appliedRepo != nil && !reflect.DeepEqual(desiredRepo, appliedRepo) {
				repoChange = true
			}
			if (desiredRepo == nil && appliedRepo != nil) || (desiredRepo != nil && appliedRepo == nil) {
				repoChange = true
			}
		}
	}

	if forceDeploy || imageChange || repoChange || agentFeaturesChanged {
		logrus.Infof("Redeploy Rancher Agents is needed for %s: forceDeploy=%v, agent/auth image changed=%v,"+
			" private repo changed=%v, agent features changed=%v", cluster.Name, forceDeploy, imageChange, repoChange,
			agentFeaturesChanged)
		return true
	}

	na, ca := getAgentImages(cluster.Name)
	if cluster.Status.AgentImage != na || cluster.Status.AgentImage != ca {
		// downstream agent does not match, kick a redeploy with settings agent
		logrus.Infof("Redeploy Rancher Agents due to Downstream Agent Image Mismatch for %s: was %s and will be %s",
			cluster.Name, na, image.ResolveWithCluster(settings.AgentImage.Get(), cluster))
		clearAgentImages(cluster.Name)
		return true
	}

	return false
}

func getDesiredImage(cluster *v3.Cluster) string {
	if cluster.Spec.AgentImageOverride != "" {
		return cluster.Spec.AgentImageOverride
	}

	return cluster.Spec.DesiredAgentImage
}

func (cd *clusterDeploy) deployAgent(cluster *v3.Cluster) error {
	desiredAgent := getDesiredImage(cluster)
	if desiredAgent == "" || desiredAgent == "fixed" {
		desiredAgent = image.ResolveWithCluster(settings.AgentImage.Get(), cluster)
	}

	var desiredAuth string
	if cluster.Spec.LocalClusterAuthEndpoint.Enabled {
		desiredAuth = cluster.Spec.DesiredAuthImage
		if desiredAuth == "" || desiredAuth == "fixed" {
			desiredAuth = image.ResolveWithCluster(settings.AuthImage.Get(), cluster)
		}
	}

	desiredFeatures := map[string]bool{
		features.Steve.Name(): features.Steve.Enabled(),
	}

	if !redeployAgent(cluster, desiredAgent, desiredAuth, desiredFeatures) {
		return nil
	}

	kubeConfig, err := cd.getKubeConfig(cluster)
	if err != nil {
		return err
	}

	if _, err = v3.ClusterConditionAgentDeployed.Do(cluster, func() (runtime.Object, error) {
		yaml, err := cd.getYAML(cluster, desiredAgent, desiredAuth, desiredFeatures)
		if err != nil {
			return cluster, err
		}
		var output []byte
		for i := 0; i < 3; i++ {
			// This will fail almost always the first time because when we create the namespace in the file
			// it won't have privileges.  Just stupidly try 3 times
			output, err = kubectl.Apply(yaml, kubeConfig)
			if err == nil {
				break
			}
			time.Sleep(2 * time.Second)
		}
		if err != nil {
			return cluster, types.NewErrors(err, errors.New(string(output)))
		}
		v3.ClusterConditionAgentDeployed.Message(cluster, string(output))
		if !cluster.Spec.LocalClusterAuthEndpoint.Enabled && cluster.Status.AppliedSpec.LocalClusterAuthEndpoint.Enabled && cluster.Status.AuthImage != "" {
			output, err = kubectl.Delete([]byte(systemtemplate.AuthDaemonSet), kubeConfig)
		}
		if err != nil {
			return cluster, types.NewErrors(err, errors.New(string(output)))
		}
		v3.ClusterConditionAgentDeployed.Message(cluster, string(output))
		return cluster, nil
	}); err != nil {
		return err
	}

	if err = cd.cacheAgentImages(cluster.Name); err != nil {
		return err
	}

	cluster.Status.AgentImage = desiredAgent
	cluster.Status.AgentFeatures = desiredFeatures
	if cluster.Spec.DesiredAgentImage == "fixed" {
		cluster.Spec.DesiredAgentImage = desiredAgent
	}
	cluster.Status.AuthImage = desiredAuth
	if cluster.Spec.DesiredAuthImage == "fixed" {
		cluster.Spec.DesiredAuthImage = desiredAuth
	}
	if cluster.Annotations[AgentForceDeployAnn] == "true" {
		cluster.Annotations[AgentForceDeployAnn] = "false"
	}

	return nil
}

func (cd *clusterDeploy) setNetworkPolicyAnn(cluster *v3.Cluster) error {
	if cluster.Spec.EnableNetworkPolicy != nil {
		return nil
	}
	// set current state for upgraded canal clusters
	if cluster.Spec.RancherKubernetesEngineConfig != nil &&
		cluster.Spec.RancherKubernetesEngineConfig.Network.Plugin == "canal" {
		enableNetworkPolicy := true
		cluster.Spec.EnableNetworkPolicy = &enableNetworkPolicy
		cluster.Annotations["networking.management.cattle.io/enable-network-policy"] = "true"
	}
	return nil
}

func (cd *clusterDeploy) getKubeConfig(cluster *v3.Cluster) (*clientcmdapi.Config, error) {
	user, err := cd.systemAccountManager.GetSystemUser(cluster.Name)
	if err != nil {
		return nil, err
	}

	token, err := cd.userManager.EnsureToken("agent-"+user.Name, "token for agent deployment", "agent", user.Name)
	if err != nil {
		return nil, err
	}

	return cd.clusterManager.KubeConfig(cluster.Name, token), nil
}

func (cd *clusterDeploy) getYAML(cluster *v3.Cluster, agentImage, authImage string, features map[string]bool) ([]byte, error) {
	logrus.Debug("Desired agent image:", agentImage)
	logrus.Debug("Desired auth image:", authImage)

	token, err := cd.systemAccountManager.GetOrCreateSystemClusterToken(cluster.Name)
	if err != nil {
		return nil, err
	}

	url := settings.ServerURL.Get()
	if url == "" {
		cd.clusters.Controller().EnqueueAfter("", cluster.Name, time.Second)
		return nil, fmt.Errorf("waiting for server-url setting to be set")
	}

	buf := &bytes.Buffer{}
	err = systemtemplate.SystemTemplate(buf, agentImage, authImage, cluster.Name, token, url, cluster.Spec.WindowsPreferedCluster,
		cluster, features)

	return buf.Bytes(), err
}

func (cd *clusterDeploy) getClusterAgentImage(name string) (string, error) {
	uc, err := cd.clusterManager.UserContext(name)
	if err != nil {
		return "", err
	}

	d, err := uc.Apps.Deployments("cattle-system").Get("cattle-cluster-agent", v1.GetOptions{})
	if err != nil {
		if !apierrors.IsNotFound(err) {
			return "", err
		}
		return "", nil
	}

	for _, c := range d.Spec.Template.Spec.Containers {
		if c.Name == "cluster-register" {
			return c.Image, nil
		}
	}

	return "", nil
}

func (cd *clusterDeploy) getNodeAgentImage(name string) (string, error) {
	uc, err := cd.clusterManager.UserContext(name)
	if err != nil {
		return "", err
	}

	ds, err := uc.Apps.DaemonSets("cattle-system").Get("cattle-node-agent", v1.GetOptions{})
	if err != nil {
		if !apierrors.IsNotFound(err) {
			return "", err
		}
		return "", nil
	}

	for _, c := range ds.Spec.Template.Spec.Containers {
		if c.Name == "agent" {
			return c.Image, nil
		}
	}

	return "", nil
}

func (cd *clusterDeploy) cacheAgentImages(name string) error {
	na, err := cd.getNodeAgentImage(name)
	if err != nil {
		return err
	}

	ca, err := cd.getClusterAgentImage(name)
	if err != nil {
		return err
	}

	agentImagesMutex.Lock()
	defer agentImagesMutex.Unlock()
	agentImages[nodeImage][name] = na
	agentImages[clusterImage][name] = ca
	return nil
}

func agentImagesCached(name string) bool {
	na, ca := getAgentImages(name)
	return na != "" && ca != ""
}

func getAgentImages(name string) (string, string) {
	agentImagesMutex.RLock()
	defer agentImagesMutex.RUnlock()
	return agentImages[nodeImage][name], agentImages[clusterImage][name]
}

func clearAgentImages(name string) {
	agentImagesMutex.Lock()
	defer agentImagesMutex.Unlock()
	delete(agentImages[nodeImage], name)
	delete(agentImages[clusterImage], name)
}
<commit_msg>Enhance log message for kubectl calls<commit_after>package clusterdeploy

import (
	"bytes"
	"context"
	"fmt"
	"reflect"
	"sync"
	"time"

	"github.com/pkg/errors"
	"github.com/rancher/norman/types"
	util "github.com/rancher/rancher/pkg/cluster"
	"github.com/rancher/rancher/pkg/clustermanager"
	"github.com/rancher/rancher/pkg/features"
	"github.com/rancher/rancher/pkg/image"
	"github.com/rancher/rancher/pkg/kubectl"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/rancher/pkg/systemaccount"
	"github.com/rancher/rancher/pkg/systemtemplate"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/rancher/types/user"
	"github.com/sirupsen/logrus"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
)

const (
	AgentForceDeployAnn = "io.cattle.agent.force.deploy"
	nodeImage           = "nodeImage"
	clusterImage        = "clusterImage"
)

var (
	agentImagesMutex sync.RWMutex
	agentImages      = map[string]map[string]string{
		nodeImage:    map[string]string{},
		clusterImage: map[string]string{},
	}
)

func Register(ctx context.Context, management *config.ManagementContext, clusterManager *clustermanager.Manager) {
	c := &clusterDeploy{
		systemAccountManager: systemaccount.NewManager(management),
		userManager:          management.UserManager,
		clusters:             management.Management.Clusters(""),
		nodeLister:           management.Management.Nodes("").Controller().Lister(),
		clusterManager:       clusterManager,
	}

	management.Management.Clusters("").AddHandler(ctx, "cluster-deploy", c.sync)
}

type clusterDeploy struct {
	systemAccountManager *systemaccount.Manager
	userManager          user.Manager
	clusters             v3.ClusterInterface
	clusterManager       *clustermanager.Manager
	nodeLister           v3.NodeLister
}

func (cd *clusterDeploy) sync(key string, cluster *v3.Cluster) (runtime.Object, error) {
	var (
		err, updateErr error
	)

	if cluster == nil || cluster.DeletionTimestamp != nil {
		// remove the system account user created for this cluster
		if err := cd.systemAccountManager.RemoveSystemAccount(key); err != nil {
			return nil, err
		}
		return nil, nil
	}

	original := cluster
	cluster = original.DeepCopy()

	if cluster.Status.Driver == v3.ClusterDriverRKE {
		if cluster.Spec.LocalClusterAuthEndpoint.Enabled {
			cluster.Spec.RancherKubernetesEngineConfig.Authentication.Strategy = "x509|webhook"
		} else {
			cluster.Spec.RancherKubernetesEngineConfig.Authentication.Strategy = "x509"
		}
	}

	err = cd.doSync(cluster)
	if cluster != nil && !reflect.DeepEqual(cluster, original) {
		_, updateErr = cd.clusters.Update(cluster)
	}

	if err != nil {
		return nil, err
	}
	return nil, updateErr
}

func (cd *clusterDeploy) doSync(cluster *v3.Cluster) error {
	if !v3.ClusterConditionProvisioned.IsTrue(cluster) {
		return nil
	}

	nodes, err := cd.nodeLister.List(cluster.Name, labels.Everything())
	if err != nil {
		return err
	}
	if len(nodes) == 0 {
		return nil
	}

	_, err = v3.ClusterConditionSystemAccountCreated.DoUntilTrue(cluster, func() (runtime.Object, error) {
		return cluster, cd.systemAccountManager.CreateSystemAccount(cluster)
	})
	if err != nil {
		return err
	}

	if cluster.Status.AgentImage != "" && !agentImagesCached(cluster.Name) {
		if err := cd.cacheAgentImages(cluster.Name); err != nil {
			return err
		}
	}

	err = cd.deployAgent(cluster)
	if err != nil {
		return err
	}

	return cd.setNetworkPolicyAnn(cluster)
}

// agentFeaturesChanged will treat a missing key as false. This means we only detect changes
// when we set a feature to true so we can't reliably set a feature to false that is enabled by default.
// This behavior makes adding new def false features not cause the agent to redeploy.
func agentFeaturesChanged(desired, actual map[string]bool) bool {
	for k, v := range desired {
		if actual[k] != v {
			return true
		}
	}

	for k, v := range actual {
		if desired[k] != v {
			return true
		}
	}

	return false
}

func redeployAgent(cluster *v3.Cluster, desiredAgent, desiredAuth string, desiredFeatures map[string]bool) bool {
	if !v3.ClusterConditionAgentDeployed.IsTrue(cluster) {
		return true
	}

	forceDeploy := cluster.Annotations[AgentForceDeployAnn] == "true"
	imageChange := cluster.Status.AgentImage != desiredAgent || cluster.Status.AuthImage != desiredAuth
	agentFeaturesChanged := agentFeaturesChanged(desiredFeatures, cluster.Status.AgentFeatures)
	repoChange := false
	if cluster.Spec.RancherKubernetesEngineConfig != nil {
		if cluster.Status.AppliedSpec.RancherKubernetesEngineConfig != nil {
			desiredRepo := util.GetPrivateRepo(cluster)
			var appliedRepo *v3.PrivateRegistry
			if len(cluster.Status.AppliedSpec.RancherKubernetesEngineConfig.PrivateRegistries) > 0 {
				appliedRepo = &cluster.Status.AppliedSpec.RancherKubernetesEngineConfig.PrivateRegistries[0]
			}
			if desiredRepo != nil && appliedRepo != nil && !reflect.DeepEqual(desiredRepo, appliedRepo) {
				repoChange = true
			}
			if (desiredRepo == nil && appliedRepo != nil) || (desiredRepo != nil && appliedRepo == nil) {
				repoChange = true
			}
		}
	}

	if forceDeploy || imageChange || repoChange || agentFeaturesChanged {
		logrus.Infof("Redeploy Rancher Agents is needed for %s: forceDeploy=%v, agent/auth image changed=%v,"+
			" private repo changed=%v, agent features changed=%v", cluster.Name, forceDeploy, imageChange, repoChange,
			agentFeaturesChanged)
		return true
	}

	na, ca := getAgentImages(cluster.Name)
	if cluster.Status.AgentImage != na || cluster.Status.AgentImage != ca {
		// downstream agent does not match, kick a redeploy with settings agent
		logrus.Infof("Redeploy Rancher Agents due to Downstream Agent Image Mismatch for %s: was %s and will be %s",
			cluster.Name, na, image.ResolveWithCluster(settings.AgentImage.Get(), cluster))
		clearAgentImages(cluster.Name)
		return true
	}

	return false
}

func getDesiredImage(cluster *v3.Cluster) string {
	if cluster.Spec.AgentImageOverride != "" {
		return cluster.Spec.AgentImageOverride
	}

	return cluster.Spec.DesiredAgentImage
}

func (cd *clusterDeploy) deployAgent(cluster *v3.Cluster) error {
	desiredAgent := getDesiredImage(cluster)
	if desiredAgent == "" || desiredAgent == "fixed" {
		desiredAgent = image.ResolveWithCluster(settings.AgentImage.Get(), cluster)
	}

	var desiredAuth string
	if cluster.Spec.LocalClusterAuthEndpoint.Enabled {
		desiredAuth = cluster.Spec.DesiredAuthImage
		if desiredAuth == "" || desiredAuth == "fixed" {
			desiredAuth = image.ResolveWithCluster(settings.AuthImage.Get(), cluster)
		}
	}

	desiredFeatures := map[string]bool{
		features.Steve.Name(): features.Steve.Enabled(),
	}

	if !redeployAgent(cluster, desiredAgent, desiredAuth, desiredFeatures) {
		return nil
	}

	kubeConfig, err := cd.getKubeConfig(cluster)
	if err != nil {
		return err
	}

	if _, err = v3.ClusterConditionAgentDeployed.Do(cluster, func() (runtime.Object, error) {
		yaml, err := cd.getYAML(cluster, desiredAgent, desiredAuth, desiredFeatures)
		if err != nil {
			return cluster, err
		}
		var output []byte
		for i := 0; i < 3; i++ {
			// This will fail almost always the first time because when we create the namespace in the file
			// it won't have privileges.  Just stupidly try 3 times
			output, err = kubectl.Apply(yaml, kubeConfig)
			if err == nil {
				break
			}
			time.Sleep(2 * time.Second)
		}
		if err != nil {
			return cluster, errors.WithMessage(types.NewErrors(err, errors.New(string(output))), "kubectl apply failed")
		}
		v3.ClusterConditionAgentDeployed.Message(cluster, string(output))
		if !cluster.Spec.LocalClusterAuthEndpoint.Enabled && cluster.Status.AppliedSpec.LocalClusterAuthEndpoint.Enabled && cluster.Status.AuthImage != "" {
			output, err = kubectl.Delete([]byte(systemtemplate.AuthDaemonSet), kubeConfig)
		}
		if err != nil {
			return cluster, errors.WithMessage(types.NewErrors(err, errors.New(string(output))), "kubectl delete failed")
		}
		v3.ClusterConditionAgentDeployed.Message(cluster, string(output))
		return cluster, nil
	}); err != nil {
		return err
	}

	if err = cd.cacheAgentImages(cluster.Name); err != nil {
		return err
	}

	cluster.Status.AgentImage = desiredAgent
	cluster.Status.AgentFeatures = desiredFeatures
	if cluster.Spec.DesiredAgentImage == "fixed" {
		cluster.Spec.DesiredAgentImage = desiredAgent
	}
	cluster.Status.AuthImage = desiredAuth
	if cluster.Spec.DesiredAuthImage == "fixed" {
		cluster.Spec.DesiredAuthImage = desiredAuth
	}
	if cluster.Annotations[AgentForceDeployAnn] == "true" {
		cluster.Annotations[AgentForceDeployAnn] = "false"
	}

	return nil
}

func (cd *clusterDeploy) setNetworkPolicyAnn(cluster *v3.Cluster) error {
	if cluster.Spec.EnableNetworkPolicy != nil {
		return nil
	}
	// set current state for upgraded canal clusters
	if cluster.Spec.RancherKubernetesEngineConfig != nil &&
		cluster.Spec.RancherKubernetesEngineConfig.Network.Plugin == "canal" {
		enableNetworkPolicy := true
		cluster.Spec.EnableNetworkPolicy = &enableNetworkPolicy
		cluster.Annotations["networking.management.cattle.io/enable-network-policy"] = "true"
	}
	return nil
}

func (cd *clusterDeploy) getKubeConfig(cluster *v3.Cluster) (*clientcmdapi.Config, error) {
	user, err := cd.systemAccountManager.GetSystemUser(cluster.Name)
	if err != nil {
		return nil, err
	}

	token, err := cd.userManager.EnsureToken("agent-"+user.Name, "token for agent deployment", "agent", user.Name)
	if err != nil {
		return nil, err
	}

	return cd.clusterManager.KubeConfig(cluster.Name, token), nil
}

func (cd *clusterDeploy) getYAML(cluster *v3.Cluster, agentImage, authImage string, features map[string]bool) ([]byte, error) {
	logrus.Debug("Desired agent image:", agentImage)
	logrus.Debug("Desired auth image:", authImage)

	token, err := cd.systemAccountManager.GetOrCreateSystemClusterToken(cluster.Name)
	if err != nil {
		return nil, err
	}

	url := settings.ServerURL.Get()
	if url == "" {
		cd.clusters.Controller().EnqueueAfter("", cluster.Name, time.Second)
		return nil, fmt.Errorf("waiting for server-url setting to be set")
	}

	buf := &bytes.Buffer{}
	err = systemtemplate.SystemTemplate(buf, agentImage, authImage, cluster.Name, token, url, cluster.Spec.WindowsPreferedCluster,
		cluster, features)

	return buf.Bytes(), err
}

func (cd *clusterDeploy) getClusterAgentImage(name string) (string, error) {
	uc, err := cd.clusterManager.UserContext(name)
	if err != nil {
		return "", err
	}

	d, err := uc.Apps.Deployments("cattle-system").Get("cattle-cluster-agent", v1.GetOptions{})
	if err != nil {
		if !apierrors.IsNotFound(err) {
			return "", err
		}
		return "", nil
	}

	for _, c := range d.Spec.Template.Spec.Containers {
		if c.Name == "cluster-register" {
			return c.Image, nil
		}
	}

	return "", nil
}

func (cd *clusterDeploy) getNodeAgentImage(name string) (string, error) {
	uc, err := cd.clusterManager.UserContext(name)
	if err != nil {
		return "", err
	}

	ds, err := uc.Apps.DaemonSets("cattle-system").Get("cattle-node-agent", v1.GetOptions{})
	if err != nil {
		if !apierrors.IsNotFound(err) {
			return "", err
		}
		return "", nil
	}

	for _, c := range ds.Spec.Template.Spec.Containers {
		if c.Name == "agent" {
			return c.Image, nil
		}
	}

	return "", nil
}

func (cd *clusterDeploy) cacheAgentImages(name string) error {
	na, err := cd.getNodeAgentImage(name)
	if err != nil {
		return err
	}

	ca, err := cd.getClusterAgentImage(name)
	if err != nil {
		return err
	}

	agentImagesMutex.Lock()
	defer agentImagesMutex.Unlock()
	agentImages[nodeImage][name] = na
	agentImages[clusterImage][name] = ca
	return nil
}

func agentImagesCached(name string) bool {
	na, ca := getAgentImages(name)
	return na != "" && ca != ""
}

func getAgentImages(name string) (string, string) {
	agentImagesMutex.RLock()
	defer agentImagesMutex.RUnlock()
	return agentImages[nodeImage][name], agentImages[clusterImage][name]
}

func clearAgentImages(name string) {
	agentImagesMutex.Lock()
	defer agentImagesMutex.Unlock()
	delete(agentImages[nodeImage], name)
	delete(agentImages[clusterImage], name)
}
<|endoftext|>


<commit_before>package packagemanifest

import (
	"context"

	k8serrors "k8s.io/apimachinery/pkg/api/errors"
	metainternalversion "k8s.io/apimachinery/pkg/apis/meta/internalversion"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	genericapirequest "k8s.io/apiserver/pkg/endpoints/request"
	"k8s.io/apiserver/pkg/registry/rest"

	"github.com/operator-framework/operator-lifecycle-manager/pkg/package-server/apis/packagemanifest/v1alpha1"
	"github.com/operator-framework/operator-lifecycle-manager/pkg/package-server/provider"
)

type PackageManifestStorage struct {
	groupResource schema.GroupResource
	prov          provider.PackageManifestProvider
}

var _ rest.KindProvider = &PackageManifestStorage{}
var _ rest.Storage = &PackageManifestStorage{}
var _ rest.Getter = &PackageManifestStorage{}
var _ rest.Lister = &PackageManifestStorage{}
var _ rest.Scoper = &PackageManifestStorage{}

// NewStorage returns an in-memory implementation of storage.Interface.
func NewStorage(groupResource schema.GroupResource, prov provider.PackageManifestProvider) *PackageManifestStorage {
	return &PackageManifestStorage{
		groupResource: groupResource,
		prov:          prov,
	}
}

// Storage interface
func (m *PackageManifestStorage) New() runtime.Object {
	return &v1alpha1.PackageManifest{}
}

// KindProvider interface
func (m *PackageManifestStorage) Kind() string {
	return "PackageManifest"
}

// Lister interface
func (m *PackageManifestStorage) NewList() runtime.Object {
	return &v1alpha1.PackageManifestList{}
}

// Lister interface
func (m *PackageManifestStorage) List(ctx context.Context, options *metainternalversion.ListOptions) (runtime.Object, error) {
	// get namespace
	namespace := genericapirequest.NamespaceValue(ctx)

	// get selectors
	labelSelector := labels.Everything()
	if options != nil && options.LabelSelector != nil {
		labelSelector = options.LabelSelector
	}

	res, err := m.prov.ListPackageManifests(namespace)
	if err != nil {
		return &v1alpha1.PackageManifestList{}, err
	}

	// filter results by label
	filtered := make([]v1alpha1.PackageManifest, len(res.Items))
	i := 0
	for _, manifest := range res.Items {
		if labelSelector.Matches(labels.Set(manifest.GetLabels())) {
			filtered[i] = manifest
			i++
		}
	}

	res.Items = filtered
	return res, nil
}

// Getter interface
func (m *PackageManifestStorage) Get(ctx context.Context, name string, opts *metav1.GetOptions) (runtime.Object, error) {
	namespace := genericapirequest.NamespaceValue(ctx)
	manifest := v1alpha1.PackageManifest{}

	pm, err := m.prov.GetPackageManifest(namespace, name)
	if err != nil {
		return nil, err
	}
	if pm != nil {
		manifest = *pm
	} else {
		return nil, k8serrors.NewNotFound(m.groupResource, name)
	}

	return &manifest, nil
}

// Scoper interface
func (m *PackageManifestStorage) NamespaceScoped() bool {
	return true
}
<commit_msg>fix(reststorage): start with empty slice when filtering<commit_after>package packagemanifest

import (
	"context"

	k8serrors "k8s.io/apimachinery/pkg/api/errors"
	metainternalversion "k8s.io/apimachinery/pkg/apis/meta/internalversion"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	genericapirequest "k8s.io/apiserver/pkg/endpoints/request"
	"k8s.io/apiserver/pkg/registry/rest"

	"github.com/operator-framework/operator-lifecycle-manager/pkg/package-server/apis/packagemanifest/v1alpha1"
	"github.com/operator-framework/operator-lifecycle-manager/pkg/package-server/provider"
)

type PackageManifestStorage struct {
	groupResource schema.GroupResource
	prov          provider.PackageManifestProvider
}

var _ rest.KindProvider = &PackageManifestStorage{}
var _ rest.Storage = &PackageManifestStorage{}
var _ rest.Getter = &PackageManifestStorage{}
var _ rest.Lister = &PackageManifestStorage{}
var _ rest.Scoper = &PackageManifestStorage{}

// NewStorage returns an in-memory implementation of storage.Interface.
func NewStorage(groupResource schema.GroupResource, prov provider.PackageManifestProvider) *PackageManifestStorage {
	return &PackageManifestStorage{
		groupResource: groupResource,
		prov:          prov,
	}
}

// Storage interface
func (m *PackageManifestStorage) New() runtime.Object {
	return &v1alpha1.PackageManifest{}
}

// KindProvider interface
func (m *PackageManifestStorage) Kind() string {
	return "PackageManifest"
}

// Lister interface
func (m *PackageManifestStorage) NewList() runtime.Object {
	return &v1alpha1.PackageManifestList{}
}

// Lister interface
func (m *PackageManifestStorage) List(ctx context.Context, options *metainternalversion.ListOptions) (runtime.Object, error) {
	// get namespace
	namespace := genericapirequest.NamespaceValue(ctx)

	// get selectors
	labelSelector := labels.Everything()
	if options != nil && options.LabelSelector != nil {
		labelSelector = options.LabelSelector
	}

	res, err := m.prov.ListPackageManifests(namespace)
	if err != nil {
		return &v1alpha1.PackageManifestList{}, err
	}

	// filter results by label
	filtered := []v1alpha1.PackageManifest{}
	for _, manifest := range res.Items {
		if labelSelector.Matches(labels.Set(manifest.GetLabels())) {
			filtered = append(filtered, manifest)
		}
	}

	res.Items = filtered
	return res, nil
}

// Getter interface
func (m *PackageManifestStorage) Get(ctx context.Context, name string, opts *metav1.GetOptions) (runtime.Object, error) {
	namespace := genericapirequest.NamespaceValue(ctx)
	manifest := v1alpha1.PackageManifest{}

	pm, err := m.prov.GetPackageManifest(namespace, name)
	if err != nil {
		return nil, err
	}
	if pm != nil {
		manifest = *pm
	} else {
		return nil, k8serrors.NewNotFound(m.groupResource, name)
	}

	return &manifest, nil
}

// Scoper interface
func (m *PackageManifestStorage) NamespaceScoped() bool {
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"io"
	"net/http"

	"github.com/NAExpire/API/src/util"
	"github.com/gorilla/mux"
)

type GetRestaurantHandler struct {
	DB *sql.DB
}

type UpdateRestaurantHandler struct {
	DB *sql.DB
}

type restaurantSchema struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	Address     string `json:"address"`
	City        string `json:"city"`
	State       string `json:"state"`
}

func (handler GetRestaurantHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	vars := mux.Vars(request)
	x := &restaurantSchema{}

	rows, err := handler.DB.Query("SELECT `name`, `description`, `address`, `city`, `state` FROM restaurants WHERE id=?", vars["id"])

	defer rows.Close()

	if err != nil {
		writer.WriteHeader(http.StatusInternalServerError)
		util.WriteErrorJSON(writer, err.Error())
		return
	}

	if !rows.Next() {
		writer.WriteHeader(http.StatusNotFound)
		util.WriteErrorJSON(writer, "Restaurant with ID "+vars["id"]+" could not be found")
		return
	}

	err = rows.Scan(&x.Name, &x.Description, &x.Address, &x.City, &x.State)
	if err != nil {
		writer.WriteHeader(http.StatusInternalServerError)
		util.WriteErrorJSON(writer, err.Error())
		return
	}

	util.EncodeJSON(writer, x)
}

func (handler UpdateRestaurantHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	vars := mux.Vars(request)
	x := &restaurantSchema{}
	err := util.DecodeJSON(request.Body, x)

	if err != nil {
		io.WriteString(writer, err.Error()+"\n")
		return
	}

	_, err = handler.DB.Exec("UPDATE restaurants SET name = ? , description = ? , address = ? , city = ? , state = ?  WHERE id = ?", x.Name, x.Description, x.Address, x.City, x.State, vars["id"])

	if err != nil {
		writer.WriteHeader(http.StatusInternalServerError)
		io.WriteString(writer, err.Error()+"\n")
		return
	}

	io.WriteString(writer, "{\"ok\": true}")
}
<commit_msg>added Business # and pickup time for restaurants<commit_after>package main

import (
	"database/sql"
	"io"
	"net/http"

	"github.com/NAExpire/API/src/util"
	"github.com/gorilla/mux"
)

type GetRestaurantHandler struct {
	DB *sql.DB
}

type UpdateRestaurantHandler struct {
	DB *sql.DB
}

type restaurantSchema struct {
	Name          string `json:"name"`
	BusinessPhone string `json:"phone-number"`
	PickupTime    string `json:"pickup-time"`
	Description   string `json:"description"`
	Address       string `json:"address"`
	City          string `json:"city"`
	State         string `json:"state"`
}

func (handler GetRestaurantHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	vars := mux.Vars(request)
	x := &restaurantSchema{}

	rows, err := handler.DB.Query("SELECT `name`, `phone-number`, `pickup-time`, `description`, `address`, `city`, `state` FROM restaurants WHERE id=?", vars["id"])

	defer rows.Close()

	if err != nil {
		writer.WriteHeader(http.StatusInternalServerError)
		util.WriteErrorJSON(writer, err.Error())
		return
	}

	if !rows.Next() {
		writer.WriteHeader(http.StatusNotFound)
		util.WriteErrorJSON(writer, "Restaurant with ID "+vars["id"]+" could not be found")
		return
	}

	err = rows.Scan(&x.Name, &x.Description, &x.Address, &x.City, &x.State)
	if err != nil {
		writer.WriteHeader(http.StatusInternalServerError)
		util.WriteErrorJSON(writer, err.Error())
		return
	}

	util.EncodeJSON(writer, x)
}

func (handler UpdateRestaurantHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	vars := mux.Vars(request)
	x := &restaurantSchema{}
	err := util.DecodeJSON(request.Body, x)

	if err != nil {
		io.WriteString(writer, err.Error()+"\n")
		return
	}

	_, err = handler.DB.Exec("UPDATE restaurants SET name = ? , phone-number = ? , pickup-time = ? , description = ? , address = ? , city = ? , state = ?  WHERE id = ?", x.Name, x.BusinessPhone, x.PickupTime, x.Description, x.Address, x.City, x.State, vars["id"])

	if err != nil {
		writer.WriteHeader(http.StatusInternalServerError)
		io.WriteString(writer, err.Error()+"\n")
		return
	}

	io.WriteString(writer, "{\"ok\": true}")
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"log"
	"net"

	"telnet"
)

const (
	cmdSE   = 240
	cmdSB   = 250
	cmdWill = 251
	cmdWont = 252
	cmdDo   = 253
	cmdDont = 254
	cmdIAC  = 255
)

const (
	optEcho           = 1
	optSupressGoAhead = 3
	optLinemode       = 34
)

const (
	IAC_NONE = iota
	IAC_CMD  = iota
	IAC_OPT  = iota
	IAC_SUB  = iota
)

const (
	QUIT = iota
	MOTD = iota
	USER = iota
	PASS = iota
	EXEC = iota
	ENAB = iota
	CONF = iota
)

type TelnetClient struct {
	//rd      *bufio.Reader
	conn       net.Conn
	wr         *bufio.Writer
	userOut    chan string // outputLoop: read from userOut and write into wr
	quit       chan int
	echo       chan bool
	status     int
	serverEcho bool
}

type Command struct {
	client *TelnetClient
	line   string
}

var cmdInput = make(chan Command)

func charReadLoop(conn net.Conn, read chan<- byte) {
	input := make([]byte, 10) // last input

	for {
		rd, err := conn.Read(input)
		if err != nil {
			log.Printf("charReadLoop: net.Read: %v", err)
			break
		}
		curr := input[:rd]
		log.Printf("charReadLoop: read len=%d [%s]", rd, curr)
		for _, b := range curr {
			read <- b
		}
	}

	log.Printf("charReadLoop: exiting")

	close(read)
}

func reader(conn net.Conn) <-chan byte {
	read := make(chan byte)
	go charReadLoop(conn, read)
	return read
}

func inputLoop(client *TelnetClient) {
	//loop:
	//	- read from rd and feed into cli interpreter
	//	- watch idle timeout
	//	- watch quitInput channel

	iac := IAC_NONE
	buf := [30]byte{} // underlying buffer
	//line := buf[:0]   // position at underlying buffer
	size := 0 // position at underlying buffer

	read := reader(client.conn)

LOOP:
	for {
		select {
		case client.serverEcho = <-client.echo:
			// do nothing
		case b, ok := <-read:
			if !ok {
				log.Printf("inputLoop: closed channel")
				break LOOP
			}

			switch iac {
			case IAC_NONE:
				switch b {
				case cmdIAC:
					// hit IAC mark?
					log.Printf("inputLoop: telnet IAC begin")
					iac = IAC_CMD
					continue
				case '\r':
					// discard
				case '\n':
					//cmdLine := string(line) // string is safe for sharing (immutable)
					cmdLine := string(buf[:size]) // string is safe for sharing (immutable)
					log.Printf("inputLoop: cmdLine len=%d [%s]", len(cmdLine), cmdLine)
					cmdInput <- Command{client, cmdLine}
					//line = buf[:0] // reset reading buffer position
					size = 0 // reset reading buffer position

					// echo newline back to client
					if client.serverEcho {
						client.userOut <- "\r\n"
					}

				default:
					// push non-commands bytes into line buffer
					//line = append(buf[:len(line)], b)
					buf[size] = b
					size++

					// echo char back to client
					if client.serverEcho {
						client.userOut <- string(b)
					}
				}

			case IAC_CMD:

				switch b {
				case cmdSB:
					log.Printf("inputLoop: telnet SUB begin")
					iac = IAC_SUB
				case cmdWill, cmdWont, cmdDo, cmdDont:
					log.Printf("inputLoop: telnet OPT begin")
					iac = IAC_OPT
				default:
					log.Printf("inputLoop: telnet IAC end")
					iac = IAC_NONE
				}

			case IAC_OPT:

				log.Printf("inputLoop: telnet OPT end")
				log.Printf("inputLoop: telnet IAC end")
				iac = IAC_NONE

			case IAC_SUB:

				if b == cmdSE {
					log.Printf("inputLoop: telnet SUB end")
					log.Printf("inputLoop: telnet IAC end")
					iac = IAC_NONE
				}

			default:
				log.Panicf("inputLoop: unexpected state iac=%d", iac)
			}

		}

		log.Printf("inputLoop: buf len=%d [%s]", size, buf[:size])
	}

	log.Printf("inputLoop: exiting")
}

func outputLoop(client *TelnetClient) {
	//loop:
	//	- read from userOut channel and write into wr
	//	- watch quitOutput channel

LOOP:
	for {
		select {
		case msg := <-client.userOut:
			if n, err := client.wr.WriteString(msg); err != nil {
				log.Printf("outputLoop: written=%d from=%d: %v", n, len(msg), err)
			}
			if err := client.wr.Flush(); err != nil {
				log.Printf("outputLoop: flush: %v", err)
			}
		case _, ok := <-client.quit:
			if !ok {
				break LOOP
			}
		}
	}

	log.Printf("outputLoop: exiting")
}

func charMode(conn net.Conn) {
	cmd := []byte{cmdIAC, cmdWill, optEcho, cmdIAC, cmdWill, optSupressGoAhead, cmdIAC, cmdDont, optLinemode}

	wr, err := conn.Write(cmd)

	log.Printf("charMode: len=%d err=%v", wr, err)
}

func handleTelnet(conn net.Conn) {
	defer conn.Close()

	log.Printf("new telnet connection from: %s", conn.RemoteAddr())

	//rd, wr := bufio.NewReader(conn), bufio.NewWriter(conn)

	charMode(conn)

	client := TelnetClient{conn, bufio.NewWriter(conn), make(chan string), make(chan int), make(chan bool), MOTD, true}

	defer close(client.userOut)

	cmdInput <- Command{&client, ""} // mock user input

	go inputLoop(&client)

	outputLoop(&client)
}

func listenTelnet(addr string) {
	telnetServer := telnet.Server{Addr: addr, Handler: handleTelnet}

	log.Printf("serving telnet on TCP %s", addr)

	if err := telnetServer.ListenAndServe(); err != nil {
		log.Fatalf("telnet server on address %s: error: %s", addr, err)
	}
}
<commit_msg>Prevent line input buffer overflow.<commit_after>package main

import (
	"bufio"
	"fmt"
	"log"
	"net"

	"telnet"
)

const (
	cmdSE   = 240
	cmdSB   = 250
	cmdWill = 251
	cmdWont = 252
	cmdDo   = 253
	cmdDont = 254
	cmdIAC  = 255
)

const (
	optEcho           = 1
	optSupressGoAhead = 3
	optLinemode       = 34
)

const (
	IAC_NONE = iota
	IAC_CMD  = iota
	IAC_OPT  = iota
	IAC_SUB  = iota
)

const (
	QUIT = iota
	MOTD = iota
	USER = iota
	PASS = iota
	EXEC = iota
	ENAB = iota
	CONF = iota
)

type TelnetClient struct {
	//rd      *bufio.Reader
	conn       net.Conn
	wr         *bufio.Writer
	userOut    chan string // outputLoop: read from userOut and write into wr
	quit       chan int
	echo       chan bool
	status     int
	serverEcho bool
}

type Command struct {
	client *TelnetClient
	line   string
}

var cmdInput = make(chan Command)

func charReadLoop(conn net.Conn, read chan<- byte) {
	input := make([]byte, 10) // last input

	for {
		rd, err := conn.Read(input)
		if err != nil {
			log.Printf("charReadLoop: net.Read: %v", err)
			break
		}
		curr := input[:rd]
		log.Printf("charReadLoop: read len=%d [%s]", rd, curr)
		for _, b := range curr {
			read <- b
		}
	}

	log.Printf("charReadLoop: exiting")

	close(read)
}

func reader(conn net.Conn) <-chan byte {
	read := make(chan byte)
	go charReadLoop(conn, read)
	return read
}

func inputLoop(client *TelnetClient) {
	//loop:
	//	- read from rd and feed into cli interpreter
	//	- watch idle timeout
	//	- watch quitInput channel

	iac := IAC_NONE
	buf := [30]byte{} // underlying buffer
	//line := buf[:0]   // position at underlying buffer
	size := 0 // position at underlying buffer

	read := reader(client.conn)

LOOP:
	for {
		select {
		case client.serverEcho = <-client.echo:
			// do nothing
		case b, ok := <-read:
			if !ok {
				log.Printf("inputLoop: closed channel")
				break LOOP
			}

			switch iac {
			case IAC_NONE:
				switch b {
				case cmdIAC:
					// hit IAC mark?
					log.Printf("inputLoop: telnet IAC begin")
					iac = IAC_CMD
					continue
				case '\r':
					// discard
				case '\n':
					//cmdLine := string(line) // string is safe for sharing (immutable)
					cmdLine := string(buf[:size]) // string is safe for sharing (immutable)
					log.Printf("inputLoop: cmdLine len=%d [%s]", len(cmdLine), cmdLine)
					cmdInput <- Command{client, cmdLine}
					//line = buf[:0] // reset reading buffer position
					size = 0 // reset reading buffer position

					// echo newline back to client
					if client.serverEcho {
						client.userOut <- "\r\n"
					}

				default:
					// push non-commands bytes into line buffer

					if size >= len(buf) {
						client.userOut <- fmt.Sprintf("\r\nline buffer overflow: size=%d max=%d\r\n", size, len(buf))
						client.userOut <- string(buf[:size]) // redisplay command to user
						continue LOOP
					}

					//line = append(buf[:len(line)], b)
					buf[size] = b
					size++

					// echo char back to client
					if client.serverEcho {
						client.userOut <- string(b)
					}
				}

			case IAC_CMD:

				switch b {
				case cmdSB:
					log.Printf("inputLoop: telnet SUB begin")
					iac = IAC_SUB
				case cmdWill, cmdWont, cmdDo, cmdDont:
					log.Printf("inputLoop: telnet OPT begin")
					iac = IAC_OPT
				default:
					log.Printf("inputLoop: telnet IAC end")
					iac = IAC_NONE
				}

			case IAC_OPT:

				log.Printf("inputLoop: telnet OPT end")
				log.Printf("inputLoop: telnet IAC end")
				iac = IAC_NONE

			case IAC_SUB:

				if b == cmdSE {
					log.Printf("inputLoop: telnet SUB end")
					log.Printf("inputLoop: telnet IAC end")
					iac = IAC_NONE
				}

			default:
				log.Panicf("inputLoop: unexpected state iac=%d", iac)
			}

		}

		log.Printf("inputLoop: buf len=%d [%s]", size, buf[:size])
	}

	log.Printf("inputLoop: exiting")
}

func outputLoop(client *TelnetClient) {
	//loop:
	//	- read from userOut channel and write into wr
	//	- watch quitOutput channel

LOOP:
	for {
		select {
		case msg := <-client.userOut:
			if n, err := client.wr.WriteString(msg); err != nil {
				log.Printf("outputLoop: written=%d from=%d: %v", n, len(msg), err)
			}
			if err := client.wr.Flush(); err != nil {
				log.Printf("outputLoop: flush: %v", err)
			}
		case _, ok := <-client.quit:
			if !ok {
				break LOOP
			}
		}
	}

	log.Printf("outputLoop: exiting")
}

func charMode(conn net.Conn) {
	cmd := []byte{cmdIAC, cmdWill, optEcho, cmdIAC, cmdWill, optSupressGoAhead, cmdIAC, cmdDont, optLinemode}

	wr, err := conn.Write(cmd)

	log.Printf("charMode: len=%d err=%v", wr, err)
}

func handleTelnet(conn net.Conn) {
	defer conn.Close()

	log.Printf("new telnet connection from: %s", conn.RemoteAddr())

	//rd, wr := bufio.NewReader(conn), bufio.NewWriter(conn)

	charMode(conn)

	client := TelnetClient{conn, bufio.NewWriter(conn), make(chan string), make(chan int), make(chan bool), MOTD, true}

	defer close(client.userOut)

	cmdInput <- Command{&client, ""} // mock user input

	go inputLoop(&client)

	outputLoop(&client)
}

func listenTelnet(addr string) {
	telnetServer := telnet.Server{Addr: addr, Handler: handleTelnet}

	log.Printf("serving telnet on TCP %s", addr)

	if err := telnetServer.ListenAndServe(); err != nil {
		log.Fatalf("telnet server on address %s: error: %s", addr, err)
	}
}
<|endoftext|>