<commit_before>package main

import (
	"bytes"
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/datawire/teleproxy/pkg/supervisor"
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh/terminal"
	"gopkg.in/natefinch/lumberjack.v2"
)

// DaemonFormatter formats log messages for the Playpen Daemon
type DaemonFormatter struct{}

// Format implement logrus.Formatter
func (f *DaemonFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	var b *bytes.Buffer
	if entry.Buffer != nil {
		b = entry.Buffer
	} else {
		b = &bytes.Buffer{}
	}

	fmt.Fprintf(b, "%s %s", entry.Time.Format("2006/01/02 15:04:05"), entry.Message)

	if len(entry.Data) > 0 {
		keys := make([]string, 0, len(entry.Data))
		for k := range entry.Data {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		for _, k := range keys {
			v := entry.Data[k]
			fmt.Fprintf(b, " %s=%+v", k, v)
		}
	}
	b.WriteByte('\n')
	return b.Bytes(), nil
}

// SetUpLogging sets up standard Playpen Daemon logging
func SetUpLogging() supervisor.Logger {
	logger := logrus.StandardLogger()
	logger.Formatter = new(DaemonFormatter)
	if !terminal.IsTerminal(int(os.Stdout.Fd())) {
		logger.SetOutput(&lumberjack.Logger{
			Filename:   logfile,
			MaxSize:    10,   // megabytes
			MaxBackups: 3,    // in the same directory
			MaxAge:     60,   // days
			LocalTime:  true, // rotated logfiles use local time names
		})
	}
	return logger
}

func doWordWrap(text string, prefix string, lineWidth int) []string {
	words := strings.Fields(strings.TrimSpace(text))
	if len(words) == 0 {
		return []string{""}
	}
	lines := make([]string, 0)
	wrapped := prefix + words[0]
	for _, word := range words[1:] {
		if len(word)+1 > lineWidth-len(wrapped) {
			lines = append(lines, wrapped)
			wrapped = prefix + word
		} else {
			wrapped += " " + word
		}
	}
	if len(wrapped) > 0 {
		lines = append(lines, wrapped)
	}
	return lines
}

var terminalWidth = 0 // Set on first use

// WordWrap returns a slice of strings with the original content wrapped at the
// terminal width or at 80 characters if no terminal is present.
func WordWrap(text string) []string {
	if terminalWidth <= 0 {
		terminalWidth = 80
		fd := int(os.Stdout.Fd())
		if terminal.IsTerminal(fd) {
			w, _, err := terminal.GetSize(fd)
			if err == nil {
				terminalWidth = w
			}
		}
	}
	return doWordWrap(text, "", terminalWidth)
}

// WordWrapString returns a string with the original content wrapped at the
// terminal width or at 80 characters if no terminal is present.
func WordWrapString(text string) string {
	return strings.Join(WordWrap(text), "\n")
}
<commit_msg>Emit compact timestamps when logging to the console<commit_after>package main

import (
	"bytes"
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/datawire/teleproxy/pkg/supervisor"
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh/terminal"
	"gopkg.in/natefinch/lumberjack.v2"
)

// DaemonFormatter formats log messages for the Playpen Daemon
type DaemonFormatter struct {
	TimestampFormat string
}

// Format implement logrus.Formatter
func (f *DaemonFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	var b *bytes.Buffer
	if entry.Buffer != nil {
		b = entry.Buffer
	} else {
		b = &bytes.Buffer{}
	}

	fmt.Fprintf(b, "%s %s", entry.Time.Format(f.TimestampFormat), entry.Message)

	if len(entry.Data) > 0 {
		keys := make([]string, 0, len(entry.Data))
		for k := range entry.Data {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		for _, k := range keys {
			v := entry.Data[k]
			fmt.Fprintf(b, " %s=%+v", k, v)
		}
	}
	b.WriteByte('\n')
	return b.Bytes(), nil
}

// SetUpLogging sets up standard Playpen Daemon logging
func SetUpLogging() supervisor.Logger {
	loggingToTerminal := terminal.IsTerminal(int(os.Stdout.Fd()))
	logger := logrus.StandardLogger()
	formatter := new(DaemonFormatter)
	logger.Formatter = formatter
	if loggingToTerminal {
		formatter.TimestampFormat = "15:04:05"
	} else {
		formatter.TimestampFormat = "2006/01/02 15:04:05"
		logger.SetOutput(&lumberjack.Logger{
			Filename:   logfile,
			MaxSize:    10,   // megabytes
			MaxBackups: 3,    // in the same directory
			MaxAge:     60,   // days
			LocalTime:  true, // rotated logfiles use local time names
		})
	}
	return logger
}

func doWordWrap(text string, prefix string, lineWidth int) []string {
	words := strings.Fields(strings.TrimSpace(text))
	if len(words) == 0 {
		return []string{""}
	}
	lines := make([]string, 0)
	wrapped := prefix + words[0]
	for _, word := range words[1:] {
		if len(word)+1 > lineWidth-len(wrapped) {
			lines = append(lines, wrapped)
			wrapped = prefix + word
		} else {
			wrapped += " " + word
		}
	}
	if len(wrapped) > 0 {
		lines = append(lines, wrapped)
	}
	return lines
}

var terminalWidth = 0 // Set on first use

// WordWrap returns a slice of strings with the original content wrapped at the
// terminal width or at 80 characters if no terminal is present.
func WordWrap(text string) []string {
	if terminalWidth <= 0 {
		terminalWidth = 80
		fd := int(os.Stdout.Fd())
		if terminal.IsTerminal(fd) {
			w, _, err := terminal.GetSize(fd)
			if err == nil {
				terminalWidth = w
			}
		}
	}
	return doWordWrap(text, "", terminalWidth)
}

// WordWrapString returns a string with the original content wrapped at the
// terminal width or at 80 characters if no terminal is present.
func WordWrapString(text string) string {
	return strings.Join(WordWrap(text), "\n")
}
<|endoftext|>


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

package platform

import (
	"crypto/rand"
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/coreos/mantle/auth"

	"github.com/coreos/mantle/Godeps/_workspace/src/golang.org/x/crypto/ssh"
	"github.com/coreos/mantle/Godeps/_workspace/src/google.golang.org/api/compute/v1"
)

type GCEOptions struct {
	Image       string
	Project     string
	Zone        string
	MachineType string
	DiskType    string
	BaseName    string
	Network     string
	ServiceAuth bool
}

type gceCluster struct {
	*baseCluster
	conf *GCEOptions
	api  *compute.Service
}

type gceMachine struct {
	gc    *gceCluster
	name  string
	intIP string
	extIP string
}

func NewGCECluster(conf GCEOptions) (Cluster, error) {
	var client *http.Client
	var err error
	if conf.ServiceAuth {
		client = auth.GoogleServiceClient()
	} else {
		client, err = auth.GoogleClient()
	}
	if err != nil {
		return nil, err
	}

	api, err := compute.New(client)
	if err != nil {
		return nil, err
	}

	bc, err := newBaseCluster()
	if err != nil {
		return nil, err
	}

	gc := &gceCluster{
		baseCluster: bc,
		api:         api,
		conf:        &conf,
	}

	return gc, nil
}

func (gc *gceCluster) Destroy() error {
	for _, gm := range gc.Machines() {
		gm.Destroy()
	}
	gc.agent.Close()
	return nil
}

// Calling in parallel is ok
func (gc *gceCluster) NewMachine(userdata string) (Machine, error) {
	conf, err := NewConf(userdata)
	if err != nil {
		return nil, err
	}

	keys, err := gc.agent.List()
	if err != nil {
		return nil, err
	}

	conf.CopyKeys(keys)

	// Create gce VM and wait for creation to succeed.
	gm, err := GCECreateVM(gc.api, gc.conf, conf.String())
	if err != nil {
		return nil, err
	}
	gm.gc = gc

	if err := commonMachineChecks(gm); err != nil {
		gm.Destroy()
		return nil, err
	}

	gc.addMach(gm)

	return Machine(gm), nil
}

func (gm *gceMachine) ID() string {
	return gm.name
}

func (gm *gceMachine) IP() string {
	return gm.extIP
}

func (gm *gceMachine) PrivateIP() string {
	return gm.intIP
}

func (gm *gceMachine) SSHClient() (*ssh.Client, error) {
	return gm.gc.SSHClient(gm.IP())
}

func (gm *gceMachine) SSH(cmd string) ([]byte, error) {
	return gm.gc.SSH(gm, cmd)
}

func (gm *gceMachine) Destroy() error {
	_, err := gm.gc.api.Instances.Delete(gm.gc.conf.Project, gm.gc.conf.Zone, gm.name).Do()
	if err != nil {
		return err
	}

	gm.gc.delMach(gm)

	return nil
}

func GCECreateVM(api *compute.Service, opts *GCEOptions, userdata string) (*gceMachine, error) {
	// generate name
	name, err := newName(opts)
	if err != nil {
		return nil, fmt.Errorf("Failed allocating unique name for vm: %v\n", err)
	}

	instance, err := gceMakeInstance(opts, userdata, name)
	if err != nil {
		return nil, err
	}

	// request instance
	op, err := api.Instances.Insert(opts.Project, opts.Zone, instance).Do()
	if err != nil {
		return nil, fmt.Errorf("Failed to create new VM: %v\n", err)
	}

	fmt.Fprintf(os.Stderr, "Instance %v requested\n", name)
	fmt.Fprintf(os.Stderr, "Waiting for creation to finish...\n")

	// wait for creation to finish
	err = gceWaitVM(api, opts.Project, opts.Zone, op.Name)
	if err != nil {
		return nil, err
	}

	inst, err := api.Instances.Get(opts.Project, opts.Zone, name).Do()
	if err != nil {
		return nil, fmt.Errorf("Error getting instance %s details after creation: %v", name, err)
	}
	intIP, extIP := instanceIPs(inst)

	gm := &gceMachine{
		name:  name,
		extIP: extIP,
		intIP: intIP,
	}

	return gm, nil
}

func GCEDestroyVM(api *compute.Service, proj, zone, name string) error {
	_, err := api.Instances.Delete(proj, zone, name).Do()
	if err != nil {
		return err
	}
	return nil
}

// Create image on GCE and and wait for completion. Will not overwrite
// existing image.
func GCECreateImage(api *compute.Service, proj, name, source string) error {
	image := &compute.Image{
		Name: name,
		RawDisk: &compute.ImageRawDisk{
			Source: source,
		},
	}

	fmt.Fprintf(os.Stderr, "Image %v requested\n", name)
	fmt.Fprintf(os.Stderr, "Waiting for image creation to finish...\n")

	op, err := api.Images.Insert(proj, image).Do()
	if err != nil {
		return err
	}

	err = gceWaitOp(api, proj, op.Name)
	if err != nil {
		return err
	}
	return nil
}

// Delete image on GCE and then recreate it.
func GCEForceCreateImage(api *compute.Service, proj, name, source string) error {
	// op xor err = nil
	op, err := api.Images.Delete(proj, name).Do()

	if op != nil {
		err = gceWaitOp(api, proj, op.Name)
		if err != nil {
			return err
		}
	}

	// don't return error when delete fails because image doesn't exist
	if err != nil && !strings.HasSuffix(err.Error(), "notFound") {
		return fmt.Errorf("deleting image: %v", err)
	}

	// create
	return GCECreateImage(api, proj, name, source)
}

func GCEListVMs(api *compute.Service, opts *GCEOptions, prefix string) ([]Machine, error) {
	var vms []Machine

	list, err := api.Instances.List(opts.Project, opts.Zone).Do()
	if err != nil {
		return nil, err
	}

	for _, inst := range list.Items {
		if !strings.HasPrefix(inst.Name, prefix) {
			continue
		}
		intIP, extIP := instanceIPs(inst)
		gm := &gceMachine{
			name:  inst.Name,
			extIP: extIP,
			intIP: intIP,
		}

		vms = append(vms, gm)
	}
	return vms, nil
}

func GCEListImages(client *http.Client, proj, prefix string) ([]string, error) {
	var images []string
	computeService, err := compute.New(client)
	list, err := computeService.Images.List(proj).Do()
	if err != nil {
		return nil, err
	}

	for _, image := range list.Items {
		if !strings.HasPrefix(image.Name, prefix) {
			continue
		}
		images = append(images, image.Name)
	}
	return images, nil
}

//Some code taken from: https://github.com/golang/build/blob/master/buildlet/gce.go
func gceMakeInstance(opts *GCEOptions, userdata string, name string) (*compute.Instance, error) {
	prefix := "https://www.googleapis.com/compute/v1/projects/" + opts.Project
	instance := &compute.Instance{
		Name:        name,
		MachineType: prefix + "/zones/" + opts.Zone + "/machineTypes/" + opts.MachineType,
		Metadata:    &compute.Metadata{},
		Disks: []*compute.AttachedDisk{
			{
				AutoDelete: true,
				Boot:       true,
				Type:       "PERSISTENT",
				InitializeParams: &compute.AttachedDiskInitializeParams{
					DiskName:    name,
					SourceImage: prefix + "/global/images/" + opts.Image,
					DiskType:    "/zones/" + opts.Zone + "/diskTypes/" + opts.DiskType,
				},
			},
		},
		NetworkInterfaces: []*compute.NetworkInterface{
			&compute.NetworkInterface{
				AccessConfigs: []*compute.AccessConfig{
					&compute.AccessConfig{
						Type: "ONE_TO_ONE_NAT",
						Name: "External NAT",
					},
				},
				Network: prefix + "/global/networks/" + opts.Network,
			},
		},
	}
	// add cloud config
	if userdata != "" {
		instance.Metadata.Items = append(instance.Metadata.Items, &compute.MetadataItems{
			Key:   "user-data",
			Value: userdata,
		})
	}

	return instance, nil
}

//Some code taken from: https://github.com/golang/build/blob/master/buildlet/gce.go
func gceWaitVM(api *compute.Service, proj, zone, opname string) error {
OpLoop:
	for {
		time.Sleep(2 * time.Second)

		op, err := api.ZoneOperations.Get(proj, zone, opname).Do()
		if err != nil {
			return fmt.Errorf("Failed to get op %s: %v", opname, err)
		}
		switch op.Status {
		case "PENDING", "RUNNING":
			continue
		case "DONE":
			if op.Error != nil {
				for _, operr := range op.Error.Errors {
					return fmt.Errorf("Error creating instance: %+v", operr)
				}
				return fmt.Errorf("Failed to start.")
			}
			break OpLoop
		default:
			return fmt.Errorf("Unknown create status %q: %+v", op.Status, op)
		}
	}

	return nil
}

func gceWaitOp(api *compute.Service, proj, opname string) error {
OpLoop:
	for {
		time.Sleep(2 * time.Second)

		op, err := api.GlobalOperations.Get(proj, opname).Do()
		if err != nil {
			return fmt.Errorf("Failed to get op %s: %v", opname, err)
		}
		switch op.Status {
		case "PENDING", "RUNNING":
			continue
		case "DONE":
			if op.Error != nil {
				for _, operr := range op.Error.Errors {
					return fmt.Errorf("Error creating instance: %+v", operr)
				}
				return fmt.Errorf("Failed to start.")
			}
			break OpLoop
		default:
			return fmt.Errorf("Unknown create status %q: %+v", op.Status, op)
		}
	}

	return nil

}

// newName returns a random name prefixed by BaseName
func newName(opts *GCEOptions) (string, error) {
	base := opts.BaseName

	randBytes := make([]byte, 16) //128 bits of entropy
	_, err := rand.Read(randBytes)
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("%v-%x", base, randBytes), nil
}

// Taken from: https://github.com/golang/build/blob/master/buildlet/gce.go#L323
func instanceIPs(inst *compute.Instance) (intIP, extIP string) {
	for _, iface := range inst.NetworkInterfaces {
		if strings.HasPrefix(iface.NetworkIP, "10.") {
			intIP = iface.NetworkIP
		}
		for _, accessConfig := range iface.AccessConfigs {
			if accessConfig.Type == "ONE_TO_ONE_NAT" {
				extIP = accessConfig.NatIP
			}
		}
	}
	return
}
<commit_msg>platform: fix compatibility with latest Google API bindings<commit_after>// Copyright 2015 CoreOS, Inc.
// Copyright 2015 The Go Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package platform

import (
	"crypto/rand"
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/coreos/mantle/auth"

	"github.com/coreos/mantle/Godeps/_workspace/src/golang.org/x/crypto/ssh"
	"github.com/coreos/mantle/Godeps/_workspace/src/google.golang.org/api/compute/v1"
)

type GCEOptions struct {
	Image       string
	Project     string
	Zone        string
	MachineType string
	DiskType    string
	BaseName    string
	Network     string
	ServiceAuth bool
}

type gceCluster struct {
	*baseCluster
	conf *GCEOptions
	api  *compute.Service
}

type gceMachine struct {
	gc    *gceCluster
	name  string
	intIP string
	extIP string
}

func NewGCECluster(conf GCEOptions) (Cluster, error) {
	var client *http.Client
	var err error
	if conf.ServiceAuth {
		client = auth.GoogleServiceClient()
	} else {
		client, err = auth.GoogleClient()
	}
	if err != nil {
		return nil, err
	}

	api, err := compute.New(client)
	if err != nil {
		return nil, err
	}

	bc, err := newBaseCluster()
	if err != nil {
		return nil, err
	}

	gc := &gceCluster{
		baseCluster: bc,
		api:         api,
		conf:        &conf,
	}

	return gc, nil
}

func (gc *gceCluster) Destroy() error {
	for _, gm := range gc.Machines() {
		gm.Destroy()
	}
	gc.agent.Close()
	return nil
}

// Calling in parallel is ok
func (gc *gceCluster) NewMachine(userdata string) (Machine, error) {
	conf, err := NewConf(userdata)
	if err != nil {
		return nil, err
	}

	keys, err := gc.agent.List()
	if err != nil {
		return nil, err
	}

	conf.CopyKeys(keys)

	// Create gce VM and wait for creation to succeed.
	gm, err := GCECreateVM(gc.api, gc.conf, conf.String())
	if err != nil {
		return nil, err
	}
	gm.gc = gc

	if err := commonMachineChecks(gm); err != nil {
		gm.Destroy()
		return nil, err
	}

	gc.addMach(gm)

	return Machine(gm), nil
}

func (gm *gceMachine) ID() string {
	return gm.name
}

func (gm *gceMachine) IP() string {
	return gm.extIP
}

func (gm *gceMachine) PrivateIP() string {
	return gm.intIP
}

func (gm *gceMachine) SSHClient() (*ssh.Client, error) {
	return gm.gc.SSHClient(gm.IP())
}

func (gm *gceMachine) SSH(cmd string) ([]byte, error) {
	return gm.gc.SSH(gm, cmd)
}

func (gm *gceMachine) Destroy() error {
	_, err := gm.gc.api.Instances.Delete(gm.gc.conf.Project, gm.gc.conf.Zone, gm.name).Do()
	if err != nil {
		return err
	}

	gm.gc.delMach(gm)

	return nil
}

func GCECreateVM(api *compute.Service, opts *GCEOptions, userdata string) (*gceMachine, error) {
	// generate name
	name, err := newName(opts)
	if err != nil {
		return nil, fmt.Errorf("Failed allocating unique name for vm: %v\n", err)
	}

	instance, err := gceMakeInstance(opts, userdata, name)
	if err != nil {
		return nil, err
	}

	// request instance
	op, err := api.Instances.Insert(opts.Project, opts.Zone, instance).Do()
	if err != nil {
		return nil, fmt.Errorf("Failed to create new VM: %v\n", err)
	}

	fmt.Fprintf(os.Stderr, "Instance %v requested\n", name)
	fmt.Fprintf(os.Stderr, "Waiting for creation to finish...\n")

	// wait for creation to finish
	err = gceWaitVM(api, opts.Project, opts.Zone, op.Name)
	if err != nil {
		return nil, err
	}

	inst, err := api.Instances.Get(opts.Project, opts.Zone, name).Do()
	if err != nil {
		return nil, fmt.Errorf("Error getting instance %s details after creation: %v", name, err)
	}
	intIP, extIP := instanceIPs(inst)

	gm := &gceMachine{
		name:  name,
		extIP: extIP,
		intIP: intIP,
	}

	return gm, nil
}

func GCEDestroyVM(api *compute.Service, proj, zone, name string) error {
	_, err := api.Instances.Delete(proj, zone, name).Do()
	if err != nil {
		return err
	}
	return nil
}

// Create image on GCE and and wait for completion. Will not overwrite
// existing image.
func GCECreateImage(api *compute.Service, proj, name, source string) error {
	image := &compute.Image{
		Name: name,
		RawDisk: &compute.ImageRawDisk{
			Source: source,
		},
	}

	fmt.Fprintf(os.Stderr, "Image %v requested\n", name)
	fmt.Fprintf(os.Stderr, "Waiting for image creation to finish...\n")

	op, err := api.Images.Insert(proj, image).Do()
	if err != nil {
		return err
	}

	err = gceWaitOp(api, proj, op.Name)
	if err != nil {
		return err
	}
	return nil
}

// Delete image on GCE and then recreate it.
func GCEForceCreateImage(api *compute.Service, proj, name, source string) error {
	// op xor err = nil
	op, err := api.Images.Delete(proj, name).Do()

	if op != nil {
		err = gceWaitOp(api, proj, op.Name)
		if err != nil {
			return err
		}
	}

	// don't return error when delete fails because image doesn't exist
	if err != nil && !strings.HasSuffix(err.Error(), "notFound") {
		return fmt.Errorf("deleting image: %v", err)
	}

	// create
	return GCECreateImage(api, proj, name, source)
}

func GCEListVMs(api *compute.Service, opts *GCEOptions, prefix string) ([]Machine, error) {
	var vms []Machine

	list, err := api.Instances.List(opts.Project, opts.Zone).Do()
	if err != nil {
		return nil, err
	}

	for _, inst := range list.Items {
		if !strings.HasPrefix(inst.Name, prefix) {
			continue
		}
		intIP, extIP := instanceIPs(inst)
		gm := &gceMachine{
			name:  inst.Name,
			extIP: extIP,
			intIP: intIP,
		}

		vms = append(vms, gm)
	}
	return vms, nil
}

func GCEListImages(client *http.Client, proj, prefix string) ([]string, error) {
	var images []string
	computeService, err := compute.New(client)
	list, err := computeService.Images.List(proj).Do()
	if err != nil {
		return nil, err
	}

	for _, image := range list.Items {
		if !strings.HasPrefix(image.Name, prefix) {
			continue
		}
		images = append(images, image.Name)
	}
	return images, nil
}

//Some code taken from: https://github.com/golang/build/blob/master/buildlet/gce.go
func gceMakeInstance(opts *GCEOptions, userdata string, name string) (*compute.Instance, error) {
	prefix := "https://www.googleapis.com/compute/v1/projects/" + opts.Project
	instance := &compute.Instance{
		Name:        name,
		MachineType: prefix + "/zones/" + opts.Zone + "/machineTypes/" + opts.MachineType,
		Metadata:    &compute.Metadata{},
		Disks: []*compute.AttachedDisk{
			{
				AutoDelete: true,
				Boot:       true,
				Type:       "PERSISTENT",
				InitializeParams: &compute.AttachedDiskInitializeParams{
					DiskName:    name,
					SourceImage: prefix + "/global/images/" + opts.Image,
					DiskType:    "/zones/" + opts.Zone + "/diskTypes/" + opts.DiskType,
				},
			},
		},
		NetworkInterfaces: []*compute.NetworkInterface{
			&compute.NetworkInterface{
				AccessConfigs: []*compute.AccessConfig{
					&compute.AccessConfig{
						Type: "ONE_TO_ONE_NAT",
						Name: "External NAT",
					},
				},
				Network: prefix + "/global/networks/" + opts.Network,
			},
		},
	}
	// add cloud config
	if userdata != "" {
		instance.Metadata.Items = append(instance.Metadata.Items, &compute.MetadataItems{
			Key:   "user-data",
			Value: &userdata,
		})
	}

	return instance, nil
}

//Some code taken from: https://github.com/golang/build/blob/master/buildlet/gce.go
func gceWaitVM(api *compute.Service, proj, zone, opname string) error {
OpLoop:
	for {
		time.Sleep(2 * time.Second)

		op, err := api.ZoneOperations.Get(proj, zone, opname).Do()
		if err != nil {
			return fmt.Errorf("Failed to get op %s: %v", opname, err)
		}
		switch op.Status {
		case "PENDING", "RUNNING":
			continue
		case "DONE":
			if op.Error != nil {
				for _, operr := range op.Error.Errors {
					return fmt.Errorf("Error creating instance: %+v", operr)
				}
				return fmt.Errorf("Failed to start.")
			}
			break OpLoop
		default:
			return fmt.Errorf("Unknown create status %q: %+v", op.Status, op)
		}
	}

	return nil
}

func gceWaitOp(api *compute.Service, proj, opname string) error {
OpLoop:
	for {
		time.Sleep(2 * time.Second)

		op, err := api.GlobalOperations.Get(proj, opname).Do()
		if err != nil {
			return fmt.Errorf("Failed to get op %s: %v", opname, err)
		}
		switch op.Status {
		case "PENDING", "RUNNING":
			continue
		case "DONE":
			if op.Error != nil {
				for _, operr := range op.Error.Errors {
					return fmt.Errorf("Error creating instance: %+v", operr)
				}
				return fmt.Errorf("Failed to start.")
			}
			break OpLoop
		default:
			return fmt.Errorf("Unknown create status %q: %+v", op.Status, op)
		}
	}

	return nil

}

// newName returns a random name prefixed by BaseName
func newName(opts *GCEOptions) (string, error) {
	base := opts.BaseName

	randBytes := make([]byte, 16) //128 bits of entropy
	_, err := rand.Read(randBytes)
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("%v-%x", base, randBytes), nil
}

// Taken from: https://github.com/golang/build/blob/master/buildlet/gce.go#L323
func instanceIPs(inst *compute.Instance) (intIP, extIP string) {
	for _, iface := range inst.NetworkInterfaces {
		if strings.HasPrefix(iface.NetworkIP, "10.") {
			intIP = iface.NetworkIP
		}
		for _, accessConfig := range iface.AccessConfigs {
			if accessConfig.Type == "ONE_TO_ONE_NAT" {
				extIP = accessConfig.NatIP
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>package vault

import (
	"reflect"
	"testing"
)

func TestCapabilities(t *testing.T) {
	c, _, token := TestCoreUnsealed(t)

	actual, err := c.Capabilities(token, "path")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	expected := []string{"root"}
	if !reflect.DeepEqual(actual, expected) {
		t.Fatalf("bad: got\n%#v\nexpected\n%#v\n", actual, expected)
	}

	// Create a policy
	policy, _ := Parse(aclPolicy)
	err = c.policyStore.SetPolicy(policy)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a token for the policy
	ent := &TokenEntry{
		ID:       "capabilitiestoken",
		Path:     "testpath",
		Policies: []string{"dev"},
	}
	if err := c.tokenStore.create(ent); err != nil {
		t.Fatalf("err: %v", err)
	}

	actual, err = c.Capabilities("capabilitiestoken", "foo/bar")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	expected = []string{"sudo", "read", "create"}
	if !reflect.DeepEqual(actual, expected) {
		t.Fatalf("bad: got\n%#v\nexpected\n%#v\n", actual, expected)
	}
}
<commit_msg>Fix capabilities test case<commit_after>package vault

import (
	"reflect"
	"testing"
)

func TestCapabilities(t *testing.T) {
	c, _, token := TestCoreUnsealed(t)

	actual, err := c.Capabilities(token, "path")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	expected := []string{"root"}
	if !reflect.DeepEqual(actual, expected) {
		t.Fatalf("bad: got\n%#v\nexpected\n%#v\n", actual, expected)
	}

	// Create a policy
	policy, _ := Parse(aclPolicy)
	err = c.policyStore.SetPolicy(policy)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a token for the policy
	ent := &TokenEntry{
		ID:       "capabilitiestoken",
		Path:     "testpath",
		Policies: []string{"dev"},
	}
	if err := c.tokenStore.create(ent); err != nil {
		t.Fatalf("err: %v", err)
	}

	actual, err = c.Capabilities("capabilitiestoken", "foo/bar")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	expected = []string{"create", "read", "sudo"}
	if !reflect.DeepEqual(actual, expected) {
		t.Fatalf("bad: got\n%#v\nexpected\n%#v\n", actual, expected)
	}
}
<|endoftext|>


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

package server

import (
	"fmt"
	"os"

	"github.com/jasonish/evebox/config"
	"github.com/jasonish/evebox/core"
	"github.com/jasonish/evebox/elasticsearch"
	"github.com/jasonish/evebox/log"
	"github.com/jasonish/evebox/server"
	"github.com/jasonish/evebox/sqlite"
	"github.com/jessevdk/go-flags"
)

const DEFAULT_ELASTICSEARCH_URL string = "http://localhost:9200"

var opts struct {
	// We don't provide a default for this one so we can easily
	// detect if its been set or not.
	ElasticSearchUri   string `long:"elasticsearch" short:"e" description:"Elastic Search URI (default: http://localhost:9200)"`
	ElasticSearchIndex string `long:"index" short:"i" description:"Elastic Search Index (default: logstash)"`
	Port               string `long:"port" short:"p" default:"5636" description:"Port to bind to"`
	Host               string `long:"host" default:"0.0.0.0" description:"Host to bind to"`
	DevServerUri       string `long:"dev" description:"Frontend development server URI"`
	Version            bool   `long:"version" description:"Show version"`
	Config             string `long:"config" short:"c" description:"Configuration filename"`
	NoCheckCertificate bool   `long:"no-check-certificate" short:"k" description:"Disable certificate check for Elastic Search"`
}

var conf *config.Config

func init() {
	conf = config.NewConfig()
}

func VersionMain() {
	fmt.Printf("EveBox Version %s (rev %s) [%s]\n",
		core.BuildVersion, core.BuildRev, core.BuildDate)
}

func getElasticSearchUrl() string {
	if opts.ElasticSearchUri != "" {
		return opts.ElasticSearchUri
	}
	if os.Getenv("ELASTICSEARCH_URL") != "" {
		return os.Getenv("ELASTICSEARCH_URL")
	}
	return DEFAULT_ELASTICSEARCH_URL
}

func getElasticSearchIndex() string {
	if opts.ElasticSearchIndex != "" {
		return opts.ElasticSearchIndex
	} else if os.Getenv("ELASTICSEARCH_INDEX") != "" {
		return os.Getenv("ELASTICSEARCH_INDEX")
	} else {
		return "logstash"
	}
}

func Main(args []string) {

	_, err := flags.ParseArgs(&opts, args)
	if err != nil {
		// flags.Parse should have already presented an error message.
		os.Exit(1)
	}

	if opts.Version {
		VersionMain()
		return
	}

	log.SetLevel(log.DEBUG)

	// If no configuration was provided, see if evebox.yaml exists
	// in the current directory.
	if opts.Config == "" {
		_, err = os.Stat("./evebox.yaml")
		if err == nil {
			opts.Config = "./evebox.yaml"
		}
	}
	if opts.Config != "" {
		log.Printf("Loading configuration file %s.\n", opts.Config)
		conf, err = config.LoadConfig(opts.Config)
		if err != nil {
			log.Fatal(err)
		}
	}

	conf.ElasticSearchIndex = getElasticSearchIndex()
	log.Info("Using ElasticSearch Index %s.", conf.ElasticSearchIndex)

	appContext := server.AppContext{
		Config: conf,
	}
	elasticSearch := elasticsearch.New(getElasticSearchUrl())
	elasticSearch.SetEventIndex(conf.ElasticSearchIndex)
	pingResponse, err := elasticSearch.Ping()
	if err != nil {
		log.Error("Failed to ping Elastic Search: %v", err)
	} else {
		log.Info("Connected to Elastic Search (version: %s)",
			pingResponse.Version.Number)
	}
	appContext.ElasticSearch = elasticSearch
	appContext.EventService = elasticsearch.NewEventService(elasticSearch)
	appContext.AlertQueryService = elasticsearch.NewAlertQueryService(elasticSearch)
	appContext.EventQueryService = elasticsearch.NewEventQueryService(elasticSearch)
	appContext.ReportService = elasticsearch.NewReportService(elasticSearch)

	//dataStoreType := "elasticsearch"
	dataStoreType := "sqlite"

	if dataStoreType == "elasticsearch" {
		appContext.DataStore, err = elasticsearch.NewDataStore(elasticSearch)
		if err != nil {
			log.Fatal(err)
		}
	} else if dataStoreType == "sqlite" {
		appContext.DataStore, err = sqlite.NewDataStore()
		if err != nil {
			log.Fatal(err)
		}
	}

	router := server.NewRouter()

	router.Handle("/api/1/archive",
		server.ApiF(appContext, server.ArchiveHandler))
	router.Handle("/api/1/escalate",
		server.ApiF(appContext, server.EscalateHandler))

	router.POST("/api/1/alert-group/add-tags",
		server.ApiF(appContext, server.AlertGroupAddTags))
	router.POST("/api/1/alert-group/remove-tags",
		server.ApiF(appContext, server.AlertGroupRemoveTags))

	router.Handle("/api/1/event/{id}",
		server.ApiF(appContext, server.GetEventByIdHandler))

	router.POST("/api/1/event/{id}/archive", server.ApiF(appContext, server.ArchiveEventHandler))
	router.POST("/api/1/event/{id}/escalate", server.ApiF(appContext, server.EscalateEventHandler))
	router.POST("/api/1/event/{id}/de-escalate", server.ApiF(appContext, server.DeEscalateEventHandler))

	router.Handle("/api/1/config",
		server.ApiF(appContext, server.ConfigHandler))
	router.Handle("/api/1/version",
		server.ApiF(appContext, server.VersionHandler))
	router.Handle("/api/1/eve2pcap", server.ApiF(appContext, server.Eve2PcapHandler))

	router.GET("/api/1/alerts", server.ApiF(appContext, server.AlertsHandler))
	router.GET("/api/1/event-query", server.ApiF(appContext, server.EventQueryHandler))

	router.Handle("/api/1/query", server.ApiF(appContext, server.QueryHandler))

	router.Handle("/api/1/_bulk", server.ApiF(appContext, server.EsBulkHandler))

	router.GET("/api/1/report/dns/requests/rrnames", server.ApiF(appContext, server.ReportDnsRequestRrnames))
	router.POST("/api/1/report/dns/requests/rrnames", server.ApiF(appContext, server.ReportDnsRequestRrnames))

	router.GET("/api/1/report/agg", server.ApiF(appContext, server.ReportAggs))
	router.GET("/api/1/report/histogram", server.ApiF(appContext, server.ReportHistogram))

	// /api/1/report/netflow/sources/bytes
	// /api/1/report/netflow/sources/packets

	// /api/1/report/netflow/destinations/bytes
	// /api/1/report/netflow/destinations/packets

	// This all needs some cleanup...

	httpServer := server.NewServer(appContext, router)
	httpServer.RegisterApiHandlers()

	// Static file server, must be last as it serves as the fallback.
	router.Prefix("/", server.StaticHandlerFactory(opts.DevServerUri))

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

package server

import (
	"fmt"
	"os"

	"github.com/jasonish/evebox/config"
	"github.com/jasonish/evebox/core"
	"github.com/jasonish/evebox/elasticsearch"
	"github.com/jasonish/evebox/log"
	"github.com/jasonish/evebox/server"
	"github.com/jasonish/evebox/sqlite"
	"github.com/jessevdk/go-flags"
)

const DEFAULT_ELASTICSEARCH_URL string = "http://localhost:9200"

var opts struct {
	// We don't provide a default for this one so we can easily
	// detect if its been set or not.
	ElasticSearchUri   string `long:"elasticsearch" short:"e" description:"Elastic Search URI (default: http://localhost:9200)"`
	ElasticSearchIndex string `long:"index" short:"i" description:"Elastic Search Index (default: logstash)"`
	Port               string `long:"port" short:"p" default:"5636" description:"Port to bind to"`
	Host               string `long:"host" default:"0.0.0.0" description:"Host to bind to"`
	DevServerUri       string `long:"dev" description:"Frontend development server URI"`
	Version            bool   `long:"version" description:"Show version"`
	Config             string `long:"config" short:"c" description:"Configuration filename"`
	NoCheckCertificate bool   `long:"no-check-certificate" short:"k" description:"Disable certificate check for Elastic Search"`
}

var conf *config.Config

func init() {
	conf = config.NewConfig()
}

func VersionMain() {
	fmt.Printf("EveBox Version %s (rev %s) [%s]\n",
		core.BuildVersion, core.BuildRev, core.BuildDate)
}

func getElasticSearchUrl() string {
	if opts.ElasticSearchUri != "" {
		return opts.ElasticSearchUri
	}
	if os.Getenv("ELASTICSEARCH_URL") != "" {
		return os.Getenv("ELASTICSEARCH_URL")
	}
	return DEFAULT_ELASTICSEARCH_URL
}

func getElasticSearchIndex() string {
	if opts.ElasticSearchIndex != "" {
		return opts.ElasticSearchIndex
	} else if os.Getenv("ELASTICSEARCH_INDEX") != "" {
		return os.Getenv("ELASTICSEARCH_INDEX")
	} else {
		return "logstash"
	}
}

func Main(args []string) {

	_, err := flags.ParseArgs(&opts, args)
	if err != nil {
		// flags.Parse should have already presented an error message.
		os.Exit(1)
	}

	if opts.Version {
		VersionMain()
		return
	}

	log.SetLevel(log.DEBUG)

	// If no configuration was provided, see if evebox.yaml exists
	// in the current directory.
	if opts.Config == "" {
		_, err = os.Stat("./evebox.yaml")
		if err == nil {
			opts.Config = "./evebox.yaml"
		}
	}
	if opts.Config != "" {
		log.Printf("Loading configuration file %s.\n", opts.Config)
		conf, err = config.LoadConfig(opts.Config)
		if err != nil {
			log.Fatal(err)
		}
	}

	conf.ElasticSearchIndex = getElasticSearchIndex()
	log.Info("Using ElasticSearch Index %s.", conf.ElasticSearchIndex)

	appContext := server.AppContext{
		Config: conf,
	}
	elasticSearch := elasticsearch.New(getElasticSearchUrl())
	elasticSearch.SetEventIndex(conf.ElasticSearchIndex)
	pingResponse, err := elasticSearch.Ping()
	if err != nil {
		log.Error("Failed to ping Elastic Search: %v", err)
	} else {
		log.Info("Connected to Elastic Search (version: %s)",
			pingResponse.Version.Number)
	}
	appContext.ElasticSearch = elasticSearch
	appContext.EventService = elasticsearch.NewEventService(elasticSearch)
	appContext.AlertQueryService = elasticsearch.NewAlertQueryService(elasticSearch)
	appContext.EventQueryService = elasticsearch.NewEventQueryService(elasticSearch)
	appContext.ReportService = elasticsearch.NewReportService(elasticSearch)

	dataStoreType := "elasticsearch"
	//dataStoreType := "sqlite"

	if dataStoreType == "elasticsearch" {
		appContext.DataStore, err = elasticsearch.NewDataStore(elasticSearch)
		if err != nil {
			log.Fatal(err)
		}
	} else if dataStoreType == "sqlite" {
		appContext.DataStore, err = sqlite.NewDataStore()
		if err != nil {
			log.Fatal(err)
		}
	}

	router := server.NewRouter()

	router.Handle("/api/1/archive",
		server.ApiF(appContext, server.ArchiveHandler))
	router.Handle("/api/1/escalate",
		server.ApiF(appContext, server.EscalateHandler))

	router.POST("/api/1/alert-group/add-tags",
		server.ApiF(appContext, server.AlertGroupAddTags))
	router.POST("/api/1/alert-group/remove-tags",
		server.ApiF(appContext, server.AlertGroupRemoveTags))

	router.Handle("/api/1/event/{id}",
		server.ApiF(appContext, server.GetEventByIdHandler))

	router.POST("/api/1/event/{id}/archive", server.ApiF(appContext, server.ArchiveEventHandler))
	router.POST("/api/1/event/{id}/escalate", server.ApiF(appContext, server.EscalateEventHandler))
	router.POST("/api/1/event/{id}/de-escalate", server.ApiF(appContext, server.DeEscalateEventHandler))

	router.Handle("/api/1/config",
		server.ApiF(appContext, server.ConfigHandler))
	router.Handle("/api/1/version",
		server.ApiF(appContext, server.VersionHandler))
	router.Handle("/api/1/eve2pcap", server.ApiF(appContext, server.Eve2PcapHandler))

	router.GET("/api/1/alerts", server.ApiF(appContext, server.AlertsHandler))
	router.GET("/api/1/event-query", server.ApiF(appContext, server.EventQueryHandler))

	router.Handle("/api/1/query", server.ApiF(appContext, server.QueryHandler))

	router.Handle("/api/1/_bulk", server.ApiF(appContext, server.EsBulkHandler))

	router.GET("/api/1/report/dns/requests/rrnames", server.ApiF(appContext, server.ReportDnsRequestRrnames))
	router.POST("/api/1/report/dns/requests/rrnames", server.ApiF(appContext, server.ReportDnsRequestRrnames))

	router.GET("/api/1/report/agg", server.ApiF(appContext, server.ReportAggs))
	router.GET("/api/1/report/histogram", server.ApiF(appContext, server.ReportHistogram))

	// /api/1/report/netflow/sources/bytes
	// /api/1/report/netflow/sources/packets

	// /api/1/report/netflow/destinations/bytes
	// /api/1/report/netflow/destinations/packets

	// This all needs some cleanup...

	httpServer := server.NewServer(appContext, router)
	httpServer.RegisterApiHandlers()

	// Static file server, must be last as it serves as the fallback.
	router.Prefix("/", server.StaticHandlerFactory(opts.DevServerUri))

	err = httpServer.Start(opts.Host + ":" + opts.Port)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


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

package main

import (
	"bytes"
	"crypto/sha1"
	"fmt"
	"io"
	"launchpad.net/tomb"
	"log"
	"os"
	"path/filepath"
)

type diskIOPeerChans struct {
	// Channels to peers
	writePiece   chan Piece
	requestPiece chan RequestPieceDisk
}

type DiskIO struct {
	metaInfo  MetaInfo
	files     []*os.File
	peerChans diskIOPeerChans
	contChans ControllerDiskIOChans
	t         tomb.Tomb
}

// checkHash accepts a byte buffer and pieceIndex, computes the SHA-1 hash of
// the buffer and returns true or false if it's correct.
func (diskio *DiskIO) checkHash(buf []byte, pieceIndex int) bool {
	h := sha1.New()
	h.Write(buf)
	if bytes.Equal(h.Sum(nil), []byte(diskio.metaInfo.Info.Pieces[pieceIndex:pieceIndex+h.Size()])) {
		return true
	}
	return false
}

// Verify reads in each file and verifies the SHA-1 checksum of each piece.
// Return the boolean list pieces that are correct.
func (diskio *DiskIO) Verify() (finishedPieces []bool) {
	log.Println("DiskIO : Verify : Started")
	defer log.Println("DiskIO : Verify : Completed")

	buf := make([]byte, diskio.metaInfo.Info.PieceLength)
	var pieceIndex, n int
	var err error

	fmt.Printf("Verifying downloaded files")
	if len(diskio.metaInfo.Info.Files) > 0 {
		// Multiple File Mode
		var m int
		// Iterate over each file
		for i, _ := range diskio.metaInfo.Info.Files {
			for offset := int64(0); ; offset += int64(n) {
				// Read from file at offset, up to buf size or
				// less if last read was incomplete due to EOF
				fmt.Printf(".")
				n, err = diskio.files[i].ReadAt(buf[m:], offset)
				if err != nil {
					if err == io.EOF {
						// Reached EOF. Increment partial read counter by bytes read
						m += n
						break
					}
					log.Fatal(err)
				}
				// We have a full buf, check the hash of buf and
				// append the result to the finished pieces
				finishedPieces = append(finishedPieces, diskio.checkHash(buf, pieceIndex))
				// Reset partial read counter
				m = 0
				// Increment piece by the length of a SHA-1 hash (20 bytes)
				pieceIndex += 20
			}
		}
		// If the final iteration resulted in a partial read, then
		// check the hash of it and append the result
		if m > 0 {
			finishedPieces = append(finishedPieces, diskio.checkHash(buf[:m], pieceIndex))
		}
	} else {
		// Single File Mode
		for offset := int64(0); ; offset += int64(n) {
			// Read from file at offset, up to buf size or
			// less if last read was incomplete due to EOF
			fmt.Printf(".")
			n, err = diskio.files[0].ReadAt(buf, offset)
			if err != nil {
				if err == io.EOF {
					// Reached EOF
					break
				}
				log.Fatal(err)
			}
			// We have a full buf, check the hash of buf and
			// append the result to the finished pieces
			finishedPieces = append(finishedPieces, diskio.checkHash(buf, pieceIndex))
			// Increment piece by the length of a SHA-1 hash (20 bytes)
			pieceIndex += 20
		}
		// If the final iteration resulted in a partial read, then compute a hash
		if n > 0 {
			finishedPieces = append(finishedPieces, diskio.checkHash(buf[:n], pieceIndex))
		}
	}
	fmt.Println()

	return finishedPieces
}

func checkError(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

// openOrCreateFile opens the named file or creates it if it doesn't already
// exist. If successful it returns a file handle that can be used for I/O.
func openOrCreateFile(name string) (file *os.File) {
	// Create the file if it doesn't exist
	if _, err := os.Stat(name); os.IsNotExist(err) {
		// Create the file and return a handle
		file, err = os.Create(name)
		checkError(err)
	} else {
		// Open the file and return a handle
		file, err = os.Open(name)
		checkError(err)
	}
	return
}

func NewDiskIO(metaInfo MetaInfo) *DiskIO {
	diskio := new(DiskIO)
	diskio.metaInfo = metaInfo
	diskio.peerChans.writePiece = make(chan Piece)
	diskio.peerChans.requestPiece = make(chan RequestPieceDisk)
	diskio.contChans.receivedPiece = make(chan ReceivedPiece)
	return diskio
}

func (diskio *DiskIO) writePiece(piece Piece) {
	offset := piece.index * diskio.metaInfo.Info.PieceLength

	for i := 0; i <= len(diskio.metaInfo.Info.Files); i++ {
		if offset > diskio.metaInfo.Info.Files[i].Length {
			offset -= diskio.metaInfo.Info.Files[i].Length
		} else {
			max := diskio.metaInfo.Info.Files[i].Length - offset
			n, err := diskio.files[i].WriteAt(piece.data[:max], int64(offset))
			if err != nil {
				log.Fatal(err)
				return
			}
			log.Printf("Wrote %d bytes for piece %x at offset %x, file %s\n", n, piece.index, offset, diskio.metaInfo.Info.Files[i].Path)
			piece.data = piece.data[max:]
			offset = 0
			if len(piece.data) == 0 {
				break
			}
		}
	}
}

func (diskio *DiskIO) Init() {
	log.Println("DiskIO : Init : Started")
	defer log.Println("DiskIO : Init : Completed")

	if len(diskio.metaInfo.Info.Files) > 0 {
		// Multiple File Mode
		directory := diskio.metaInfo.Info.Name
		// Create the directory if it doesn't exist
		if _, err := os.Stat(directory); os.IsNotExist(err) {
			err = os.Mkdir(directory, os.ModeDir|os.ModePerm)
			checkError(err)
		}
		err := os.Chdir(directory)
		checkError(err)
		for _, file := range diskio.metaInfo.Info.Files {
			// Create any sub-directories if required
			if len(file.Path) > 1 {
				directory = filepath.Join(file.Path[1:]...)
				if _, err := os.Stat(directory); os.IsNotExist(err) {
					err = os.MkdirAll(directory, os.ModeDir|os.ModePerm)
					checkError(err)
				}
			}
			// Create the file if it doesn't exist
			name := filepath.Join(file.Path...)
			diskio.files = append(diskio.files, openOrCreateFile(name))
		}
	} else {
		// Single File Mode
		diskio.files = append(diskio.files, openOrCreateFile(diskio.metaInfo.Info.Name))
	}
}

func (diskio *DiskIO) Stop() error {
	log.Println("DiskIO : Stop : Stopping")
	diskio.t.Kill(nil)
	return diskio.t.Wait()
}

func (diskio *DiskIO) Run() {
	log.Println("DiskIO : Run : Started")
	defer diskio.t.Done()
	defer log.Println("DiskIO : Run : Completed")

	for {
		select {
		case piece := <-diskio.peerChans.writePiece:
			go func() {
				diskio.writePiece(piece)
				diskio.contChans.receivedPiece <- ReceivedPiece{pieceNum: piece.index, peerName: piece.peerName}
			}()
		case request := <-diskio.peerChans.requestPiece:
			fmt.Println(request)
		case <-diskio.t.Dying():
			return
		}
	}
}
<commit_msg>Redundant return<commit_after>// Copyright 2013 Jari Takkala. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"crypto/sha1"
	"fmt"
	"io"
	"launchpad.net/tomb"
	"log"
	"os"
	"path/filepath"
)

type diskIOPeerChans struct {
	// Channels to peers
	writePiece   chan Piece
	requestPiece chan RequestPieceDisk
}

type DiskIO struct {
	metaInfo  MetaInfo
	files     []*os.File
	peerChans diskIOPeerChans
	contChans ControllerDiskIOChans
	t         tomb.Tomb
}

// checkHash accepts a byte buffer and pieceIndex, computes the SHA-1 hash of
// the buffer and returns true or false if it's correct.
func (diskio *DiskIO) checkHash(buf []byte, pieceIndex int) bool {
	h := sha1.New()
	h.Write(buf)
	if bytes.Equal(h.Sum(nil), []byte(diskio.metaInfo.Info.Pieces[pieceIndex:pieceIndex+h.Size()])) {
		return true
	}
	return false
}

// Verify reads in each file and verifies the SHA-1 checksum of each piece.
// Return the boolean list pieces that are correct.
func (diskio *DiskIO) Verify() (finishedPieces []bool) {
	log.Println("DiskIO : Verify : Started")
	defer log.Println("DiskIO : Verify : Completed")

	buf := make([]byte, diskio.metaInfo.Info.PieceLength)
	var pieceIndex, n int
	var err error

	fmt.Printf("Verifying downloaded files")
	if len(diskio.metaInfo.Info.Files) > 0 {
		// Multiple File Mode
		var m int
		// Iterate over each file
		for i, _ := range diskio.metaInfo.Info.Files {
			for offset := int64(0); ; offset += int64(n) {
				// Read from file at offset, up to buf size or
				// less if last read was incomplete due to EOF
				fmt.Printf(".")
				n, err = diskio.files[i].ReadAt(buf[m:], offset)
				if err != nil {
					if err == io.EOF {
						// Reached EOF. Increment partial read counter by bytes read
						m += n
						break
					}
					log.Fatal(err)
				}
				// We have a full buf, check the hash of buf and
				// append the result to the finished pieces
				finishedPieces = append(finishedPieces, diskio.checkHash(buf, pieceIndex))
				// Reset partial read counter
				m = 0
				// Increment piece by the length of a SHA-1 hash (20 bytes)
				pieceIndex += 20
			}
		}
		// If the final iteration resulted in a partial read, then
		// check the hash of it and append the result
		if m > 0 {
			finishedPieces = append(finishedPieces, diskio.checkHash(buf[:m], pieceIndex))
		}
	} else {
		// Single File Mode
		for offset := int64(0); ; offset += int64(n) {
			// Read from file at offset, up to buf size or
			// less if last read was incomplete due to EOF
			fmt.Printf(".")
			n, err = diskio.files[0].ReadAt(buf, offset)
			if err != nil {
				if err == io.EOF {
					// Reached EOF
					break
				}
				log.Fatal(err)
			}
			// We have a full buf, check the hash of buf and
			// append the result to the finished pieces
			finishedPieces = append(finishedPieces, diskio.checkHash(buf, pieceIndex))
			// Increment piece by the length of a SHA-1 hash (20 bytes)
			pieceIndex += 20
		}
		// If the final iteration resulted in a partial read, then compute a hash
		if n > 0 {
			finishedPieces = append(finishedPieces, diskio.checkHash(buf[:n], pieceIndex))
		}
	}
	fmt.Println()

	return finishedPieces
}

func checkError(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

// openOrCreateFile opens the named file or creates it if it doesn't already
// exist. If successful it returns a file handle that can be used for I/O.
func openOrCreateFile(name string) (file *os.File) {
	// Create the file if it doesn't exist
	if _, err := os.Stat(name); os.IsNotExist(err) {
		// Create the file and return a handle
		file, err = os.Create(name)
		checkError(err)
	} else {
		// Open the file and return a handle
		file, err = os.Open(name)
		checkError(err)
	}
	return
}

func NewDiskIO(metaInfo MetaInfo) *DiskIO {
	diskio := new(DiskIO)
	diskio.metaInfo = metaInfo
	diskio.peerChans.writePiece = make(chan Piece)
	diskio.peerChans.requestPiece = make(chan RequestPieceDisk)
	diskio.contChans.receivedPiece = make(chan ReceivedPiece)
	return diskio
}

func (diskio *DiskIO) writePiece(piece Piece) {
	offset := piece.index * diskio.metaInfo.Info.PieceLength

	for i := 0; i <= len(diskio.metaInfo.Info.Files); i++ {
		if offset > diskio.metaInfo.Info.Files[i].Length {
			offset -= diskio.metaInfo.Info.Files[i].Length
		} else {
			max := diskio.metaInfo.Info.Files[i].Length - offset
			n, err := diskio.files[i].WriteAt(piece.data[:max], int64(offset))
			if err != nil {
				log.Fatal(err)
			}
			log.Printf("Wrote %d bytes for piece %x at offset %x, file %s\n", n, piece.index, offset, diskio.metaInfo.Info.Files[i].Path)
			piece.data = piece.data[max:]
			offset = 0
			if len(piece.data) == 0 {
				break
			}
		}
	}
}

func (diskio *DiskIO) Init() {
	log.Println("DiskIO : Init : Started")
	defer log.Println("DiskIO : Init : Completed")

	if len(diskio.metaInfo.Info.Files) > 0 {
		// Multiple File Mode
		directory := diskio.metaInfo.Info.Name
		// Create the directory if it doesn't exist
		if _, err := os.Stat(directory); os.IsNotExist(err) {
			err = os.Mkdir(directory, os.ModeDir|os.ModePerm)
			checkError(err)
		}
		err := os.Chdir(directory)
		checkError(err)
		for _, file := range diskio.metaInfo.Info.Files {
			// Create any sub-directories if required
			if len(file.Path) > 1 {
				directory = filepath.Join(file.Path[1:]...)
				if _, err := os.Stat(directory); os.IsNotExist(err) {
					err = os.MkdirAll(directory, os.ModeDir|os.ModePerm)
					checkError(err)
				}
			}
			// Create the file if it doesn't exist
			name := filepath.Join(file.Path...)
			diskio.files = append(diskio.files, openOrCreateFile(name))
		}
	} else {
		// Single File Mode
		diskio.files = append(diskio.files, openOrCreateFile(diskio.metaInfo.Info.Name))
	}
}

func (diskio *DiskIO) Stop() error {
	log.Println("DiskIO : Stop : Stopping")
	diskio.t.Kill(nil)
	return diskio.t.Wait()
}

func (diskio *DiskIO) Run() {
	log.Println("DiskIO : Run : Started")
	defer diskio.t.Done()
	defer log.Println("DiskIO : Run : Completed")

	for {
		select {
		case piece := <-diskio.peerChans.writePiece:
			go func() {
				diskio.writePiece(piece)
				diskio.contChans.receivedPiece <- ReceivedPiece{pieceNum: piece.index, peerName: piece.peerName}
			}()
		case request := <-diskio.peerChans.requestPiece:
			fmt.Println(request)
		case <-diskio.t.Dying():
			return
		}
	}
}
<|endoftext|>


<commit_before>/*package catalog contains funcitons for reading from and writing to
files containing catalogs of particle locations and velocities.

Catalogs written by external programs are standardized to a single type when
read. Currently the only such extrenal catalog type that is supported the
Gadget 2 particle catalog.

The binary format used is as follows:
    |-- 1 --||-- 2 --||-- 3 --||-- ... 4 ... --||-- ... 5 ... --|

    1 - (int32) Flag indicating the endianness of the file. 0 indicates a big
        endian byte ordering and -1 indicates a little endian byte order.
    2 - (int32) Size of a Header struct. Should be checked for consistency.
    3 - (int32) Size of a Particle struct. Should be checked for consistency.
    4 - (tetra.Header) Header file containing meta-information about the
        particle catalog.
    5 - ([]tetra.Particle) Contiguous block of particles. Garuanteed to be
        of size unsafe.Sizeof(Particle{}) * header.Count.
*/
package catalog

import (
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"os"

	"unsafe"

	tetra "github.com/phil-mansfield/gotetra"
)

const (
	// Endianness used by default when writing catalogs. Catalogs of any
	// endianness can be read.
	DefaultEndiannessFlag int32 = -1
)

// gadgetHeader is the formatting for meta-information used by Gadget 2.
type gadgetHeader struct {
	NPart                                     [6]uint32
	Mass                                      [6]float64
	Time, Redshift                            float64
	FlagSfr, FlagFeedback                     int32
	NPartTotal                                [6]uint32
	FlagCooling, NumFiles                     int32
	BoxSize, Omega0, OmegaLambda, HubbleParam float64
	FlagStellarAge, HashTabSize               int32

	Padding [88]byte
}

// readInt32 returns single 32-bit interger from the given file using the
// given endianness.
func readInt32(r io.Reader, order binary.ByteOrder) int32 {
	var n int32
	if err := binary.Read(r, order, &n); err != nil {
		panic(err)
	}
	return n
}

// Standardize returns a tetra.Header that corresponds to the source
// Gadget 2 header.
func (gh *gadgetHeader) Standardize() *tetra.Header {
	h := &tetra.Header{}

	h.Count = int64(gh.NPart[1] + gh.NPart[0]<<32)
	h.TotalCount = int64(gh.NPartTotal[1] + gh.NPartTotal[0]<<32)
	h.Mass = float64(gh.Mass[1])
	h.TotalWidth = float64(gh.BoxSize)
	h.Width = -1.0

	h.Cosmo.Z = gh.Redshift
	h.Cosmo.OmegaM = gh.Omega0
	h.Cosmo.OmegaL = gh.OmegaLambda
	h.Cosmo.H100 = gh.HubbleParam

	return h
}

// WrapDistance takes a value and interprets it as a position defined within
// a periodic domain of width h.BoxSize.
func (h *gadgetHeader) WrapDistance(x float64) float64 {
	if x < 0 {
		return x + h.BoxSize
	} else if x >= h.BoxSize {
		return x - h.BoxSize
	}
	return x
}

func ReadGadgetHeader(path string, order binary.ByteOrder) *tetra.Header {
	f, err := os.Open(path)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	gh := &gadgetHeader{}

	_ = readInt32(f, order)
	binary.Read(f, binary.LittleEndian, gh)
	h := gh.Standardize()

	return h
}

func ReadGadgetParticlesAt(
	path string,
	order binary.ByteOrder,
	floatBuf []float32,
	intBuf []int64,
	ps []tetra.Particle,
) {
	f, err := os.Open(path)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	gh := &gadgetHeader{}

	_ = readInt32(f, order)
	binary.Read(f, binary.LittleEndian, gh)
	_ = readInt32(f, order)

	h := gh.Standardize()

	if int64(len(floatBuf)) != 3*h.Count {
		panic(fmt.Sprintf(
			"Incorrect length for float buffer. Found %d, expected %d",
			len(floatBuf), 3*h.Count,
		))
	} else if int64(len(intBuf)) != h.Count {
		panic(fmt.Sprintf(
			"Incorrect length for int buffer. Found %d, expected %d",
			len(intBuf), h.Count,
		))
	} else if int64(len(ps)) != h.Count {
		panic(fmt.Sprintf(
			"Incorrect length for Particle buffer. Found %d, expected %d",
			len(ps), h.Count,
		))
	}

	_ = readInt32(f, order)
	binary.Read(f, order, floatBuf)
	_ = readInt32(f, order)

	for i := range ps {
		ps[i].Xs[0] = gh.WrapDistance(float64(floatBuf[3*i+0]))
		ps[i].Xs[1] = gh.WrapDistance(float64(floatBuf[3*i+1]))
		ps[i].Xs[2] = gh.WrapDistance(float64(floatBuf[3*i+2]))
	}

	_ = readInt32(f, order)
	binary.Read(f, order, floatBuf)
	_ = readInt32(f, order)

	rootA := float32(math.Sqrt(float64(gh.Time)))
	for i := range ps {
		ps[i].Vs[0] = float64(floatBuf[3*i+0] * rootA)
		ps[i].Vs[1] = float64(floatBuf[3*i+1] * rootA)
		ps[i].Vs[2] = float64(floatBuf[3*i+2] * rootA)
	}

	_ = readInt32(f, order)
	binary.Read(f, order, intBuf)
	_ = readInt32(f, order)

	for i := range ps {
		ps[i].Id = intBuf[i]
	}
}

// ReadGadget reads the gadget particle catalog located at the given location
// and written with the given endianness. Its header and particle sequence
// are returned in a standardized format.
func ReadGadget(path string, order binary.ByteOrder) (*tetra.Header, []tetra.Particle) {
	f, err := os.Open(path)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	gh := &gadgetHeader{}

	_ = readInt32(f, order)
	binary.Read(f, binary.LittleEndian, gh)
	_ = readInt32(f, order)

	h := gh.Standardize()
	floatBuf := make([]float32, 3*h.Count)
	ps := make([]tetra.Particle, h.Count)

	_ = readInt32(f, order)
	binary.Read(f, order, floatBuf)
	_ = readInt32(f, order)

	for i := range ps {
		ps[i].Xs[0] = gh.WrapDistance(float64(floatBuf[3*i+0]))
		ps[i].Xs[1] = gh.WrapDistance(float64(floatBuf[3*i+1]))
		ps[i].Xs[2] = gh.WrapDistance(float64(floatBuf[3*i+2]))
	}

	_ = readInt32(f, order)
	binary.Read(f, order, floatBuf)
	_ = readInt32(f, order)

	rootA := float32(math.Sqrt(float64(gh.Time)))
	for i := range ps {
		ps[i].Vs[0] = float64(floatBuf[3*i+0] * rootA)
		ps[i].Vs[1] = float64(floatBuf[3*i+1] * rootA)
		ps[i].Vs[2] = float64(floatBuf[3*i+2] * rootA)
	}

	ids := make([]int64, h.Count)

	_ = readInt32(f, order)
	binary.Read(f, order, ids)
	_ = readInt32(f, order)

	for i := range ps {
		ps[i].Id = ids[i]
	}

	return h, ps
}

// Write writes the given header and particle sequence to the specified file.
func Write(path string, h *tetra.Header, ps []tetra.Particle) {
	f, err := os.Create(path)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	h.Count = int64(len(ps))
	order := endianness(DefaultEndiannessFlag)

	err = binary.Write(f, order, DefaultEndiannessFlag)
	if err != nil {
		panic(err.Error())
	}
	err = binary.Write(f, order, int32(unsafe.Sizeof(tetra.Header{})))
	if err != nil {
		panic(err.Error())
	}
	err = binary.Write(f, order, int32(unsafe.Sizeof(ps[0])))
	if err != nil {
		panic(err.Error())
	}
	err = binary.Write(f, order, h)
	if err != nil {
		panic(err.Error())
	}
	err = binary.Write(f, order, ps)
	if err != nil {
		panic(err.Error())
	}
}

// readHeader reads only the header from the given file.
func readHeader(path string, flag int) (*tetra.Header, *os.File, binary.ByteOrder) {
	f, err := os.OpenFile(path, flag, os.ModePerm)
	if err != nil {
		panic(err)
	}

	order := endianness(readInt32(f, binary.LittleEndian))

	// Sanity checks:
	headerSize := readInt32(f, order)
	particleSize := readInt32(f, order)
	if int32(unsafe.Sizeof(tetra.Header{})) != headerSize {
		panic(fmt.Sprintf(
			"Size of header in code, %d, does not match size in catalog, %d.",
			unsafe.Sizeof(tetra.Header{}), headerSize,
		))
	} else if int32(unsafe.Sizeof(tetra.Particle{})) != particleSize {
		panic(fmt.Sprintf(
			"Size of header in code, %d, does not match size in catalog, %d.",
			unsafe.Sizeof(tetra.Particle{}), particleSize,
		))
	}

	h := &tetra.Header{}
	err = binary.Read(f, order, h)
	if err != nil {
		panic(err.Error())
	}

	return h, f, order
}

func ReadHeader(path string) *tetra.Header {
	h, f, _ := readHeader(path, os.O_RDONLY)
	if err := f.Close(); err != nil {
		panic(err)
	}
	return h
}

// Append appends a particle sequence to the end of the given file.
func Append(path string, ps []tetra.Particle) {
	if len(ps) == 0 {
		return
	}

	h, f, order := readHeader(path, os.O_RDWR)
	defer f.Close()

	_, err := f.Seek(0, 2)
	if err != nil {
		panic(err.Error())
	}
	err = binary.Write(f, order, ps)
	if err != nil {
		panic(err.Error())
	}

	_, err = f.Seek(12, 0)
	if err != nil {
		panic(err.Error())
	}
	h.Count += int64(len(ps))
	err = binary.Write(f, order, h)
	if err != nil {
		panic(err.Error())
	}
}

// Read reads a header and particle sequence from the given file.
func ReadParticlesAt(path string, ps []tetra.Particle) {
	h, f, order := readHeader(path, os.O_RDONLY)
	defer f.Close()

	if int64(len(ps)) != h.Count {
		panic("Incorrect Particle buffer length.")
	}

	err := binary.Read(f, order, ps)
	if err != nil {
		panic(err.Error())
	}
}

// endianness is a utility function converting an endianness flag to a
// byte order.
func endianness(flag int32) binary.ByteOrder {
	if flag == 0 {
		return binary.LittleEndian
	} else if flag == -1 {
		return binary.BigEndian
	} else {
		panic("Unrecognized endianness flag.")
	}
}
<commit_msg>Added Read funciton to catalog package.<commit_after>/*package catalog contains funcitons for reading from and writing to
files containing catalogs of particle locations and velocities.

Catalogs written by external programs are standardized to a single type when
read. Currently the only such extrenal catalog type that is supported the
Gadget 2 particle catalog.

The binary format used is as follows:
    |-- 1 --||-- 2 --||-- 3 --||-- ... 4 ... --||-- ... 5 ... --|

    1 - (int32) Flag indicating the endianness of the file. 0 indicates a big
        endian byte ordering and -1 indicates a little endian byte order.
    2 - (int32) Size of a Header struct. Should be checked for consistency.
    3 - (int32) Size of a Particle struct. Should be checked for consistency.
    4 - (tetra.Header) Header file containing meta-information about the
        particle catalog.
    5 - ([]tetra.Particle) Contiguous block of particles. Garuanteed to be
        of size unsafe.Sizeof(Particle{}) * header.Count.
*/
package catalog

import (
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"os"

	"unsafe"

	tetra "github.com/phil-mansfield/gotetra"
)

const (
	// Endianness used by default when writing catalogs. Catalogs of any
	// endianness can be read.
	DefaultEndiannessFlag int32 = -1
)

// gadgetHeader is the formatting for meta-information used by Gadget 2.
type gadgetHeader struct {
	NPart                                     [6]uint32
	Mass                                      [6]float64
	Time, Redshift                            float64
	FlagSfr, FlagFeedback                     int32
	NPartTotal                                [6]uint32
	FlagCooling, NumFiles                     int32
	BoxSize, Omega0, OmegaLambda, HubbleParam float64
	FlagStellarAge, HashTabSize               int32

	Padding [88]byte
}

// readInt32 returns single 32-bit interger from the given file using the
// given endianness.
func readInt32(r io.Reader, order binary.ByteOrder) int32 {
	var n int32
	if err := binary.Read(r, order, &n); err != nil {
		panic(err)
	}
	return n
}

// Standardize returns a tetra.Header that corresponds to the source
// Gadget 2 header.
func (gh *gadgetHeader) Standardize() *tetra.Header {
	h := &tetra.Header{}

	h.Count = int64(gh.NPart[1] + gh.NPart[0]<<32)
	h.TotalCount = int64(gh.NPartTotal[1] + gh.NPartTotal[0]<<32)
	h.Mass = float64(gh.Mass[1])
	h.TotalWidth = float64(gh.BoxSize)
	h.Width = -1.0

	h.Cosmo.Z = gh.Redshift
	h.Cosmo.OmegaM = gh.Omega0
	h.Cosmo.OmegaL = gh.OmegaLambda
	h.Cosmo.H100 = gh.HubbleParam

	return h
}

// WrapDistance takes a value and interprets it as a position defined within
// a periodic domain of width h.BoxSize.
func (h *gadgetHeader) WrapDistance(x float64) float64 {
	if x < 0 {
		return x + h.BoxSize
	} else if x >= h.BoxSize {
		return x - h.BoxSize
	}
	return x
}

func ReadGadgetHeader(path string, order binary.ByteOrder) *tetra.Header {
	f, err := os.Open(path)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	gh := &gadgetHeader{}

	_ = readInt32(f, order)
	binary.Read(f, binary.LittleEndian, gh)
	h := gh.Standardize()

	return h
}

func ReadGadgetParticlesAt(
	path string,
	order binary.ByteOrder,
	floatBuf []float32,
	intBuf []int64,
	ps []tetra.Particle,
) {
	f, err := os.Open(path)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	gh := &gadgetHeader{}

	_ = readInt32(f, order)
	binary.Read(f, binary.LittleEndian, gh)
	_ = readInt32(f, order)

	h := gh.Standardize()

	if int64(len(floatBuf)) != 3*h.Count {
		panic(fmt.Sprintf(
			"Incorrect length for float buffer. Found %d, expected %d",
			len(floatBuf), 3*h.Count,
		))
	} else if int64(len(intBuf)) != h.Count {
		panic(fmt.Sprintf(
			"Incorrect length for int buffer. Found %d, expected %d",
			len(intBuf), h.Count,
		))
	} else if int64(len(ps)) != h.Count {
		panic(fmt.Sprintf(
			"Incorrect length for Particle buffer. Found %d, expected %d",
			len(ps), h.Count,
		))
	}

	_ = readInt32(f, order)
	binary.Read(f, order, floatBuf)
	_ = readInt32(f, order)

	for i := range ps {
		ps[i].Xs[0] = gh.WrapDistance(float64(floatBuf[3*i+0]))
		ps[i].Xs[1] = gh.WrapDistance(float64(floatBuf[3*i+1]))
		ps[i].Xs[2] = gh.WrapDistance(float64(floatBuf[3*i+2]))
	}

	_ = readInt32(f, order)
	binary.Read(f, order, floatBuf)
	_ = readInt32(f, order)

	rootA := float32(math.Sqrt(float64(gh.Time)))
	for i := range ps {
		ps[i].Vs[0] = float64(floatBuf[3*i+0] * rootA)
		ps[i].Vs[1] = float64(floatBuf[3*i+1] * rootA)
		ps[i].Vs[2] = float64(floatBuf[3*i+2] * rootA)
	}

	_ = readInt32(f, order)
	binary.Read(f, order, intBuf)
	_ = readInt32(f, order)

	for i := range ps {
		ps[i].Id = intBuf[i]
	}
}

// ReadGadget reads the gadget particle catalog located at the given location
// and written with the given endianness. Its header and particle sequence
// are returned in a standardized format.
func ReadGadget(path string, order binary.ByteOrder) (*tetra.Header, []tetra.Particle) {
	f, err := os.Open(path)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	gh := &gadgetHeader{}

	_ = readInt32(f, order)
	binary.Read(f, binary.LittleEndian, gh)
	_ = readInt32(f, order)

	h := gh.Standardize()
	floatBuf := make([]float32, 3*h.Count)
	ps := make([]tetra.Particle, h.Count)

	_ = readInt32(f, order)
	binary.Read(f, order, floatBuf)
	_ = readInt32(f, order)

	for i := range ps {
		ps[i].Xs[0] = gh.WrapDistance(float64(floatBuf[3*i+0]))
		ps[i].Xs[1] = gh.WrapDistance(float64(floatBuf[3*i+1]))
		ps[i].Xs[2] = gh.WrapDistance(float64(floatBuf[3*i+2]))
	}

	_ = readInt32(f, order)
	binary.Read(f, order, floatBuf)
	_ = readInt32(f, order)

	rootA := float32(math.Sqrt(float64(gh.Time)))
	for i := range ps {
		ps[i].Vs[0] = float64(floatBuf[3*i+0] * rootA)
		ps[i].Vs[1] = float64(floatBuf[3*i+1] * rootA)
		ps[i].Vs[2] = float64(floatBuf[3*i+2] * rootA)
	}

	ids := make([]int64, h.Count)

	_ = readInt32(f, order)
	binary.Read(f, order, ids)
	_ = readInt32(f, order)

	for i := range ps {
		ps[i].Id = ids[i]
	}

	return h, ps
}

// Write writes the given header and particle sequence to the specified file.
func Write(path string, h *tetra.Header, ps []tetra.Particle) {
	f, err := os.Create(path)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	h.Count = int64(len(ps))
	order := endianness(DefaultEndiannessFlag)

	err = binary.Write(f, order, DefaultEndiannessFlag)
	if err != nil {
		panic(err.Error())
	}
	err = binary.Write(f, order, int32(unsafe.Sizeof(tetra.Header{})))
	if err != nil {
		panic(err.Error())
	}
	err = binary.Write(f, order, int32(unsafe.Sizeof(ps[0])))
	if err != nil {
		panic(err.Error())
	}
	err = binary.Write(f, order, h)
	if err != nil {
		panic(err.Error())
	}
	err = binary.Write(f, order, ps)
	if err != nil {
		panic(err.Error())
	}
}

// readHeader reads only the header from the given file.
func readHeader(path string, flag int) (*tetra.Header, *os.File, binary.ByteOrder) {
	f, err := os.OpenFile(path, flag, os.ModePerm)
	if err != nil {
		panic(err)
	}

	order := endianness(readInt32(f, binary.LittleEndian))

	// Sanity checks:
	headerSize := readInt32(f, order)
	particleSize := readInt32(f, order)
	if int32(unsafe.Sizeof(tetra.Header{})) != headerSize {
		panic(fmt.Sprintf(
			"Size of header in code, %d, does not match size in catalog, %d.",
			unsafe.Sizeof(tetra.Header{}), headerSize,
		))
	} else if int32(unsafe.Sizeof(tetra.Particle{})) != particleSize {
		panic(fmt.Sprintf(
			"Size of header in code, %d, does not match size in catalog, %d.",
			unsafe.Sizeof(tetra.Particle{}), particleSize,
		))
	}

	h := &tetra.Header{}
	err = binary.Read(f, order, h)
	if err != nil {
		panic(err.Error())
	}

	return h, f, order
}

func ReadHeader(path string) *tetra.Header {
	h, f, _ := readHeader(path, os.O_RDONLY)
	if err := f.Close(); err != nil {
		panic(err)
	}
	return h
}

// Append appends a particle sequence to the end of the given file.
func Append(path string, ps []tetra.Particle) {
	if len(ps) == 0 {
		return
	}

	h, f, order := readHeader(path, os.O_RDWR)
	defer f.Close()

	_, err := f.Seek(0, 2)
	if err != nil {
		panic(err.Error())
	}
	err = binary.Write(f, order, ps)
	if err != nil {
		panic(err.Error())
	}

	_, err = f.Seek(12, 0)
	if err != nil {
		panic(err.Error())
	}
	h.Count += int64(len(ps))
	err = binary.Write(f, order, h)
	if err != nil {
		panic(err.Error())
	}
}

// Read reads a header and particle sequence from the given file.
func ReadParticlesAt(path string, ps []tetra.Particle) {
	h, f, order := readHeader(path, os.O_RDONLY)
	defer f.Close()

	if int64(len(ps)) != h.Count {
		panic("Incorrect Particle buffer length.")
	}

	err := binary.Read(f, order, ps)
	if err != nil {
		panic(err.Error())
	}
}

func Read(path string) (*tetra.Header, []tetra.Particle) {
	h := ReadHeader(path)
	ps := make([]tetra.Particle, h.Count)
	ReadParticlesAt(path, ps)
	return h, ps
}

// endianness is a utility function converting an endianness flag to a
// byte order.
func endianness(flag int32) binary.ByteOrder {
	if flag == 0 {
		return binary.LittleEndian
	} else if flag == -1 {
		return binary.BigEndian
	} else {
		panic("Unrecognized endianness flag.")
	}
}
<|endoftext|>


<commit_before>package daemon

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	_ "github.com/go-sql-driver/mysql"
	"github.com/prism-river/killy/collect"
)

// TCPMessage defines what a message that can be
// sent or received to/from LUA scripts
type TCPMessage struct {
	Cmd  string   `json:"cmd,omitempty"`
	Args []string `json:"args,omitempty"`
	// Id is used to associate requests & responses
	ID   int         `json:"id,omitempty"`
	Data interface{} `json:"data,omitempty"`
}

// ContainerEvent is one kind of Data that can
// be transported by a TCPMessage in the Data field.
// It describes a Docker container event. (start, stop, destroy...)
type TidbEvent struct {
	Action string `json:"action,omitempty"`
	ID     string `json:"id,omitempty"`
	Name   string `json:"name,omitempty"`
	CPU    string `json:"cpu,omitempty"`
	RAM    string `json:"ram,omitempty"`
	sync.RWMutex
}

// Table represents a table in TiDB
type Table struct {
	Name    string     `json:"name"`
	Columns []string   `json:"columns"`
	Data    [][]string `json:"data"`
}

// Config for the daemon
type Config struct {
	Database struct {
		User     string `json:"user"`
		Password string `json:"password"`
		Address  string `json:"address"`
		Name     string `json:"name"`
	} `json:"database"`
	Interval int `json:"interval"` // in seconds
}

// Daemon maintains state when the dockercraft daemon is running
type Daemon struct {
	Client *collect.Collect
	// The configuration
	Config *Config

	// tcpMessages can be used to send bytes to the Lua
	// plugin from any go routine.
	tcpMessages chan []byte

	sync.Mutex
}

// NewDaemon returns a new instance of Daemon
func NewDaemon(address string) *Daemon {
	client := collect.NewCollect("http://" + address + "/api/v1/query?query=")
	return &Daemon{
		Client: client,
	}
}

// Init initializes a Daemon
func (d *Daemon) Init() error {
	var err error
	// load configuration
	d.Config = new(Config)
	var configFile *os.File
	configFile, err = os.Open("config.json")
	defer configFile.Close()
	if err != nil {
		log.Fatal(err.Error())
	}
	jsonParser := json.NewDecoder(configFile)
	jsonParser.Decode(d.Config)

	d.tcpMessages = make(chan []byte)

	return nil
}

// Serve exposes a TCP server on port 25566 to handle
// connections from the LUA scripts
func (d *Daemon) Serve() {

	tcpAddr, err := net.ResolveTCPAddr("tcp", ":25566")

	ln, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		log.Fatalln("listen tcp error:", err)
	}
	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Fatalln("tcp conn accept error:", err)
		}
		// no need to handle connection in a go routine
		// goproxy is used as support for one single Lua plugin.
		d.handleConn(conn)
	}
}

// StartMonitoringEvents listens for events from the
// Docker daemon and uses callback to transmit them
// to LUA scripts.
func (d *Daemon) StartMonitoringEvents() {
	log.Info("Monitoring Database Events")

	// mysql test
	go func() {
		db, err := sql.Open("mysql", fmt.Sprintf("%v:%v@tcp(%v)/%v", d.Config.Database.User, d.Config.Database.Password, d.Config.Database.Address, d.Config.Database.Name))
		if err != nil {
			log.Fatal(err)
		}
		defer db.Close()

		for {
			tables := make([]string, 0)
			rows, err := db.Query("SHOW TABLES")
			if err != nil {
				log.Fatal(err)
			}
			defer rows.Close()
			for rows.Next() {
				var name string
				err := rows.Scan(&name)
				if err != nil {
					log.Fatal(err)
				}
				tables = append(tables, name)
			}
			err = rows.Err()
			if err != nil {
				log.Fatal(err)
			}

			res := make([]Table, 0)
			for _, tableName := range tables {
				rows, err := db.Query("SELECT * FROM " + tableName)
				if err != nil {
					log.Fatal(err)
				}
				defer rows.Close()
				columns, err := rows.Columns()
				if err != nil {
					log.Fatal(err)
				}
				table := Table{
					Name:    tableName,
					Columns: columns,
					Data:    make([][]string, 0),
				}
				for rows.Next() {
					fields := make([]string, len(columns))
					pointers := make([]interface{}, len(columns))
					for i := range fields {
						pointers[i] = &fields[i]
					}
					err := rows.Scan(pointers...)
					if err != nil {
						log.Fatal(err)
					}
					table.Data = append(table.Data, fields)
				}
				err = rows.Err()
				if err != nil {
					log.Fatal(err)
				}
				res = append(res, table)
			}

			tcpMsg := TCPMessage{}
			tcpMsg.Cmd = "event"
			tcpMsg.Args = []string{"table"}
			tcpMsg.ID = 0
			tcpMsg.Data = &res

			data, err := json.Marshal(&tcpMsg)
			if err != nil {
				log.Println("statCallback error:", err)
				return
			}

			separator := []byte(string('\n'))

			d.tcpMessages <- append(data, separator...)

			time.Sleep(time.Duration(d.Config.Interval) * time.Second)
		}
	}()
}

// handleConn handles a TCP connection
// with a Dockercraft Lua plugin.
func (d *Daemon) handleConn(conn net.Conn) {

	go func() {
		separator := []byte(string('\n'))

		buf := make([]byte, 256)
		cursor := 0
		for {
			// resize buf if needed
			if len(buf)-cursor < 256 {
				buf = append(buf, make([]byte, 256-(len(buf)-cursor))...)
			}
			n, err := conn.Read(buf[cursor:])
			if err != nil && err != io.EOF {
				log.Fatalln("conn read error: ", err)
			}
			cursor += n

			// TODO(aduermael): check cNetwork plugin implementation
			// conn.Read doesn't seem to be blocking if there's nothing
			// to read. Maybe the broken pipe is due to an implementation
			// problem on cNetwork plugin side
			if cursor == 0 {
				<-time.After(500 * time.Millisecond)
				continue
			}
			// log.Println("TCP data read:", string(buf[:cursor]), "cursor:", cursor)

			// see if there's a complete json message in buf.
			// messages are separated with \n characters
			messages := bytes.Split(buf[:cursor], separator)
			// if one complete message and seperator is found
			// then we should have len(messages) > 1, the
			// last entry being an incomplete message or empty array.
			if len(messages) > 1 {
				shiftLen := 0
				for i := 0; i < len(messages)-1; i++ {
					// log.Println(string(messages[i]))

					msgCopy := make([]byte, len(messages[i]))
					copy(msgCopy, messages[i])

					go d.handleMessage(msgCopy)
					shiftLen += len(messages[i]) + 1
				}
				copy(buf, buf[shiftLen:])
				cursor -= shiftLen
			}
		}
	}()

	for {
		tcpMessage := <-d.tcpMessages
		log.Debug("tcpMessage:", string(tcpMessage))
		_, err := conn.Write(tcpMessage)
		if err != nil {
			log.Fatal("conn write error:", err)
		}
	}
}

// handleMessage handles a message read
// from TCP connection
func (d *Daemon) handleMessage(message []byte) {

	var tcpMsg TCPMessage

	err := json.Unmarshal(message, &tcpMsg)
	if err != nil {
		log.Println("json unmarshal error:", err)
		return
	}

	log.Debugf("handleMessage: %#v \n", tcpMsg)

	switch tcpMsg.Cmd {
	}
}

// Utility functions
func splitRepoAndTag(repoTag string) (string, string) {

	repo := ""
	tag := ""

	repoAndTag := strings.Split(repoTag, ":")

	if len(repoAndTag) > 0 {
		repo = repoAndTag[0]
	}

	if len(repoAndTag) > 1 {
		tag = repoAndTag[1]
	}

	return repo, tag
}
<commit_msg>Add query handler<commit_after>package daemon

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	_ "github.com/go-sql-driver/mysql"
	"github.com/prism-river/killy/collect"
)

// TCPMessage defines what a message that can be
// sent or received to/from LUA scripts
type TCPMessage struct {
	Cmd  string   `json:"cmd,omitempty"`
	Args []string `json:"args,omitempty"`
	// Id is used to associate requests & responses
	ID   int         `json:"id,omitempty"`
	Data interface{} `json:"data,omitempty"`
}

// ContainerEvent is one kind of Data that can
// be transported by a TCPMessage in the Data field.
// It describes a Docker container event. (start, stop, destroy...)
type TidbEvent struct {
	Action string `json:"action,omitempty"`
	ID     string `json:"id,omitempty"`
	Name   string `json:"name,omitempty"`
	CPU    string `json:"cpu,omitempty"`
	RAM    string `json:"ram,omitempty"`
	sync.RWMutex
}

// Table represents a table in TiDB
type Table struct {
	Name    string     `json:"name"`
	Columns []string   `json:"columns"`
	Data    [][]string `json:"data"`
}

// Config for the daemon
type Config struct {
	Database struct {
		User     string `json:"user"`
		Password string `json:"password"`
		Address  string `json:"address"`
		Name     string `json:"name"`
	} `json:"database"`
	Interval int `json:"interval"` // in seconds
}

// Daemon maintains state when the dockercraft daemon is running
type Daemon struct {
	Client *collect.Collect
	// The configuration
	Config *Config

	// tcpMessages can be used to send bytes to the Lua
	// plugin from any go routine.
	tcpMessages chan []byte

	sync.Mutex
}

// NewDaemon returns a new instance of Daemon
func NewDaemon(address string) *Daemon {
	client := collect.NewCollect("http://" + address + "/api/v1/query?query=")
	return &Daemon{
		Client: client,
	}
}

// Init initializes a Daemon
func (d *Daemon) Init() error {
	var err error
	// load configuration
	d.Config = new(Config)
	var configFile *os.File
	configFile, err = os.Open("config.json")
	defer configFile.Close()
	if err != nil {
		log.Fatal(err.Error())
	}
	jsonParser := json.NewDecoder(configFile)
	jsonParser.Decode(d.Config)

	d.tcpMessages = make(chan []byte)

	return nil
}

// Serve exposes a TCP server on port 25566 to handle
// connections from the LUA scripts
func (d *Daemon) Serve() {

	tcpAddr, err := net.ResolveTCPAddr("tcp", ":25566")

	ln, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		log.Fatalln("listen tcp error:", err)
	}
	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Fatalln("tcp conn accept error:", err)
		}
		// no need to handle connection in a go routine
		// goproxy is used as support for one single Lua plugin.
		d.handleConn(conn)
	}
}

// StartMonitoringEvents listens for events from the
// Docker daemon and uses callback to transmit them
// to LUA scripts.
func (d *Daemon) StartMonitoringEvents() {
	log.Info("Monitoring Database Events")

	// mysql test
	go func() {
		db, err := sql.Open("mysql", fmt.Sprintf("%v:%v@tcp(%v)/%v", d.Config.Database.User, d.Config.Database.Password, d.Config.Database.Address, d.Config.Database.Name))
		if err != nil {
			log.Fatal(err)
		}
		defer db.Close()

		for {
			tables := make([]string, 0)
			rows, err := db.Query("SHOW TABLES")
			if err != nil {
				log.Fatal(err)
			}
			defer rows.Close()
			for rows.Next() {
				var name string
				err := rows.Scan(&name)
				if err != nil {
					log.Fatal(err)
				}
				tables = append(tables, name)
			}
			err = rows.Err()
			if err != nil {
				log.Fatal(err)
			}

			res := make([]Table, 0)
			for _, tableName := range tables {
				table, err := sqlQuery(db, "SELECT * FROM "+tableName)
				if err != nil {
					log.Fatal(err)
				}
				table.Name = tableName
				res = append(res, *table)
			}

			tcpMsg := TCPMessage{}
			tcpMsg.Cmd = "event"
			tcpMsg.Args = []string{"table"}
			tcpMsg.ID = 0
			tcpMsg.Data = &res

			data, err := json.Marshal(&tcpMsg)
			if err != nil {
				log.Println("statCallback error:", err)
				return
			}

			separator := []byte(string('\n'))

			d.tcpMessages <- append(data, separator...)

			time.Sleep(time.Duration(d.Config.Interval) * time.Second)
		}
	}()
}

// handleConn handles a TCP connection
// with a Dockercraft Lua plugin.
func (d *Daemon) handleConn(conn net.Conn) {

	go func() {
		separator := []byte(string('\n'))

		buf := make([]byte, 256)
		cursor := 0
		for {
			// resize buf if needed
			if len(buf)-cursor < 256 {
				buf = append(buf, make([]byte, 256-(len(buf)-cursor))...)
			}
			n, err := conn.Read(buf[cursor:])
			if err != nil && err != io.EOF {
				log.Fatalln("conn read error: ", err)
			}
			cursor += n

			// TODO(aduermael): check cNetwork plugin implementation
			// conn.Read doesn't seem to be blocking if there's nothing
			// to read. Maybe the broken pipe is due to an implementation
			// problem on cNetwork plugin side
			if cursor == 0 {
				<-time.After(500 * time.Millisecond)
				continue
			}
			// log.Println("TCP data read:", string(buf[:cursor]), "cursor:", cursor)

			// see if there's a complete json message in buf.
			// messages are separated with \n characters
			messages := bytes.Split(buf[:cursor], separator)
			// if one complete message and seperator is found
			// then we should have len(messages) > 1, the
			// last entry being an incomplete message or empty array.
			if len(messages) > 1 {
				shiftLen := 0
				for i := 0; i < len(messages)-1; i++ {
					// log.Println(string(messages[i]))

					msgCopy := make([]byte, len(messages[i]))
					copy(msgCopy, messages[i])

					go d.handleMessage(msgCopy)
					shiftLen += len(messages[i]) + 1
				}
				copy(buf, buf[shiftLen:])
				cursor -= shiftLen
			}
		}
	}()

	for {
		tcpMessage := <-d.tcpMessages
		log.Debug("tcpMessage:", string(tcpMessage))
		_, err := conn.Write(tcpMessage)
		if err != nil {
			log.Fatal("conn write error:", err)
		}
	}
}

// handleMessage handles a message read
// from TCP connection
func (d *Daemon) handleMessage(message []byte) {

	var tcpMsg TCPMessage

	err := json.Unmarshal(message, &tcpMsg)
	if err != nil {
		log.Println("json unmarshal error:", err)
		return
	}

	log.Debugf("handleMessage: %#v \n", tcpMsg)

	switch tcpMsg.Cmd {
	case "query":
		query := tcpMsg.Data.(string)
		db, err := sql.Open("mysql", fmt.Sprintf("%v:%v@tcp(%v)/%v", d.Config.Database.User, d.Config.Database.Password, d.Config.Database.Address, d.Config.Database.Name))
		var msg TCPMessage
		if err != nil {
			msg = TCPMessage{
				Cmd:  "event",
				Args: []string{"error"},
				ID:   0,
				Data: err.Error(),
			}
		} else {
			defer db.Close()
			res, err := sqlQuery(db, query)
			if err != nil {
				msg = TCPMessage{
					Cmd:  "event",
					Args: []string{"error"},
					ID:   0,
					Data: err.Error(),
				}
			} else {
				msg = TCPMessage{
					Cmd:  "event",
					Args: []string{"result"},
					ID:   0,
					Data: res,
				}
			}
		}
		data, err := json.Marshal(&msg)
		if err != nil {
			log.Println("query error:", err)
			return
		}
		separator := []byte(string('\n'))
		d.tcpMessages <- append(data, separator...)
	}
}

// help function
func sqlQuery(db *sql.DB, query string) (*Table, error) {
	rows, err := db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	columns, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	table := Table{
		Name:    "result",
		Columns: columns,
		Data:    make([][]string, 0),
	}
	for rows.Next() {
		fields := make([]string, len(columns))
		pointers := make([]interface{}, len(columns))
		for i := range fields {
			pointers[i] = &fields[i]
		}
		err := rows.Scan(pointers...)
		if err != nil {
			return nil, err
		}
		table.Data = append(table.Data, fields)
	}
	err = rows.Err()
	if err != nil {
		return nil, err
	}
	return &table, nil
}
<|endoftext|>


<commit_before>package envconfig

import (
	"bufio"
	"bytes"
	"io"
	"strings"
)

type sliceTokenizer struct {
	err      error
	r        *bufio.Reader
	buf      bytes.Buffer
	inBraces bool
}

var eof = rune(0)

func newSliceTokenizer(str string) *sliceTokenizer {
	return &sliceTokenizer{
		r: bufio.NewReader(strings.NewReader(str)),
	}
}

func (t *sliceTokenizer) scan() bool {
	for {
		if t.err == io.EOF && t.buf.Len() == 0 {
			return false
		}

		ch := t.readRune()
		if ch == eof {
			return true
		}

		if ch == '{' {
			t.inBraces = true
		}
		if ch == '}' {
			t.inBraces = false
		}

		if ch == ',' && !t.inBraces {
			return true
		}

		_, t.err = t.buf.WriteRune(ch)
		if t.err != nil {
			return false
		}
	}
}

func (t *sliceTokenizer) readRune() rune {
	ch, _, err := t.r.ReadRune()
	if err != nil {
		t.err = err
		return eof
	}

	return ch
}

func (t *sliceTokenizer) text() string {
	str := t.buf.String()
	t.buf.Reset()

	return str
}

func (t *sliceTokenizer) Err() error {
	if t.err == io.EOF {
		return nil
	}
	return t.err
}
<commit_msg>Ignore the WriteRune error here<commit_after>package envconfig

import (
	"bufio"
	"bytes"
	"io"
	"strings"
)

type sliceTokenizer struct {
	err      error
	r        *bufio.Reader
	buf      bytes.Buffer
	inBraces bool
}

var eof = rune(0)

func newSliceTokenizer(str string) *sliceTokenizer {
	return &sliceTokenizer{
		r: bufio.NewReader(strings.NewReader(str)),
	}
}

func (t *sliceTokenizer) scan() bool {
	for {
		if t.err == io.EOF && t.buf.Len() == 0 {
			return false
		}

		ch := t.readRune()
		if ch == eof {
			return true
		}

		if ch == '{' {
			t.inBraces = true
		}
		if ch == '}' {
			t.inBraces = false
		}

		if ch == ',' && !t.inBraces {
			return true
		}

		// NOTE(vincent): we ignore the WriteRune error here because there is NO WAY
		// for WriteRune to return an error.
		// Yep. Seriously. Look here http://golang.org/src/bytes/buffer.go?s=7661:7714#L227
		_, _ = t.buf.WriteRune(ch)
	}
}

func (t *sliceTokenizer) readRune() rune {
	ch, _, err := t.r.ReadRune()
	if err != nil {
		t.err = err
		return eof
	}

	return ch
}

func (t *sliceTokenizer) text() string {
	str := t.buf.String()
	t.buf.Reset()

	return str
}

func (t *sliceTokenizer) Err() error {
	if t.err == io.EOF {
		return nil
	}
	return t.err
}
<|endoftext|>