<commit_before>/*
Copyright 2014 The Kubernetes 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 runtime

import (
	"fmt"
	"reflect"

	"k8s.io/apimachinery/pkg/runtime/schema"
)

type notRegisteredErr struct {
	gvk    schema.GroupVersionKind
	target GroupVersioner
	t      reflect.Type
}

func NewNotRegisteredErrForKind(gvk schema.GroupVersionKind) error {
	return &notRegisteredErr{gvk: gvk}
}

func NewNotRegisteredErrForType(t reflect.Type) error {
	return &notRegisteredErr{t: t}
}

func NewNotRegisteredErrForTarget(t reflect.Type, target GroupVersioner) error {
	return &notRegisteredErr{t: t, target: target}
}

func (k *notRegisteredErr) Error() string {
	if k.t != nil && k.target != nil {
		return fmt.Sprintf("%v is not suitable for converting to %q", k.t, k.target)
	}
	if k.t != nil {
		return fmt.Sprintf("no kind is registered for the type %v", k.t)
	}
	if len(k.gvk.Kind) == 0 {
		return fmt.Sprintf("no version %q has been registered", k.gvk.GroupVersion())
	}
	if k.gvk.Version == APIVersionInternal {
		return fmt.Sprintf("no kind %q is registered for the internal version of group %q", k.gvk.Kind, k.gvk.Group)
	}

	return fmt.Sprintf("no kind %q is registered for version %q", k.gvk.Kind, k.gvk.GroupVersion())
}

// IsNotRegisteredError returns true if the error indicates the provided
// object or input data is not registered.
func IsNotRegisteredError(err error) bool {
	if err == nil {
		return false
	}
	_, ok := err.(*notRegisteredErr)
	return ok
}

type missingKindErr struct {
	data string
}

func NewMissingKindErr(data string) error {
	return &missingKindErr{data}
}

func (k *missingKindErr) Error() string {
	return fmt.Sprintf("Object 'Kind' is missing in '%s'", k.data)
}

// IsMissingKind returns true if the error indicates that the provided object
// is missing a 'Kind' field.
func IsMissingKind(err error) bool {
	if err == nil {
		return false
	}
	_, ok := err.(*missingKindErr)
	return ok
}

type missingVersionErr struct {
	data string
}

// IsMissingVersion returns true if the error indicates that the provided object
// is missing a 'Version' field.
func NewMissingVersionErr(data string) error {
	return &missingVersionErr{data}
}

func (k *missingVersionErr) Error() string {
	return fmt.Sprintf("Object 'apiVersion' is missing in '%s'", k.data)
}

func IsMissingVersion(err error) bool {
	if err == nil {
		return false
	}
	_, ok := err.(*missingVersionErr)
	return ok
}
<commit_msg>move IsMissingVersion comments<commit_after>/*
Copyright 2014 The Kubernetes 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 runtime

import (
	"fmt"
	"reflect"

	"k8s.io/apimachinery/pkg/runtime/schema"
)

type notRegisteredErr struct {
	gvk    schema.GroupVersionKind
	target GroupVersioner
	t      reflect.Type
}

func NewNotRegisteredErrForKind(gvk schema.GroupVersionKind) error {
	return &notRegisteredErr{gvk: gvk}
}

func NewNotRegisteredErrForType(t reflect.Type) error {
	return &notRegisteredErr{t: t}
}

func NewNotRegisteredErrForTarget(t reflect.Type, target GroupVersioner) error {
	return &notRegisteredErr{t: t, target: target}
}

func (k *notRegisteredErr) Error() string {
	if k.t != nil && k.target != nil {
		return fmt.Sprintf("%v is not suitable for converting to %q", k.t, k.target)
	}
	if k.t != nil {
		return fmt.Sprintf("no kind is registered for the type %v", k.t)
	}
	if len(k.gvk.Kind) == 0 {
		return fmt.Sprintf("no version %q has been registered", k.gvk.GroupVersion())
	}
	if k.gvk.Version == APIVersionInternal {
		return fmt.Sprintf("no kind %q is registered for the internal version of group %q", k.gvk.Kind, k.gvk.Group)
	}

	return fmt.Sprintf("no kind %q is registered for version %q", k.gvk.Kind, k.gvk.GroupVersion())
}

// IsNotRegisteredError returns true if the error indicates the provided
// object or input data is not registered.
func IsNotRegisteredError(err error) bool {
	if err == nil {
		return false
	}
	_, ok := err.(*notRegisteredErr)
	return ok
}

type missingKindErr struct {
	data string
}

func NewMissingKindErr(data string) error {
	return &missingKindErr{data}
}

func (k *missingKindErr) Error() string {
	return fmt.Sprintf("Object 'Kind' is missing in '%s'", k.data)
}

// IsMissingKind returns true if the error indicates that the provided object
// is missing a 'Kind' field.
func IsMissingKind(err error) bool {
	if err == nil {
		return false
	}
	_, ok := err.(*missingKindErr)
	return ok
}

type missingVersionErr struct {
	data string
}

func NewMissingVersionErr(data string) error {
	return &missingVersionErr{data}
}

func (k *missingVersionErr) Error() string {
	return fmt.Sprintf("Object 'apiVersion' is missing in '%s'", k.data)
}

// IsMissingVersion returns true if the error indicates that the provided object
// is missing a 'Version' field.
func IsMissingVersion(err error) bool {
	if err == nil {
		return false
	}
	_, ok := err.(*missingVersionErr)
	return ok
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package vm

import (
	"fmt"

	"github.com/spf13/cobra"
	k8smetav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/tools/clientcmd"

	"k8s.io/apimachinery/pkg/types"

	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/virtctl/templates"
)

const (
	COMMAND_START = "start"
	COMMAND_STOP  = "stop"
)

func NewStartCommand(clientConfig clientcmd.ClientConfig) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "start (vmi)",
		Short:   "Start a virtual machine which is managed by an offline virtual machine.",
		Example: usage(COMMAND_START),
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			c := Command{command: COMMAND_START, clientConfig: clientConfig}
			return c.Run(cmd, args)
		},
	}
	cmd.SetUsageTemplate(templates.UsageTemplate())
	return cmd
}

func NewStopCommand(clientConfig clientcmd.ClientConfig) *cobra.Command {
	return &cobra.Command{
		Use:     "stop (vmi)",
		Short:   "Stop a virtual machine which is managed by an offline virtual machine.",
		Example: usage(COMMAND_STOP),
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			c := Command{command: COMMAND_STOP, clientConfig: clientConfig}
			return c.Run(cmd, args)
		},
	}
}

type Command struct {
	clientConfig clientcmd.ClientConfig
	command      string
}

func NewCommand(command string) *Command {
	return &Command{command: command}
}

func usage(cmd string) string {
	usage := "#Start a virtual machine called 'myvmi':\n"
	usage += fmt.Sprintf("virtctl %s myvmi", cmd)
	return usage
}

func (o *Command) Run(cmd *cobra.Command, args []string) error {

	vmiName := args[0]

	namespace, _, err := o.clientConfig.Namespace()
	if err != nil {
		return err
	}

	virtClient, err := kubecli.GetKubevirtClientFromClientConfig(o.clientConfig)
	if err != nil {
		return fmt.Errorf("Cannot obtain KubeVirt client: %v", err)
	}

	options := k8smetav1.GetOptions{}
	vm, err := virtClient.VirtualMachine(namespace).Get(vmiName, &options)
	if err != nil {
		return fmt.Errorf("Error fetching VirtualMachine: %v", err)
	}

	var running bool
	if o.command == COMMAND_START {
		running = true
	} else if o.command == COMMAND_STOP {
		running = false
	}

	if vm.Spec.Running != running {
		bodyStr := fmt.Sprintf("{\"spec\":{\"running\":%t}}", running)

		_, err := virtClient.VirtualMachine(namespace).Patch(vm.Name, types.MergePatchType,
			[]byte(bodyStr))

		if err != nil {
			return fmt.Errorf("Error updating VirtualMachine: %v", err)
		}

	} else {
		stateMsg := "stopped"
		if running {
			stateMsg = "running"
		}
		return fmt.Errorf("Error: VirtualMachineInstance '%s' is already %s", vmiName, stateMsg)
	}

	cmd.Printf("VM %s was scheduled to %s\n", vmiName, o.command)
	return nil
}
<commit_msg>Fix virtctl command references to ovms<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package vm

import (
	"fmt"

	"github.com/spf13/cobra"
	k8smetav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/tools/clientcmd"

	"k8s.io/apimachinery/pkg/types"

	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/virtctl/templates"
)

const (
	COMMAND_START = "start"
	COMMAND_STOP  = "stop"
)

func NewStartCommand(clientConfig clientcmd.ClientConfig) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "start (vm)",
		Short:   "Start a VirtualMachine.",
		Example: usage(COMMAND_START),
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			c := Command{command: COMMAND_START, clientConfig: clientConfig}
			return c.Run(cmd, args)
		},
	}
	cmd.SetUsageTemplate(templates.UsageTemplate())
	return cmd
}

func NewStopCommand(clientConfig clientcmd.ClientConfig) *cobra.Command {
	return &cobra.Command{
		Use:     "stop (vm)",
		Short:   "Stop a VirtualMachine.",
		Example: usage(COMMAND_STOP),
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			c := Command{command: COMMAND_STOP, clientConfig: clientConfig}
			return c.Run(cmd, args)
		},
	}
}

type Command struct {
	clientConfig clientcmd.ClientConfig
	command      string
}

func NewCommand(command string) *Command {
	return &Command{command: command}
}

func usage(cmd string) string {
	usage := "#Start a virtual machine called 'myvmi':\n"
	usage += fmt.Sprintf("virtctl %s myvmi", cmd)
	return usage
}

func (o *Command) Run(cmd *cobra.Command, args []string) error {

	vmiName := args[0]

	namespace, _, err := o.clientConfig.Namespace()
	if err != nil {
		return err
	}

	virtClient, err := kubecli.GetKubevirtClientFromClientConfig(o.clientConfig)
	if err != nil {
		return fmt.Errorf("Cannot obtain KubeVirt client: %v", err)
	}

	options := k8smetav1.GetOptions{}
	vm, err := virtClient.VirtualMachine(namespace).Get(vmiName, &options)
	if err != nil {
		return fmt.Errorf("Error fetching VirtualMachine: %v", err)
	}

	var running bool
	if o.command == COMMAND_START {
		running = true
	} else if o.command == COMMAND_STOP {
		running = false
	}

	if vm.Spec.Running != running {
		bodyStr := fmt.Sprintf("{\"spec\":{\"running\":%t}}", running)

		_, err := virtClient.VirtualMachine(namespace).Patch(vm.Name, types.MergePatchType,
			[]byte(bodyStr))

		if err != nil {
			return fmt.Errorf("Error updating VirtualMachine: %v", err)
		}

	} else {
		stateMsg := "stopped"
		if running {
			stateMsg = "running"
		}
		return fmt.Errorf("Error: VirtualMachineInstance '%s' is already %s", vmiName, stateMsg)
	}

	cmd.Printf("VM %s was scheduled to %s\n", vmiName, o.command)
	return nil
}
<|endoftext|>


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

// Package signer is a sample implementation of a commit-log based signer.
package signer

import (
	"flag"
	"fmt"
	"time"

	"github.com/golang/glog"
	"github.com/google/trillian/docs/storage/commit_log/simelection"
	"github.com/google/trillian/docs/storage/commit_log/simkafka"
)

var batchSize = flag.Int("batch_size", 5, "Maximum leaves to sign in one run")
var pessimizeInterval = flag.Duration("signer_pessimize", 10*time.Millisecond, "Pause interval in signing to induce inter-signer problems")

// Signer is a simulated signer instance.
type Signer struct {
	Name      string
	election  *simelection.Election
	epoch     int64
	dbSTHInfo STHInfo
	db        FakeDatabase
}

// New creates a simulated signer that uses the provided election.
func New(name string, election *simelection.Election, epoch int64) *Signer {
	return &Signer{
		Name:     name,
		election: election,
		epoch:    epoch,
		dbSTHInfo: STHInfo{
			treeRevision: -1,
			sthOffset:    -1,
		},
	}
}

func (s *Signer) String() string {
	prefix := "  "
	if s.IsMaster() {
		prefix = "**"
	}
	return fmt.Sprintf("%s Signer %s up to STH{offset=%d, rev=%d} = %s\n", prefix, s.Name, s.dbSTHInfo.sthOffset, s.dbSTHInfo.treeRevision, s.dbSTHInfo.sth.String())
}

// LatestSTHInfo returns the most recent STHInfo known about by the signer
func (s *Signer) LatestSTHInfo() STHInfo {
	return s.dbSTHInfo
}

// StoreSTHInfo updates the STHInfo known about by the signer.
func (s *Signer) StoreSTHInfo(info STHInfo) {
	s.dbSTHInfo = info
}

// IsMaster indicates if this signer is master.
func (s *Signer) IsMaster() bool {
	return s.election.IsMaster(s.Name)
}

// Run performs a single signing run.
func (s *Signer) Run() {
	// Read from local DB to see what STH we know about locally.
	dbSTHInfo := s.LatestSTHInfo()
	glog.V(2).Infof("%s: our DB has data upto STH at %d", s.Name, dbSTHInfo.sthOffset)

	// Sanity check that the STH table has what we already know.
	if dbSTHInfo.sth.TreeSize > 0 {
		ourSTH := sthFromString(simkafka.Read("STHs/<treeID>", dbSTHInfo.sthOffset))
		if ourSTH == nil {
			glog.Errorf("%s: local DB has data ahead of STHs topic!!", s.Name)
			return
		}
		if ourSTH.Offset != dbSTHInfo.sthOffset {
			glog.Errorf("%s: local DB recorded offset %d but that has inconsistent STH %s!!", s.Name, dbSTHInfo.sthOffset, ourSTH)
			return
		}
		if ourSTH.TimeStamp != dbSTHInfo.sth.TimeStamp || ourSTH.TreeSize != dbSTHInfo.sth.TreeSize {
			glog.Errorf("%s: local DB has different data than STHs topic!!", s.Name)
			return
		}
		glog.V(2).Infof("%s: our DB at %v, matches STH at that offset", s.Name, dbSTHInfo.sthOffset)
	}

	// Look to see if anyone else has already stored data just ahead of our STH.  This will
	// normally be the next entry, but we need to ignore any entries that have inconsistent
	// offsets.
	nextOffset := dbSTHInfo.sthOffset
	var nextSTH *STH
	for {
		nextOffset++
		nextSTH = sthFromString(simkafka.Read("STHs/<treeID>", nextOffset))
		if nextSTH == nil {
			break
		}
		if nextSTH != nil && nextSTH.Offset != nextOffset {
			glog.V(2).Infof("%s: ignoring inconsistent STH %s at offset %d", s.Name, nextSTH.String(), nextOffset)
			continue
		}
		if nextSTH.TimeStamp < dbSTHInfo.sth.TimeStamp || nextSTH.TreeSize < dbSTHInfo.sth.TreeSize {
			glog.Errorf("%s: next STH %s has earlier timestamp than in local DB (%s)!!", s.Name, nextSTH.String(), dbSTHInfo.sth)
			return
		}
		break
	}

	time.Sleep(*pessimizeInterval)
	if nextSTH == nil {
		// We're up-to-date with the STHs topic (as of a moment ago) ...
		if !s.IsMaster() {
			glog.V(2).Infof("%s: up-to-date with STHs but not master, so exit", s.Name)
			return
		}
		// ... and we're the master. Move the STHs topic along to encompass any unincorporated leaves.
		offset := dbSTHInfo.sth.TreeSize
		batch := simkafka.ReadMultiple("Leaves/<treeID>", offset, *batchSize)
		glog.V(2).Infof("%s: nothing at next offset %d and we are master, so read %d more leaves", s.Name, nextOffset, len(batch))
		if len(batch) == 0 {
			glog.V(2).Infof("%s: nothing to do", s.Name)
			return
		}
		timestamp := (time.Now().UnixNano() / int64(time.Millisecond)) - s.epoch
		newSTHInfo := STHInfo{
			sth: STH{
				TreeSize:  s.db.Size() + len(batch),
				TimeStamp: timestamp,
				Offset:    nextOffset, // The offset we expect this to end at
			},
			treeRevision: dbSTHInfo.treeRevision + 1,
		}
		newSTHInfo.sthOffset = simkafka.Append("STHs/<treeID>", newSTHInfo.sth.String())
		if newSTHInfo.sthOffset != nextOffset {
			// The STH didn't get stored at the offset we expected, presumable because someone else got there first
			glog.Warningf("%s: stored new STH %s at offset %d, which is unexpected; give up", s.Name, newSTHInfo.sth.String(), newSTHInfo.sthOffset)
			return
		}
		glog.V(2).Infof("%s: stored new STH %s at expected offset, including %d new leaves", s.Name, newSTHInfo.sth.String(), len(batch))

		// Now the STH topic is updated (correctly), do our local DB
		s.db.AddLeaves(timestamp, nextOffset, batch)
		s.StoreSTHInfo(newSTHInfo)
	} else {
		// There is an STH one ahead of us that we're not caught up with yet.
		// Read the leaves between what we have in our DB, and that STH...
		count := nextSTH.TreeSize - dbSTHInfo.sth.TreeSize
		glog.V(2).Infof("%s: our DB is %d leaves behind the next STH at %s, so update it", s.Name, count, nextSTH.String())
		batch := simkafka.ReadMultiple("Leaves/<treeID>", dbSTHInfo.sth.TreeSize, count)
		if len(batch) != count {
			glog.Errorf("%s: expected to read leaves [%d, %d) but only got %d!!", s.Name, dbSTHInfo.sth.TreeSize, dbSTHInfo.sth.TreeSize+count, len(batch))
			return
		}
		// ... and store it in our local DB
		newSTHInfo := STHInfo{
			sth:          s.db.AddLeaves(nextSTH.TimeStamp, nextOffset, batch),
			treeRevision: dbSTHInfo.treeRevision + 1,
			sthOffset:    nextOffset,
		}
		glog.V(2).Infof("%s: update our DB to %s", s.Name, newSTHInfo.sth.String())
		s.StoreSTHInfo(newSTHInfo)
		// We may still not be caught up, but that's for the next time around.
	}
}
<commit_msg>docs: drop superfluous if test<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package signer is a sample implementation of a commit-log based signer.
package signer

import (
	"flag"
	"fmt"
	"time"

	"github.com/golang/glog"
	"github.com/google/trillian/docs/storage/commit_log/simelection"
	"github.com/google/trillian/docs/storage/commit_log/simkafka"
)

var batchSize = flag.Int("batch_size", 5, "Maximum leaves to sign in one run")
var pessimizeInterval = flag.Duration("signer_pessimize", 10*time.Millisecond, "Pause interval in signing to induce inter-signer problems")

// Signer is a simulated signer instance.
type Signer struct {
	Name      string
	election  *simelection.Election
	epoch     int64
	dbSTHInfo STHInfo
	db        FakeDatabase
}

// New creates a simulated signer that uses the provided election.
func New(name string, election *simelection.Election, epoch int64) *Signer {
	return &Signer{
		Name:     name,
		election: election,
		epoch:    epoch,
		dbSTHInfo: STHInfo{
			treeRevision: -1,
			sthOffset:    -1,
		},
	}
}

func (s *Signer) String() string {
	prefix := "  "
	if s.IsMaster() {
		prefix = "**"
	}
	return fmt.Sprintf("%s Signer %s up to STH{offset=%d, rev=%d} = %s\n", prefix, s.Name, s.dbSTHInfo.sthOffset, s.dbSTHInfo.treeRevision, s.dbSTHInfo.sth.String())
}

// LatestSTHInfo returns the most recent STHInfo known about by the signer
func (s *Signer) LatestSTHInfo() STHInfo {
	return s.dbSTHInfo
}

// StoreSTHInfo updates the STHInfo known about by the signer.
func (s *Signer) StoreSTHInfo(info STHInfo) {
	s.dbSTHInfo = info
}

// IsMaster indicates if this signer is master.
func (s *Signer) IsMaster() bool {
	return s.election.IsMaster(s.Name)
}

// Run performs a single signing run.
func (s *Signer) Run() {
	// Read from local DB to see what STH we know about locally.
	dbSTHInfo := s.LatestSTHInfo()
	glog.V(2).Infof("%s: our DB has data upto STH at %d", s.Name, dbSTHInfo.sthOffset)

	// Sanity check that the STH table has what we already know.
	if dbSTHInfo.sth.TreeSize > 0 {
		ourSTH := sthFromString(simkafka.Read("STHs/<treeID>", dbSTHInfo.sthOffset))
		if ourSTH == nil {
			glog.Errorf("%s: local DB has data ahead of STHs topic!!", s.Name)
			return
		}
		if ourSTH.Offset != dbSTHInfo.sthOffset {
			glog.Errorf("%s: local DB recorded offset %d but that has inconsistent STH %s!!", s.Name, dbSTHInfo.sthOffset, ourSTH)
			return
		}
		if ourSTH.TimeStamp != dbSTHInfo.sth.TimeStamp || ourSTH.TreeSize != dbSTHInfo.sth.TreeSize {
			glog.Errorf("%s: local DB has different data than STHs topic!!", s.Name)
			return
		}
		glog.V(2).Infof("%s: our DB at %v, matches STH at that offset", s.Name, dbSTHInfo.sthOffset)
	}

	// Look to see if anyone else has already stored data just ahead of our STH.  This will
	// normally be the next entry, but we need to ignore any entries that have inconsistent
	// offsets.
	nextOffset := dbSTHInfo.sthOffset
	var nextSTH *STH
	for {
		nextOffset++
		nextSTH = sthFromString(simkafka.Read("STHs/<treeID>", nextOffset))
		if nextSTH == nil {
			break
		}
		if nextSTH.Offset != nextOffset {
			glog.V(2).Infof("%s: ignoring inconsistent STH %s at offset %d", s.Name, nextSTH.String(), nextOffset)
			continue
		}
		if nextSTH.TimeStamp < dbSTHInfo.sth.TimeStamp || nextSTH.TreeSize < dbSTHInfo.sth.TreeSize {
			glog.Errorf("%s: next STH %s has earlier timestamp than in local DB (%s)!!", s.Name, nextSTH.String(), dbSTHInfo.sth)
			return
		}
		break
	}

	time.Sleep(*pessimizeInterval)
	if nextSTH == nil {
		// We're up-to-date with the STHs topic (as of a moment ago) ...
		if !s.IsMaster() {
			glog.V(2).Infof("%s: up-to-date with STHs but not master, so exit", s.Name)
			return
		}
		// ... and we're the master. Move the STHs topic along to encompass any unincorporated leaves.
		offset := dbSTHInfo.sth.TreeSize
		batch := simkafka.ReadMultiple("Leaves/<treeID>", offset, *batchSize)
		glog.V(2).Infof("%s: nothing at next offset %d and we are master, so read %d more leaves", s.Name, nextOffset, len(batch))
		if len(batch) == 0 {
			glog.V(2).Infof("%s: nothing to do", s.Name)
			return
		}
		timestamp := (time.Now().UnixNano() / int64(time.Millisecond)) - s.epoch
		newSTHInfo := STHInfo{
			sth: STH{
				TreeSize:  s.db.Size() + len(batch),
				TimeStamp: timestamp,
				Offset:    nextOffset, // The offset we expect this to end at
			},
			treeRevision: dbSTHInfo.treeRevision + 1,
		}
		newSTHInfo.sthOffset = simkafka.Append("STHs/<treeID>", newSTHInfo.sth.String())
		if newSTHInfo.sthOffset != nextOffset {
			// The STH didn't get stored at the offset we expected, presumable because someone else got there first
			glog.Warningf("%s: stored new STH %s at offset %d, which is unexpected; give up", s.Name, newSTHInfo.sth.String(), newSTHInfo.sthOffset)
			return
		}
		glog.V(2).Infof("%s: stored new STH %s at expected offset, including %d new leaves", s.Name, newSTHInfo.sth.String(), len(batch))

		// Now the STH topic is updated (correctly), do our local DB
		s.db.AddLeaves(timestamp, nextOffset, batch)
		s.StoreSTHInfo(newSTHInfo)
	} else {
		// There is an STH one ahead of us that we're not caught up with yet.
		// Read the leaves between what we have in our DB, and that STH...
		count := nextSTH.TreeSize - dbSTHInfo.sth.TreeSize
		glog.V(2).Infof("%s: our DB is %d leaves behind the next STH at %s, so update it", s.Name, count, nextSTH.String())
		batch := simkafka.ReadMultiple("Leaves/<treeID>", dbSTHInfo.sth.TreeSize, count)
		if len(batch) != count {
			glog.Errorf("%s: expected to read leaves [%d, %d) but only got %d!!", s.Name, dbSTHInfo.sth.TreeSize, dbSTHInfo.sth.TreeSize+count, len(batch))
			return
		}
		// ... and store it in our local DB
		newSTHInfo := STHInfo{
			sth:          s.db.AddLeaves(nextSTH.TimeStamp, nextOffset, batch),
			treeRevision: dbSTHInfo.treeRevision + 1,
			sthOffset:    nextOffset,
		}
		glog.V(2).Infof("%s: update our DB to %s", s.Name, newSTHInfo.sth.String())
		s.StoreSTHInfo(newSTHInfo)
		// We may still not be caught up, but that's for the next time around.
	}
}
<|endoftext|>


<commit_before>package depot

import (
	"os"

	"github.com/cloudfoundry-incubator/executor/api"
	"github.com/cloudfoundry-incubator/executor/registry"
	"github.com/cloudfoundry-incubator/executor/sequence"
	"github.com/cloudfoundry-incubator/executor/transformer"
	"github.com/cloudfoundry-incubator/garden/warden"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/ifrit"
)

type client struct {
	containerOwnerName    string
	containerMaxCPUShares uint64
	containerInodeLimit   uint64
	wardenClient          warden.Client
	registry              registry.Registry
	transformer           *transformer.Transformer
	logger                lager.Logger
}

func NewClient(
	containerOwnerName string,
	containerMaxCPUShares uint64,
	containerInodeLimit uint64,
	wardenClient warden.Client,
	registry registry.Registry,
	transformer *transformer.Transformer,
	logger lager.Logger,
) api.Client {
	return &client{
		containerOwnerName:    containerOwnerName,
		containerMaxCPUShares: containerMaxCPUShares,
		containerInodeLimit:   containerInodeLimit,
		wardenClient:          wardenClient,
		registry:              registry,
		transformer:           transformer,
		logger:                logger.Session("depot-client"),
	}
}

func (c *client) InitializeContainer(guid string, request api.ContainerInitializationRequest) (api.Container, error) {
	if request.CpuPercent > 100 || request.CpuPercent < 0 {
		return api.Container{}, api.ErrLimitsInvalid
	}

	initLog := c.logger.Session("initialize", lager.Data{
		"guid": guid,
	})

	container, err := c.registry.FindByGuid(guid)
	if err != nil {
		initLog.Error("failed-to-find-container", err)
		return api.Container{}, api.ErrContainerNotFound
	}

	container, err = c.registry.Initialize(guid)
	if err != nil {
		initLog.Error("failed-to-initialize-registry-container", err)
		return api.Container{}, err
	}

	containerClient, err := c.wardenClient.Create(warden.ContainerSpec{
		RootFSPath: request.RootFSPath,
		Properties: warden.Properties{
			"owner": c.containerOwnerName,
		},
	})
	if err != nil {
		initLog.Error("failed-to-create-container", err)
		return api.Container{}, err
	}

	defer func() {
		if err != nil {
			initLog.Error("destroying-container-after-failed-init", err)
			destroyErr := c.wardenClient.Destroy(containerClient.Handle())
			if destroyErr != nil {
				initLog.Error("destroying-container-after-failed-init-also-failed", destroyErr)
			}
		}
	}()

	err = c.limitContainerDiskAndMemory(container, containerClient)
	if err != nil {
		initLog.Error("failed-to-limit-memory-and-disk", err)
		return api.Container{}, err
	}

	err = c.limitContainerCPU(request, containerClient)
	if err != nil {
		initLog.Error("failed-to-limit-cpu", err)
		return api.Container{}, err
	}

	portMapping, err := c.mapPorts(request, containerClient)
	if err != nil {
		initLog.Error("failed-to-map-ports", err)
		return api.Container{}, err
	}

	request.Ports = portMapping

	container, err = c.registry.Create(guid, containerClient.Handle(), request)
	if err != nil {
		initLog.Error("failed-to-register-container", err)
		return api.Container{}, err
	}

	return container, nil
}

func (c *client) limitContainerDiskAndMemory(reg api.Container, containerClient warden.Container) error {
	if reg.MemoryMB != 0 {
		err := containerClient.LimitMemory(warden.MemoryLimits{
			LimitInBytes: uint64(reg.MemoryMB * 1024 * 1024),
		})
		if err != nil {
			return err
		}
	}

	err := containerClient.LimitDisk(warden.DiskLimits{
		ByteHard:  uint64(reg.DiskMB * 1024 * 1024),
		InodeHard: c.containerInodeLimit,
	})
	if err != nil {
		return err
	}

	return nil
}

func (c *client) limitContainerCPU(request api.ContainerInitializationRequest, containerClient warden.Container) error {
	if request.CpuPercent != 0 {
		err := containerClient.LimitCPU(warden.CPULimits{
			LimitInShares: uint64(float64(c.containerMaxCPUShares) * float64(request.CpuPercent) / 100.0),
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *client) mapPorts(request api.ContainerInitializationRequest, containerClient warden.Container) ([]api.PortMapping, error) {
	var result []api.PortMapping
	for _, mapping := range request.Ports {
		hostPort, containerPort, err := containerClient.NetIn(mapping.HostPort, mapping.ContainerPort)
		if err != nil {
			return nil, err
		}

		result = append(result, api.PortMapping{
			HostPort:      hostPort,
			ContainerPort: containerPort,
		})
	}

	return result, nil
}

func (c *client) AllocateContainer(guid string, request api.ContainerAllocationRequest) (api.Container, error) {
	allocLog := c.logger.Session("allocate", lager.Data{
		"guid": guid,
	})

	container, err := c.registry.Reserve(guid, request)
	if err == registry.ErrContainerAlreadyExists {
		allocLog.Error("container-already-allocated", err)
		return api.Container{}, api.ErrContainerGuidNotAvailable
	}

	if err != nil {
		allocLog.Error("full", err)
		return api.Container{}, api.ErrInsufficientResourcesAvailable
	}

	return container, nil
}

func (c *client) GetContainer(guid string) (api.Container, error) {
	getLog := c.logger.Session("get", lager.Data{
		"guid": guid,
	})

	container, err := c.registry.FindByGuid(guid)
	if err != nil {
		getLog.Error("container-not-found", err)
		return api.Container{}, api.ErrContainerNotFound
	}
	return container, nil
}

func (c *client) Run(guid string, request api.ContainerRunRequest) error {
	runLog := c.logger.Session("run", lager.Data{
		"guid": guid,
	})

	registration, err := c.registry.FindByGuid(guid)
	if err != nil {
		runLog.Error("container-not-found", err)
		return api.ErrContainerNotFound
	}

	container, err := c.wardenClient.Lookup(registration.ContainerHandle)
	if err != nil {
		runLog.Error("lookup-failed", err)
		return err
	}

	var result string
	steps, err := c.transformer.StepsFor(registration.Log, request.Actions, request.Env, container, &result)
	if err != nil {
		runLog.Error("steps-invalid", err)
		return api.ErrStepsInvalid
	}

	run := RunSequence{
		CompleteURL:  request.CompleteURL,
		Registration: registration,
		Sequence:     sequence.New(steps),
		Result:       &result,
		Registry:     c.registry,
		Logger:       c.logger,
	}
	process := ifrit.Envoke(run)
	c.registry.Start(run.Registration.Guid, process)

	runLog.Info("started", lager.Data{
		"handle": registration.ContainerHandle,
	})

	return nil
}

func (c *client) ListContainers() ([]api.Container, error) {
	return c.registry.GetAllContainers(), nil
}

func (c *client) DeleteContainer(guid string) error {
	logData := lager.Data{
		"guid": guid,
	}

	deleteLog := c.logger.Session("delete", logData)

	reg, err := c.registry.MarkForDelete(guid)
	if err != nil {
		return handleDeleteError(err, deleteLog, logData)
	}

	logData["handle"] = reg.ContainerHandle

	deleteLog.Debug("deleting", logData)

	if reg.Process != nil {
		deleteLog.Debug("interrupting", logData)

		reg.Process.Signal(os.Interrupt)
		<-reg.Process.Wait()

		deleteLog.Info("interrupted", logData)
	}

	if reg.ContainerHandle != "" {
		deleteLog.Debug("destroying", logData)

		err = c.wardenClient.Destroy(reg.ContainerHandle)
		if err != nil {
			return handleDeleteError(err, deleteLog, logData)
		}

		deleteLog.Info("destroyed", logData)
	}

	deleteLog.Debug("unregistering", logData)

	err = c.registry.Delete(guid)
	if err != nil {
		return handleDeleteError(err, deleteLog, logData)
	}

	deleteLog.Info("unregistered", logData)

	return nil
}

func (c *client) RemainingResources() (api.ExecutorResources, error) {
	cap := c.registry.CurrentCapacity()

	return api.ExecutorResources{
		MemoryMB:   cap.MemoryMB,
		DiskMB:     cap.DiskMB,
		Containers: cap.Containers,
	}, nil
}

func (c *client) Ping() error {
	return c.wardenClient.Ping()
}

func (c *client) TotalResources() (api.ExecutorResources, error) {
	totalCapacity := c.registry.TotalCapacity()

	resources := api.ExecutorResources{
		MemoryMB:   totalCapacity.MemoryMB,
		DiskMB:     totalCapacity.DiskMB,
		Containers: totalCapacity.Containers,
	}
	return resources, nil
}

func handleDeleteError(err error, logger lager.Logger, logData lager.Data) error {
	if err == registry.ErrContainerNotFound {
		logger.Error("container-not-found", err, logData)
		return api.ErrContainerNotFound
	}

	logger.Error("failed-to-delete-container", err, logData)

	return err
}
<commit_msg>add container-handle to logging<commit_after>package depot

import (
	"os"

	"github.com/cloudfoundry-incubator/executor/api"
	"github.com/cloudfoundry-incubator/executor/registry"
	"github.com/cloudfoundry-incubator/executor/sequence"
	"github.com/cloudfoundry-incubator/executor/transformer"
	"github.com/cloudfoundry-incubator/garden/warden"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/ifrit"
)

type client struct {
	containerOwnerName    string
	containerMaxCPUShares uint64
	containerInodeLimit   uint64
	wardenClient          warden.Client
	registry              registry.Registry
	transformer           *transformer.Transformer
	logger                lager.Logger
}

func NewClient(
	containerOwnerName string,
	containerMaxCPUShares uint64,
	containerInodeLimit uint64,
	wardenClient warden.Client,
	registry registry.Registry,
	transformer *transformer.Transformer,
	logger lager.Logger,
) api.Client {
	return &client{
		containerOwnerName:    containerOwnerName,
		containerMaxCPUShares: containerMaxCPUShares,
		containerInodeLimit:   containerInodeLimit,
		wardenClient:          wardenClient,
		registry:              registry,
		transformer:           transformer,
		logger:                logger.Session("depot-client"),
	}
}

func (c *client) InitializeContainer(guid string, request api.ContainerInitializationRequest) (api.Container, error) {
	if request.CpuPercent > 100 || request.CpuPercent < 0 {
		return api.Container{}, api.ErrLimitsInvalid
	}

	initLog := c.logger.Session("initialize", lager.Data{
		"guid": guid,
	})

	container, err := c.registry.FindByGuid(guid)
	if err != nil {
		initLog.Error("failed-to-find-container", err)
		return api.Container{}, api.ErrContainerNotFound
	}

	container, err = c.registry.Initialize(guid)
	if err != nil {
		initLog.Error("failed-to-initialize-registry-container", err)
		return api.Container{}, err
	}

	containerClient, err := c.wardenClient.Create(warden.ContainerSpec{
		RootFSPath: request.RootFSPath,
		Properties: warden.Properties{
			"owner": c.containerOwnerName,
		},
	})
	if err != nil {
		initLog.Error("failed-to-create-container", err)
		return api.Container{}, err
	}

	defer func() {
		if err != nil {
			initLog.Error("destroying-container-after-failed-init", err)
			destroyErr := c.wardenClient.Destroy(containerClient.Handle())
			if destroyErr != nil {
				initLog.Error("destroying-container-after-failed-init-also-failed", destroyErr)
			}
		}
	}()

	err = c.limitContainerDiskAndMemory(container, containerClient)
	if err != nil {
		initLog.Error("failed-to-limit-memory-and-disk", err)
		return api.Container{}, err
	}

	err = c.limitContainerCPU(request, containerClient)
	if err != nil {
		initLog.Error("failed-to-limit-cpu", err)
		return api.Container{}, err
	}

	portMapping, err := c.mapPorts(request, containerClient)
	if err != nil {
		initLog.Error("failed-to-map-ports", err)
		return api.Container{}, err
	}

	request.Ports = portMapping

	container, err = c.registry.Create(guid, containerClient.Handle(), request)
	if err != nil {
		initLog.Error("failed-to-register-container", err, lager.Data{
			"container-handle": container.ContainerHandle,
		})
		return api.Container{}, err
	}

	return container, nil
}

func (c *client) limitContainerDiskAndMemory(reg api.Container, containerClient warden.Container) error {
	if reg.MemoryMB != 0 {
		err := containerClient.LimitMemory(warden.MemoryLimits{
			LimitInBytes: uint64(reg.MemoryMB * 1024 * 1024),
		})
		if err != nil {
			return err
		}
	}

	err := containerClient.LimitDisk(warden.DiskLimits{
		ByteHard:  uint64(reg.DiskMB * 1024 * 1024),
		InodeHard: c.containerInodeLimit,
	})
	if err != nil {
		return err
	}

	return nil
}

func (c *client) limitContainerCPU(request api.ContainerInitializationRequest, containerClient warden.Container) error {
	if request.CpuPercent != 0 {
		err := containerClient.LimitCPU(warden.CPULimits{
			LimitInShares: uint64(float64(c.containerMaxCPUShares) * float64(request.CpuPercent) / 100.0),
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *client) mapPorts(request api.ContainerInitializationRequest, containerClient warden.Container) ([]api.PortMapping, error) {
	var result []api.PortMapping
	for _, mapping := range request.Ports {
		hostPort, containerPort, err := containerClient.NetIn(mapping.HostPort, mapping.ContainerPort)
		if err != nil {
			return nil, err
		}

		result = append(result, api.PortMapping{
			HostPort:      hostPort,
			ContainerPort: containerPort,
		})
	}

	return result, nil
}

func (c *client) AllocateContainer(guid string, request api.ContainerAllocationRequest) (api.Container, error) {
	allocLog := c.logger.Session("allocate", lager.Data{
		"guid": guid,
	})

	container, err := c.registry.Reserve(guid, request)
	if err == registry.ErrContainerAlreadyExists {
		allocLog.Error("container-already-allocated", err)
		return api.Container{}, api.ErrContainerGuidNotAvailable
	}

	if err != nil {
		allocLog.Error("full", err)
		return api.Container{}, api.ErrInsufficientResourcesAvailable
	}

	return container, nil
}

func (c *client) GetContainer(guid string) (api.Container, error) {
	getLog := c.logger.Session("get", lager.Data{
		"guid": guid,
	})

	container, err := c.registry.FindByGuid(guid)
	if err != nil {
		getLog.Error("container-not-found", err)
		return api.Container{}, api.ErrContainerNotFound
	}
	return container, nil
}

func (c *client) Run(guid string, request api.ContainerRunRequest) error {
	logData := lager.Data{
		"guid": guid,
	}

	runLog := c.logger.Session("run", logData)

	registration, err := c.registry.FindByGuid(guid)
	if err != nil {
		runLog.Error("container-not-found", err)
		return api.ErrContainerNotFound
	}

	logData["container-handle"] = registration.ContainerHandle

	container, err := c.wardenClient.Lookup(registration.ContainerHandle)
	if err != nil {
		runLog.Error("lookup-failed", err, lager.Data{
			"container-handle": registration.ContainerHandle,
		})
		return err
	}

	var result string
	steps, err := c.transformer.StepsFor(registration.Log, request.Actions, request.Env, container, &result)
	if err != nil {
		runLog.Error("steps-invalid", err, logData)
		return api.ErrStepsInvalid
	}

	run := RunSequence{
		CompleteURL:  request.CompleteURL,
		Registration: registration,
		Sequence:     sequence.New(steps),
		Result:       &result,
		Registry:     c.registry,
		Logger:       c.logger,
	}
	process := ifrit.Envoke(run)
	c.registry.Start(run.Registration.Guid, process)

	runLog.Info("started", logData)

	return nil
}

func (c *client) ListContainers() ([]api.Container, error) {
	return c.registry.GetAllContainers(), nil
}

func (c *client) DeleteContainer(guid string) error {
	logData := lager.Data{
		"guid": guid,
	}

	deleteLog := c.logger.Session("delete", logData)

	reg, err := c.registry.MarkForDelete(guid)
	if err != nil {
		return handleDeleteError(err, deleteLog, logData)
	}

	logData["handle"] = reg.ContainerHandle

	deleteLog.Debug("deleting", logData)

	if reg.Process != nil {
		deleteLog.Debug("interrupting", logData)

		reg.Process.Signal(os.Interrupt)
		<-reg.Process.Wait()

		deleteLog.Info("interrupted", logData)
	}

	if reg.ContainerHandle != "" {
		deleteLog.Debug("destroying", logData)

		err = c.wardenClient.Destroy(reg.ContainerHandle)
		if err != nil {
			return handleDeleteError(err, deleteLog, logData)
		}

		deleteLog.Info("destroyed", logData)
	}

	deleteLog.Debug("unregistering", logData)

	err = c.registry.Delete(guid)
	if err != nil {
		return handleDeleteError(err, deleteLog, logData)
	}

	deleteLog.Info("unregistered", logData)

	return nil
}

func (c *client) RemainingResources() (api.ExecutorResources, error) {
	cap := c.registry.CurrentCapacity()

	return api.ExecutorResources{
		MemoryMB:   cap.MemoryMB,
		DiskMB:     cap.DiskMB,
		Containers: cap.Containers,
	}, nil
}

func (c *client) Ping() error {
	return c.wardenClient.Ping()
}

func (c *client) TotalResources() (api.ExecutorResources, error) {
	totalCapacity := c.registry.TotalCapacity()

	resources := api.ExecutorResources{
		MemoryMB:   totalCapacity.MemoryMB,
		DiskMB:     totalCapacity.DiskMB,
		Containers: totalCapacity.Containers,
	}
	return resources, nil
}

func handleDeleteError(err error, logger lager.Logger, logData lager.Data) error {
	if err == registry.ErrContainerNotFound {
		logger.Error("container-not-found", err, logData)
		return api.ErrContainerNotFound
	}

	logger.Error("failed-to-delete-container", err, logData)

	return err
}
<|endoftext|>


<commit_before>package margelet

import (
	"fmt"
	"github.com/Syfaro/telegram-bot-api"
)

// PanicResponder - test responder that panics
type PanicResponder struct {
	Margelet *Margelet
}

// Response sends default help message
func (responder PanicResponder) Response(bot MargeletAPI, message tgbotapi.Message) error {
	panic("TEST")
	return nil
}

// HelpMessage return help string for HelpResponder
func (responder PanicResponder) HelpMessage() string {
	return "Panic!"
}
<commit_msg>build fixed<commit_after>package margelet

import (
	"github.com/Syfaro/telegram-bot-api"
)

// PanicResponder - test responder that panics
type PanicResponder struct {
	Margelet *Margelet
}

// Response sends default help message
func (responder PanicResponder) Response(bot MargeletAPI, message tgbotapi.Message) error {
	panic("TEST")
	return nil
}

// HelpMessage return help string for HelpResponder
func (responder PanicResponder) HelpMessage() string {
	return "Panic!"
}
<|endoftext|>


<commit_before>package checks

import (
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
)

func runCommand(cmd string) (string, string, error) {
	file, err := ioutil.TempFile(os.TempDir(), "hakim")
	if err != nil {
		return "", "", err
	}
	defer os.Remove(file.Name())

	_, err = file.Write([]byte(cmd))
	if err != nil {
		return "", "", err
	}

	err = file.Close()
	if err != nil {
		return "", "", err
	}

	err = os.Rename(file.Name(), file.Name()+".ps1")
	if err != nil {
		return "", "", err
	}

	execCmd := exec.Command("powershell", "-noprofile", "-noninteractive", file.Name()+".ps1")

	execCmd.Stdin = bytes.NewBufferString(cmd)

	stdout := bytes.NewBufferString("")
	execCmd.Stdout = stdout

	stderr := bytes.NewBufferString("")
	execCmd.Stderr = stderr

	err = execCmd.Run()

	stdoutStr := string(stdout.Bytes())
	stderrStr := string(stderr.Bytes())

	execCmd.Wait()

	return stdoutStr, stderrStr, err
}
<commit_msg>Fix runCommand execution<commit_after>package checks

import (
	"bytes"
	"os/exec"
)

func runCommand(cmd string) (string, string, error) {
	execCmd := exec.Command("powershell", "-noprofile", "-noninteractive", "-command", cmd)

	var stdout bytes.Buffer
	execCmd.Stdout = &stdout

	var stderr bytes.Buffer
	execCmd.Stderr = &stderr

	err := execCmd.Run()
	return stdout.String(), stderr.String(), err
}
<|endoftext|>


<commit_before>package workstation

import (
	"fmt"
	"io"
	"os/exec"

	"math"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"encoding/csv"
	"github.com/xshellinc/tools/dialogs"
	"github.com/xshellinc/tools/lib/help"
)

var cleanTemplate = `
select disk %s
clean
create partition primary
active
assign letter=N
remove letter=N
format fs=fat32 label=New quick
`

type windows struct {
	*workstation
	ddPath string
}

func newWorkstation() WorkStation {
	m := new(MountInfo)
	var ms []*MountInfo
	return &windows{&workstation{runtime.GOOS, true, m, ms}, ""}
}

// Lists available mounts
func (w *windows) ListRemovableDisk() error {
	log.Debug("Listing disks...")
	var out = []*MountInfo{}

	stdout, err := help.ExecCmd("wmic", []string{"diskdrive", "get", "DeviceID,index,InterfaceType,MediaType,Model,Size", "/format:csv"})
	log.Debug(stdout)
	if err != nil {
		stdout = ""
	}
	r := csv.NewReader(strings.NewReader(strings.TrimSpace(stdout)))
	r.TrimLeadingSpace = true
	r.Read() //skip the first line
	for {
		if record, err := r.Read(); err == io.EOF {
			break
		} else if err == nil {
			if !strings.Contains(record[4], "Removable") || strings.Contains(record[3], "IDE") {
				continue
			}
			var p = &MountInfo{}
			size := record[6]
			p.deviceSize = size
			sizeInt, _ := strconv.Atoi(size)
			sizeFloat := math.Ceil(float64(sizeInt) / 1024 / 1024 / 1024)
			p.deviceName = record[5] + " [" + strconv.Itoa(int(sizeFloat)) + "GB]"
			p.diskName = `\\?\Device\Harddisk` + record[2] + `\Partition0`
			p.diskNameRaw = record[2]
			out = append(out, p)
		}
	}
	log.WithField("out", out).Debug("got drives")
	if !(len(out) > 0) {
		return fmt.Errorf("[-] No removable disks found, please insert your SD card and try again.\n[-] Please remember to run this tool as an administrator.")
	}
	w.workstation.mounts = out
	return nil
}

// Unmounts the disk
func (w *windows) Unmount() error {
	return nil
}

// Ejects the mounted disk
func (w *windows) Eject() error {
	return nil
}

const diskSelectionTries = 3
const writeAttempts = 5

// Notifies user to choose a mount, after that it tries to write the data with `diskSelectionTries` number of retries
func (w *windows) WriteToDisk(img string) (job *help.BackgroundJob, err error) {
	for attempt := 0; attempt < diskSelectionTries; attempt++ {
		if attempt > 0 && !dialogs.YesNoDialog("Continue?") {
			break
		}

		err = w.ListRemovableDisk()
		if err != nil {
			fmt.Println("[-] SD card not found, please insert an unlocked SD card")
			continue
		}

		rng := make([]string, len(w.workstation.mounts))
		for i, e := range w.workstation.mounts {
			rng[i] = fmt.Sprintf(dialogs.PrintColored("%s")+" - "+dialogs.PrintColored("%s"), e.deviceName, e.diskName)
		}
		num := dialogs.SelectOneDialog("Select disk to use: ", rng)

		w.workstation.mount = w.workstation.mounts[num]
		break
	}

	if err != nil {
		return nil, err
	}

	if w.ddPath == "" {
		if err := w.getDDBinary(); err != nil {
			log.Error(err)
			fmt.Println("[-] Error downloading dd binary")
			return nil, err
		}
	}

	if dialogs.YesNoDialog("Are you sure? ") {
		fmt.Printf("[+] Writing %s to %s\n", img, w.workstation.mount.deviceName)

		job = help.NewBackgroundJob()

		go func() {
			defer job.Close()

			var err error
			// may be save state here? So if user canceled operation or error happened he wouldn't need to start the process from the beginning
			for attempt := 0; attempt < writeAttempts; attempt++ {
				if attempt > 0 {
					if !dialogs.YesNoDialog("Retry flashing?") {
						break
					}
				}
				job.Active(true)
				var out []byte
				if out, err = exec.Command(w.ddPath,
					"--filter=removable",
					fmt.Sprintf("if=%s", img),
					fmt.Sprintf("of=%s", w.workstation.mount.diskName),
					"bs=1M").CombinedOutput(); err != nil {
					log.WithField("out", string(out)).Error("Error while executing: `", w.ddPath)
					job.Active(false)
					fmt.Println("\r[-] Can't write to disk.")
				} else {
					sout := string(out)
					job.Active(false)
					log.WithField("out", sout).Debug("dd finished")
					if strings.Contains(sout, "Error ") {
						if strings.Contains(sout, "Access is denied") || strings.Contains(sout, "The device is not ready") {
							fmt.Println("\n[-] Can't write to disk. Please make sure to run this tool as administrator, close all Explorer windows, try reconnecting your disk and finally reboot your computer.\n [-] You can run this tool with `clean` to clean your disk before applying image.")
							if dialogs.YesNoDialog("Or we can try to clean it's partitions right now, should we proceed?") {
								if derr := w.CleanDisk(); derr != nil {
									fmt.Println("[-] Disk cleaning failed:", derr)
									continue
								} else {
									for !dialogs.YesNoDialog("[+] Disk formatted, now please reconnect the device. Type yes once you've done it.") {
									}
								}
							}
							continue
						} else {
							fmt.Println(sout)
							continue
						}
					}
					fmt.Printf("\r[+] Done writing %s to %s \n", img, w.workstation.mount.diskName)
					return
				}
			}

			if err != nil {
				job.Error(err)
			}
			job.Error(fmt.Errorf("Image wasn't flashed"))
		}()
		return job, nil
	}
	return nil, nil
}

func (w *windows) getDDBinary() error {
	dst := help.GetTempDir() + help.Separator()
	url := "https://cdn.isaax.io/isaax-distro/utilities/dd/ddrelease64.zip"

	if help.Exists(dst + "ddrelease64.exe") {
		w.ddPath = dst + "ddrelease64.exe"
		return nil
	}

	wg := &sync.WaitGroup{}
	fileName, bar, err := help.DownloadFromUrlWithAttemptsAsync(url, dst, 5, wg)
	if err != nil {
		return err
	}

	bar.Prefix(fmt.Sprintf("[+] Download %-15s", fileName))
	bar.Start()
	wg.Wait()
	bar.Finish()
	time.Sleep(time.Second)

	log.WithField("dst", dst).Debug("Extracting")
	if out, err := exec.Command("unzip", "-o", dst+"ddrelease64.zip", "-d", dst).CombinedOutput(); err != nil {
		return err
	} else {
		log.Debug(string(out))
	}
	w.ddPath = dst + "ddrelease64.exe"
	return nil
}

// CleanDisk cleans target disk partitions
func (w *windows) CleanDisk() error {
	fmt.Println("[+] Cleaning disk...")
	var last error
	for attempt := 0; attempt < diskSelectionTries; attempt++ {
		if attempt > 0 && !dialogs.YesNoDialog("Continue?") {
			break
		}

		if err := w.ListRemovableDisk(); err != nil {
			fmt.Println("[-] SD card not found, please insert an unlocked SD card")
			last = err
			continue
		}

		rng := make([]string, len(w.workstation.mounts))
		for i, e := range w.workstation.mounts {
			rng[i] = fmt.Sprintf(dialogs.PrintColored("%s")+" - "+dialogs.PrintColored("%s"), e.deviceName, e.diskName)
		}
		num := dialogs.SelectOneDialog("Select disk to clean: ", rng)

		w.workstation.mount = w.workstation.mounts[num]
		break
	}

	if last != nil {
		return last
	}

	dst := help.GetTempDir() + help.Separator() + "clean_script.txt"
	if dialogs.YesNoDialog("Are you sure you want to clean this disk? ") {
		fmt.Printf("[+] Cleaning disk %s (%s)\n", w.workstation.mount.diskNameRaw, w.workstation.mount.deviceName)
		help.CreateFile(dst)
		help.WriteFile(dst, fmt.Sprintf(cleanTemplate, w.workstation.mount.diskNameRaw))

		if help.Exists(dst) {
			if out, err := exec.Command("diskpart", "/s", dst).CombinedOutput(); err != nil {
				return err
			} else {
				log.Debug(string(out))
				fmt.Println(string(out))
				return nil
			}
		}
	}
	return err
}
<commit_msg>Fix windows build error<commit_after>package workstation

import (
	"fmt"
	"io"
	"os/exec"

	"math"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"encoding/csv"
	"github.com/xshellinc/tools/dialogs"
	"github.com/xshellinc/tools/lib/help"
)

var cleanTemplate = `
select disk %s
clean
create partition primary
active
assign letter=N
remove letter=N
format fs=fat32 label=New quick
`

type windows struct {
	*workstation
	ddPath string
}

func newWorkstation() WorkStation {
	m := new(MountInfo)
	var ms []*MountInfo
	return &windows{&workstation{runtime.GOOS, true, m, ms}, ""}
}

// Lists available mounts
func (w *windows) ListRemovableDisk() error {
	log.Debug("Listing disks...")
	var out = []*MountInfo{}

	stdout, err := help.ExecCmd("wmic", []string{"diskdrive", "get", "DeviceID,index,InterfaceType,MediaType,Model,Size", "/format:csv"})
	log.Debug(stdout)
	if err != nil {
		stdout = ""
	}
	r := csv.NewReader(strings.NewReader(strings.TrimSpace(stdout)))
	r.TrimLeadingSpace = true
	r.Read() //skip the first line
	for {
		if record, err := r.Read(); err == io.EOF {
			break
		} else if err == nil {
			if !strings.Contains(record[4], "Removable") || strings.Contains(record[3], "IDE") {
				continue
			}
			var p = &MountInfo{}
			size := record[6]
			p.deviceSize = size
			sizeInt, _ := strconv.Atoi(size)
			sizeFloat := math.Ceil(float64(sizeInt) / 1024 / 1024 / 1024)
			p.deviceName = record[5] + " [" + strconv.Itoa(int(sizeFloat)) + "GB]"
			p.diskName = `\\?\Device\Harddisk` + record[2] + `\Partition0`
			p.diskNameRaw = record[2]
			out = append(out, p)
		}
	}
	log.WithField("out", out).Debug("got drives")
	if !(len(out) > 0) {
		return fmt.Errorf("[-] No removable disks found, please insert your SD card and try again.\n[-] Please remember to run this tool as an administrator.")
	}
	w.workstation.mounts = out
	return nil
}

// Unmounts the disk
func (w *windows) Unmount() error {
	return nil
}

// Ejects the mounted disk
func (w *windows) Eject() error {
	return nil
}

const diskSelectionTries = 3
const writeAttempts = 5

// Notifies user to choose a mount, after that it tries to write the data with `diskSelectionTries` number of retries
func (w *windows) WriteToDisk(img string) (job *help.BackgroundJob, err error) {
	for attempt := 0; attempt < diskSelectionTries; attempt++ {
		if attempt > 0 && !dialogs.YesNoDialog("Continue?") {
			break
		}

		err = w.ListRemovableDisk()
		if err != nil {
			fmt.Println("[-] SD card not found, please insert an unlocked SD card")
			continue
		}

		rng := make([]string, len(w.workstation.mounts))
		for i, e := range w.workstation.mounts {
			rng[i] = fmt.Sprintf(dialogs.PrintColored("%s")+" - "+dialogs.PrintColored("%s"), e.deviceName, e.diskName)
		}
		num := dialogs.SelectOneDialog("Select disk to use: ", rng)

		w.workstation.mount = w.workstation.mounts[num]
		break
	}

	if err != nil {
		return nil, err
	}

	if w.ddPath == "" {
		if err := w.getDDBinary(); err != nil {
			log.Error(err)
			fmt.Println("[-] Error downloading dd binary")
			return nil, err
		}
	}

	if dialogs.YesNoDialog("Are you sure? ") {
		fmt.Printf("[+] Writing %s to %s\n", img, w.workstation.mount.deviceName)

		job = help.NewBackgroundJob()

		go func() {
			defer job.Close()

			var err error
			// may be save state here? So if user canceled operation or error happened he wouldn't need to start the process from the beginning
			for attempt := 0; attempt < writeAttempts; attempt++ {
				if attempt > 0 {
					if !dialogs.YesNoDialog("Retry flashing?") {
						break
					}
				}
				job.Active(true)
				var out []byte
				if out, err = exec.Command(w.ddPath,
					"--filter=removable",
					fmt.Sprintf("if=%s", img),
					fmt.Sprintf("of=%s", w.workstation.mount.diskName),
					"bs=1M").CombinedOutput(); err != nil {
					log.WithField("out", string(out)).Error("Error while executing: `", w.ddPath)
					job.Active(false)
					fmt.Println("\r[-] Can't write to disk.")
				} else {
					sout := string(out)
					job.Active(false)
					log.WithField("out", sout).Debug("dd finished")
					if strings.Contains(sout, "Error ") {
						if strings.Contains(sout, "Access is denied") || strings.Contains(sout, "The device is not ready") {
							fmt.Println("\n[-] Can't write to disk. Please make sure to run this tool as administrator, close all Explorer windows, try reconnecting your disk and finally reboot your computer.\n [-] You can run this tool with `clean` to clean your disk before applying image.")
							if dialogs.YesNoDialog("Or we can try to clean it's partitions right now, should we proceed?") {
								if derr := w.CleanDisk(); derr != nil {
									fmt.Println("[-] Disk cleaning failed:", derr)
									continue
								} else {
									for !dialogs.YesNoDialog("[+] Disk formatted, now please reconnect the device. Type yes once you've done it.") {
									}
								}
							}
							continue
						} else {
							fmt.Println(sout)
							continue
						}
					}
					fmt.Printf("\r[+] Done writing %s to %s \n", img, w.workstation.mount.diskName)
					return
				}
			}

			if err != nil {
				job.Error(err)
			}
			job.Error(fmt.Errorf("Image wasn't flashed"))
		}()
		return job, nil
	}
	return nil, nil
}

func (w *windows) getDDBinary() error {
	dst := help.GetTempDir() + help.Separator()
	url := "https://cdn.isaax.io/isaax-distro/utilities/dd/ddrelease64.zip"

	if help.Exists(dst + "ddrelease64.exe") {
		w.ddPath = dst + "ddrelease64.exe"
		return nil
	}

	wg := &sync.WaitGroup{}
	fileName, bar, err := help.DownloadFromUrlWithAttemptsAsync(url, dst, 5, wg)
	if err != nil {
		return err
	}

	bar.Prefix(fmt.Sprintf("[+] Download %-15s", fileName))
	bar.Start()
	wg.Wait()
	bar.Finish()
	time.Sleep(time.Second)

	log.WithField("dst", dst).Debug("Extracting")
	if out, err := exec.Command("unzip", "-o", dst+"ddrelease64.zip", "-d", dst).CombinedOutput(); err != nil {
		return err
	} else {
		log.Debug(string(out))
	}
	w.ddPath = dst + "ddrelease64.exe"
	return nil
}

// CleanDisk cleans target disk partitions
func (w *windows) CleanDisk() error {
	fmt.Println("[+] Cleaning disk...")
	var last error
	for attempt := 0; attempt < diskSelectionTries; attempt++ {
		if attempt > 0 && !dialogs.YesNoDialog("Continue?") {
			break
		}

		if err := w.ListRemovableDisk(); err != nil {
			fmt.Println("[-] SD card not found, please insert an unlocked SD card")
			last = err
			continue
		}

		rng := make([]string, len(w.workstation.mounts))
		for i, e := range w.workstation.mounts {
			rng[i] = fmt.Sprintf(dialogs.PrintColored("%s")+" - "+dialogs.PrintColored("%s"), e.deviceName, e.diskName)
		}
		num := dialogs.SelectOneDialog("Select disk to clean: ", rng)

		w.workstation.mount = w.workstation.mounts[num]
		break
	}

	if last != nil {
		return last
	}

	dst := help.GetTempDir() + help.Separator() + "clean_script.txt"
	if dialogs.YesNoDialog("Are you sure you want to clean this disk? ") {
		fmt.Printf("[+] Cleaning disk %s (%s)\n", w.workstation.mount.diskNameRaw, w.workstation.mount.deviceName)
		help.CreateFile(dst)
		help.WriteFile(dst, fmt.Sprintf(cleanTemplate, w.workstation.mount.diskNameRaw))

		if help.Exists(dst) {
			if out, err := exec.Command("diskpart", "/s", dst).CombinedOutput(); err != nil {
				return err
			} else {
				log.Debug(string(out))
				fmt.Println(string(out))
				return nil
			}
		}
	}
	return nil
}
<|endoftext|>


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

package rootio_test

import (
	"testing"

	"go-hep.org/x/hep/rootio"
)

func TestChain(t *testing.T) {
	for _, tc := range []struct {
		fnames  []string
		entries int64
		name    string
		title   string
	}{
		{
			fnames:  nil,
			entries: 0,
			name:    "",
			title:   "",
		},
		{
			fnames:  []string{"testdata/chain.1.root"},
			entries: 10,
			name:    "tree",
			title:   "my tree title",
		},
		{
			// twice the same tree
			fnames:  []string{"testdata/chain.1.root", "testdata/chain.1.root"},
			entries: 20,
			name:    "tree",
			title:   "my tree title",
		},
		{
			// two different trees (with the same schema)
			fnames:  []string{"testdata/chain.1.root", "testdata/chain.2.root"},
			entries: 20,
			name:    "tree",
			title:   "my tree title",
		},
		// TODO(sbinet): add a test with 2 trees with different schemas)
	} {
		t.Run("", func(t *testing.T) {
			files := make([]*rootio.File, len(tc.fnames))
			trees := make([]rootio.Tree, len(tc.fnames))
			for i, fname := range tc.fnames {
				f, err := rootio.Open(fname)
				if err != nil {
					t.Fatalf("could not open ROOT file %q: %v", fname, err)
				}
				defer f.Close()
				files[i] = f

				obj, err := f.Get(tc.name)
				if err != nil {
					t.Fatal(err)
				}

				trees[i] = obj.(rootio.Tree)
			}

			chain := rootio.Chain(trees...)

			if got, want := chain.Name(), tc.name; got != want {
				t.Fatalf("names differ\ngot = %q, want= %q", got, want)
			}
			if got, want := chain.Title(), tc.title; got != want {
				t.Fatalf("titles differ\ngot = %q, want= %q", got, want)
			}
			if got, want := chain.Entries(), tc.entries; got != want {
				t.Fatalf("titles differ\ngot = %v, want= %v", got, want)
			}
		})
	}
}
<commit_msg>rootio: add Scanner test for Chain<commit_after>// Copyright 2018 The go-hep Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rootio_test

import (
	"fmt"
	"io"
	"reflect"
	"testing"

	"go-hep.org/x/hep/rootio"
)

func TestChain(t *testing.T) {
	for _, tc := range []struct {
		fnames  []string
		entries int64
		name    string
		title   string
	}{
		{
			fnames:  nil,
			entries: 0,
			name:    "",
			title:   "",
		},
		{
			fnames:  []string{"testdata/chain.1.root"},
			entries: 10,
			name:    "tree",
			title:   "my tree title",
		},
		{
			// twice the same tree
			fnames:  []string{"testdata/chain.1.root", "testdata/chain.1.root"},
			entries: 20,
			name:    "tree",
			title:   "my tree title",
		},
		{
			// two different trees (with the same schema)
			fnames:  []string{"testdata/chain.1.root", "testdata/chain.2.root"},
			entries: 20,
			name:    "tree",
			title:   "my tree title",
		},
		// TODO(sbinet): add a test with 2 trees with different schemas)
	} {
		t.Run("", func(t *testing.T) {
			files := make([]*rootio.File, len(tc.fnames))
			trees := make([]rootio.Tree, len(tc.fnames))
			for i, fname := range tc.fnames {
				f, err := rootio.Open(fname)
				if err != nil {
					t.Fatalf("could not open ROOT file %q: %v", fname, err)
				}
				defer f.Close()
				files[i] = f

				obj, err := f.Get(tc.name)
				if err != nil {
					t.Fatal(err)
				}

				trees[i] = obj.(rootio.Tree)
			}

			chain := rootio.Chain(trees...)

			if got, want := chain.Name(), tc.name; got != want {
				t.Fatalf("names differ\ngot = %q, want= %q", got, want)
			}
			if got, want := chain.Title(), tc.title; got != want {
				t.Fatalf("titles differ\ngot = %q, want= %q", got, want)
			}
			if got, want := chain.Entries(), tc.entries; got != want {
				t.Fatalf("titles differ\ngot = %v, want= %v", got, want)
			}
		})
	}
}

func TestChainScan(t *testing.T) {
	files := []string{
		"testdata/chain.1.root",
		//	"testdata/chain.2.root", // FIXME(sbinet): implement for >1 tree
	}
	trees := make([]rootio.Tree, len(files))
	for i, fname := range files {
		f, err := rootio.Open(fname)
		if err != nil {
			t.Fatalf("could not open ROOT file %q: %v", fname, err)
		}
		defer f.Close()

		obj, err := f.Get("tree")
		if err != nil {
			t.Fatal(err)
		}

		trees[i] = obj.(rootio.Tree)
	}

	chain := rootio.Chain(trees...)

	type Data struct {
		Event struct {
			Beg       string      `rootio:"Beg"`
			F64       float64     `rootio:"F64"`
			ArrF64    [10]float64 `rootio:"ArrayF64"`
			N         int32       `rootio:"N"`
			SliF64    []float64   `rootio:"SliceF64"`
			StdStr    string      `rootio:"StdStr"`
			StlVecF64 []float64   `rootio:"StlVecF64"`
			StlVecStr []string    `rootio:"StlVecStr"`
			End       string      `rootio:"End"`
		} `rootio:"evt"`
	}

	want := func(i int64) (data Data) {
		evt := &data.Event
		evt.Beg = fmt.Sprintf("beg-%03d", i)
		evt.F64 = float64(i)
		for j := range evt.ArrF64 {
			evt.ArrF64[j] = float64(i)
		}
		evt.N = int32(i) % 10
		evt.SliF64 = make([]float64, evt.N)
		evt.StdStr = fmt.Sprintf("std-%03d", i)
		evt.StlVecF64 = make([]float64, int(evt.N))
		evt.StlVecStr = make([]string, int(evt.N))
		for ii := 0; ii < int(evt.N); ii++ {
			evt.SliF64[ii] = float64(i)
			evt.StlVecF64[ii] = float64(i)
			evt.StlVecStr[ii] = fmt.Sprintf("vec-%03d", i)
		}
		evt.End = fmt.Sprintf("end-%03d", i)

		return data
	}

	sc, err := rootio.NewTreeScanner(chain, &Data{})
	if err != nil {
		t.Fatal(err)
	}
	defer sc.Close()

	for sc.Next() {
		var d1 Data
		err := sc.Scan(&d1)
		if err != nil {
			t.Fatal(err)
		}
		i := sc.Entry()
		if !reflect.DeepEqual(d1, want(i)) {
			t.Fatalf("entry[%d]:\ngot= %#v\nwant=%#v\n", i, d1, want(i))
		}

		var d2 Data
		err = sc.Scan(&d2)
		if err != nil {
			t.Fatal(err)
		}
		if !reflect.DeepEqual(d2, want(i)) {
			t.Fatalf("entry[%d]:\ngot= %#v\nwant=%#v\n", i, d2, want(i))
		}
	}
	if err := sc.Err(); err != nil && err != io.EOF {
		t.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Only files including IPv4, IPv6, and Location (in english)
// will be read and parsed into lists.
package parser

import (
	"encoding/binary"
	"encoding/csv"
	"errors"
	"io"
	"log"
	"net"
	"regexp"
	"strconv"
	"strings"
)

const ipNumColumnsGlite2 = 10
const locationNumColumnsGlite2 = 13
const ipNumColumnsGliteLatest = 3
const mapMax = 200000
const gLiteLatestPrefix = "GeoLiteCity-Blocks"
const gLite2CityPrefix = "GeoLite2-City-Blocks"

// IPNode defines IPv4 and IPv6 databases
type IPNode struct {
	IPAddressLow  net.IP
	IPAddressHigh net.IP
	LocationIndex int // Index to slice of locations
	PostalCode    string
	Latitude      float64
	Longitude     float64
}

// LocationNode defines Location databases
type LocationNode struct {
	GeonameID     int
	ContinentCode string
	CountryCode   string
	CountryName   string
	MetroCode     int64
	CityName      string
}

// Creates a List of IPNodes
func CreateIPList(reader io.Reader, idMap map[int]int, file string) ([]IPNode, error) {
	list := []IPNode{}
	r := csv.NewReader(reader)
	r.TrimLeadingSpace = true
	// Skip first line
	_, err := r.Read()
	if err == io.EOF {
		log.Println("Empty input data")
		return list, errors.New("Empty input data")
	}

	if file == "GeoLiteCity-Blocks.csv" {
		_, err := r.Read()
		if err == io.EOF {
			log.Println("Empty input data")
			return list, errors.New("Empty input data")
		}
	}
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		var newNode IPNode
		switch {
		case strings.HasPrefix(file, gLiteLatestPrefix):
			err = checkColumnLength(record, ipNumColumnsGliteLatest)
			if err != nil {
				return nil, err
			}
			var newNode IPNode
			newNode.IPAddressLow, err = Int2ip(record[0])
			if err != nil {
				return nil, err
			}
			newNode.IPAddressHigh, err = Int2ip(record[1])
			if err != nil {
				return nil, err
			}
			index, err := validateGeoId(record[2], idMap)
			if err != nil {
				return nil, err
			}
			newNode.LocationIndex = index
			list = append(list, newNode)
		case strings.HasPrefix(file, gLite2CityPrefix):
			err = checkColumnLength(record, ipNumColumnsGlite2)
			if err != nil {
				return nil, err
			}
			newNode.IPAddressLow, err = RangeCIDR(record[0], true)
			if err != nil {
				return nil, err
			}
			newNode.IPAddressHigh, err = RangeCIDR(record[0], false)
			if err != nil {
				return nil, err
			}
			index, err := validateGeoId(record[1], idMap)
			if err != nil {
				return nil, err
			}
			newNode.LocationIndex = index
			newNode.PostalCode = record[6]
			newNode.Latitude, err = stringToFloat(record[7], "Latitude")
			if err != nil {
				return nil, err
			}
			newNode.Longitude, err = stringToFloat(record[8], "Longitude")
			if err != nil {
				return nil, err
			}
			list = append(list, newNode)
		default:
			log.Println("Unaccepted csv file provided: ", file)
			return list, errors.New("Unaccepted csv file provided")
		}
	}
	return list, nil
}

// Verify column length
func checkColumnLength(record []string, size int) error {
	if len(record) != size {
		log.Println("Incorrect number of columns in IP list", size, " got: ", len(record), record)
		return errors.New("Corrupted Data: wrong number of columns")
	}
	return nil
}

// Converts integer to IPv4
func Int2ip(str string) (net.IP, error) {
	num, err := strconv.Atoi(str)
	if err != nil {
		log.Println("Provided IP should be a number")
		return nil, errors.New("Inputed string cannot be converted to a number")
	}
	ip := make(net.IP, 4)
	binary.BigEndian.PutUint32(ip, uint32(num))
	return ip, nil
}

// Finds the smallest and largest net.IP from a CIDR range
// Example: "1.0.0.0/24" -> Low bound: 1.0.0.0 High bound: 1.0.0.255
func RangeCIDR(cidr string, lowBound bool) (net.IP, error) {
	ip, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		return nil, errors.New("Invalid CIDR IP range")
	}
	if lowBound {
		return ip, nil
	}
	mask := ipnet.Mask
	for x, _ := range ip {
		if len(mask) == 4 {
			if x < 12 {
				ip[x] |= 0
			} else {
				ip[x] |= ^mask[x-12]
			}
		} else {
			ip[x] |= ^mask[x]
		}
	}
	return ip, nil
}

// Returns the index of geonameId within idMap
func validateGeoId(gnid string, idMap map[int]int) (int, error) {
	geonameId, err := strconv.Atoi(gnid)
	if err != nil {
		log.Println("geonameID should be a number")
		return 0, errors.New("Corrupted Data: geonameID should be a number")
	}
	loadIndex, ok := idMap[geonameId]
	if !ok {
		log.Println("geonameID not found ", geonameId)
		return 0, errors.New("Corrupted Data: geonameId not found")
	}
	return loadIndex, nil
}

func stringToFloat(str, field string) (float64, error) {
	flt, err := strconv.ParseFloat(str, 64)
	if err != nil {
		if len(str) > 0 {
			log.Println(field, " was not a number")
			output := strings.Join([]string{"Corrupted Data: ", field, " should be an int"}, "")
			return 0, errors.New(output)

		}
	}
	return flt, nil
}

func checkAllCaps(str, field string) (string, error) {
	match, _ := regexp.MatchString("^[A-Z]*$", str)
	if match {
		return str, nil
	} else {
		log.Println(field, "should be all capitals and no numbers")
		output := strings.Join([]string{"Corrupted Data: ", field, " should be all caps"}, "")
		return "", errors.New(output)

	}
}

// Creates list for location databases
func CreateLocationList(reader io.Reader) ([]LocationNode, map[int]int, error) {
	idMap := make(map[int]int, mapMax)
	list := []LocationNode{}
	r := csv.NewReader(reader)
	r.TrimLeadingSpace = true
	// Skip the first line
	_, err := r.Read()
	if err == io.EOF {
		log.Println("Empty input data")
		return nil, nil, errors.New("Empty input data")
	}
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		if len(record) != locationNumColumnsGlite2 {
			log.Println("Incorrect number of columns in Location list\n\twanted: ", locationNumColumnsGlite2, " got: ", len(record), record)
			return nil, nil, errors.New("Corrupted Data: wrong number of columns")
		}
		var newNode LocationNode
		newNode.GeonameID, err = strconv.Atoi(record[0])
		if err != nil {
			if len(record[0]) > 0 {
				log.Println("GeonameID should be a number")
				return nil, nil, errors.New("Corrupted Data: GeonameID should be a number")
			}
		}
		newNode.ContinentCode, err = checkAllCaps(record[2], "Continent code")
		if err != nil {
			return nil, nil, err
		}
		newNode.CountryCode, err = checkAllCaps(record[4], "Country code")
		if err != nil {
			return nil, nil, err
		}
		match, _ := regexp.MatchString("^[a-zA-Z]*$", record[5])
		if match {
			newNode.CountryName = record[5]
		} else {
			log.Println("Country name should be letters only")
			return nil, nil, errors.New("Corrupted Data: country name should be letters")
		}
		newNode.MetroCode, err = strconv.ParseInt(record[11], 10, 64)
		if err != nil {
			if len(record[11]) > 0 {
				log.Println("MetroCode should be a number")
				return nil, nil, errors.New("Corrupted Data: metrocode should be a number")
			}
		}
		newNode.CityName = record[10]
		list = append(list, newNode)
		idMap[newNode.GeonameID] = len(list) - 1
	}
	return list, idMap, nil
}
<commit_msg>simplify parseCIDR code<commit_after>// Only files including IPv4, IPv6, and Location (in english)
// will be read and parsed into lists.
package parser

import (
	"encoding/binary"
	"encoding/csv"
	"errors"
	"io"
	"log"
	"net"
	"regexp"
	"strconv"
	"strings"
)

const ipNumColumnsGlite2 = 10
const locationNumColumnsGlite2 = 13
const ipNumColumnsGliteLatest = 3
const mapMax = 200000
const gLiteLatestPrefix = "GeoLiteCity-Blocks"
const gLite2CityPrefix = "GeoLite2-City-Blocks"

// IPNode defines IPv4 and IPv6 databases
type IPNode struct {
	IPAddressLow  net.IP
	IPAddressHigh net.IP
	LocationIndex int // Index to slice of locations
	PostalCode    string
	Latitude      float64
	Longitude     float64
}

// LocationNode defines Location databases
type LocationNode struct {
	GeonameID     int
	ContinentCode string
	CountryCode   string
	CountryName   string
	MetroCode     int64
	CityName      string
}

// Creates a List of IPNodes
func CreateIPList(reader io.Reader, idMap map[int]int, file string) ([]IPNode, error) {
	list := []IPNode{}
	r := csv.NewReader(reader)
	r.TrimLeadingSpace = true
	// Skip first line
	_, err := r.Read()
	if err == io.EOF {
		log.Println("Empty input data")
		return list, errors.New("Empty input data")
	}

	if file == "GeoLiteCity-Blocks.csv" {
		_, err := r.Read()
		if err == io.EOF {
			log.Println("Empty input data")
			return list, errors.New("Empty input data")
		}
	}
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		var newNode IPNode
		switch {
		case strings.HasPrefix(file, gLiteLatestPrefix):
			err = checkColumnLength(record, ipNumColumnsGliteLatest)
			if err != nil {
				return nil, err
			}
			var newNode IPNode
			newNode.IPAddressLow, err = Int2ip(record[0])
			if err != nil {
				return nil, err
			}
			newNode.IPAddressHigh, err = Int2ip(record[1])
			if err != nil {
				return nil, err
			}
			index, err := validateGeoId(record[2], idMap)
			if err != nil {
				return nil, err
			}
			newNode.LocationIndex = index
			list = append(list, newNode)
		case strings.HasPrefix(file, gLite2CityPrefix):
			err = checkColumnLength(record, ipNumColumnsGlite2)
			if err != nil {
				return nil, err
			}
			lowIp, highIp, err := RangeCIDR(record[0])
			if err != nil {
				return nil, err
			}
			newNode.IPAddressLow = lowIp
			newNode.IPAddressHigh = highIp
			index, err := validateGeoId(record[1], idMap)
			if err != nil {
				return nil, err
			}
			newNode.LocationIndex = index
			newNode.PostalCode = record[6]
			newNode.Latitude, err = stringToFloat(record[7], "Latitude")
			if err != nil {
				return nil, err
			}
			newNode.Longitude, err = stringToFloat(record[8], "Longitude")
			if err != nil {
				return nil, err
			}
			list = append(list, newNode)
		default:
			log.Println("Unaccepted csv file provided: ", file)
			return list, errors.New("Unaccepted csv file provided")
		}
	}
	return list, nil
}

// Verify column length
func checkColumnLength(record []string, size int) error {
	if len(record) != size {
		log.Println("Incorrect number of columns in IP list", size, " got: ", len(record), record)
		return errors.New("Corrupted Data: wrong number of columns")
	}
	return nil
}

// Converts integer to net.IP
func Int2ip(str string) (net.IP, error) {
	num, err := strconv.Atoi(str)
	if err != nil {
		log.Println("Provided IP should be a number")
		return nil, errors.New("Inputed string cannot be converted to a number")
	}
	ip := make(net.IP, 4)
	binary.BigEndian.PutUint32(ip, uint32(num))
	return ip, nil
}

// Finds the smallest and largest net.IP from a CIDR range
// Example: "1.0.0.0/24" -> 1.0.0.0 , 1.0.0.255
func RangeCIDR(cidr string) (net.IP, net.IP, error) {
	ip, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		return nil, nil, errors.New("Invalid CIDR IP range")
	}
	lowIp := make(net.IP, len(ip))
	copy(lowIp, ip)
	mask := ipnet.Mask
	for x, _ := range ip {
		if len(mask) == 4 {
			if x < 12 {
				ip[x] |= 0
			} else {
				ip[x] |= ^mask[x-12]
			}
		} else {
			ip[x] |= ^mask[x]
		}
	}
	return lowIp, ip, nil
}

// Returns the index of geonameId within idMap
func validateGeoId(gnid string, idMap map[int]int) (int, error) {
	geonameId, err := strconv.Atoi(gnid)
	if err != nil {
		log.Println("geonameID should be a number")
		return 0, errors.New("Corrupted Data: geonameID should be a number")
	}
	loadIndex, ok := idMap[geonameId]
	if !ok {
		log.Println("geonameID not found ", geonameId)
		return 0, errors.New("Corrupted Data: geonameId not found")
	}
	return loadIndex, nil
}

func stringToFloat(str, field string) (float64, error) {
	flt, err := strconv.ParseFloat(str, 64)
	if err != nil {
		if len(str) > 0 {
			log.Println(field, " was not a number")
			output := strings.Join([]string{"Corrupted Data: ", field, " should be an int"}, "")
			return 0, errors.New(output)

		}
	}
	return flt, nil
}

func checkAllCaps(str, field string) (string, error) {
	match, _ := regexp.MatchString("^[A-Z]*$", str)
	if match {
		return str, nil
	} else {
		log.Println(field, "should be all capitals and no numbers")
		output := strings.Join([]string{"Corrupted Data: ", field, " should be all caps"}, "")
		return "", errors.New(output)

	}
}

// Creates list for location databases
func CreateLocationList(reader io.Reader) ([]LocationNode, map[int]int, error) {
	idMap := make(map[int]int, mapMax)
	list := []LocationNode{}
	r := csv.NewReader(reader)
	r.TrimLeadingSpace = true
	// Skip the first line
	_, err := r.Read()
	if err == io.EOF {
		log.Println("Empty input data")
		return nil, nil, errors.New("Empty input data")
	}
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		if len(record) != locationNumColumnsGlite2 {
			log.Println("Incorrect number of columns in Location list\n\twanted: ", locationNumColumnsGlite2, " got: ", len(record), record)
			return nil, nil, errors.New("Corrupted Data: wrong number of columns")
		}
		var newNode LocationNode
		newNode.GeonameID, err = strconv.Atoi(record[0])
		if err != nil {
			if len(record[0]) > 0 {
				log.Println("GeonameID should be a number")
				return nil, nil, errors.New("Corrupted Data: GeonameID should be a number")
			}
		}
		newNode.ContinentCode, err = checkAllCaps(record[2], "Continent code")
		if err != nil {
			return nil, nil, err
		}
		newNode.CountryCode, err = checkAllCaps(record[4], "Country code")
		if err != nil {
			return nil, nil, err
		}
		match, _ := regexp.MatchString("^[a-zA-Z]*$", record[5])
		if match {
			newNode.CountryName = record[5]
		} else {
			log.Println("Country name should be letters only")
			return nil, nil, errors.New("Corrupted Data: country name should be letters")
		}
		newNode.MetroCode, err = strconv.ParseInt(record[11], 10, 64)
		if err != nil {
			if len(record[11]) > 0 {
				log.Println("MetroCode should be a number")
				return nil, nil, errors.New("Corrupted Data: metrocode should be a number")
			}
		}
		newNode.CityName = record[10]
		list = append(list, newNode)
		idMap[newNode.GeonameID] = len(list) - 1
	}
	return list, idMap, nil
}
<|endoftext|>


<commit_before>package parser

import (
	"errors"
	"fmt"
	"regexp"
	"strconv"
	"strings"

	"github.com/twtiger/gosecco/tree"
)

var (
	allowRE      = regexp.MustCompile(`^ *1$`)
	returnRE     = regexp.MustCompile(`^ *return *([[:word:]]+)$`)
	exprReturnRE = regexp.MustCompile(`; *return *([[:word:]]+)$`)
)

type traceData struct {
	file string
	line int
}

type parser struct {
	forBinding bool
	traceData  *traceData
}

func parseExpressionForBinding(expr string) (tree.Expression, bool, uint16, error) {
	return newParser(true, nil).parseExpression(expr)
}

func parseExpression(expr string) (tree.Expression, bool, uint16, error) {
	return newParser(false, nil).parseExpression(expr)
}

func newParser(forBinding bool, td *traceData) *parser {
	return &parser{
		forBinding,
		td,
	}
}

func (p *parser) parseSpecialCases(expr string) (tree.Expression, bool, uint16, bool, string, error) {
	hasRet := false
	ret := uint16(0)
	newExpr := expr
	if !p.forBinding {
		if match := allowRE.FindStringSubmatch(expr); match != nil {
			return tree.BooleanLiteral{true}, false, 0, true, newExpr, nil
		}

		if match := returnRE.FindStringSubmatch(expr); match != nil {
			errno, err := strconv.ParseUint(match[1], 0, 16)
			if err == nil {
				return nil, true, uint16(errno), true, newExpr, nil
			}
			return nil, false, 0, true, newExpr, err
		}

		if match := exprReturnRE.FindStringSubmatch(expr); match != nil {
			newExpr = strings.TrimSuffix(expr, match[0])
			errno, err := strconv.ParseUint(match[1], 0, 16)
			if err == nil {
				hasRet = true
				ret = uint16(errno)
			} else {
				return nil, false, 0, true, newExpr, err
			}
		}
	}
	return nil, hasRet, ret, false, newExpr, nil
}

func (p *parser) parseExpression(expr string) (tree.Expression, bool, uint16, error) {
	expression, hasRet, ret, done, expr, err := p.parseSpecialCases(expr)
	if done {
		return expression, hasRet, ret, err
	}

	tokens, err := tokenize(expr, func(ts, te int, data []byte) error {
		trace := "<input>:-1"
		if p.traceData != nil {
			trace = fmt.Sprintf("%s:%d", p.traceData.file, p.traceData.line)
		}
		return fmt.Errorf("unexpected token at %s:%d: '%s'", trace, ts, data[ts:te])

	})
	if err != nil {
		return nil, false, 0, err
	}
	ctx := parseContext{0, tokens, len(tokens) == 0, p}

	if ctx.atEnd {
		return nil, hasRet, ret, nil
	}

	expression, err = ctx.logicalORExpression()
	if err != nil {
		return nil, false, 0, err
	}

	if err = ctx.end(); err != nil {
		return nil, false, 0, err
	}

	return expression, hasRet, ret, nil
}

func (ctx *parseContext) genErr(exp string) error {
	found := "EOF"
	if !ctx.atEnd {
		td := ""
		if ctx.tokens[ctx.index].td != nil {
			td = " " + string(ctx.tokens[ctx.index].td)
		}
		found = fmt.Sprintf("'%s'%s", tokens[ctx.tokens[ctx.index].t], td)
	}
	return fmt.Errorf("expression is invalid. unable to parse: expected %s, found %s", exp, found)
}

func (ctx *parseContext) end() error {
	if !ctx.atEnd {
		return ctx.genErr("EOF")
	}
	return nil
}

func (ctx *parseContext) logicalORExpression() (tree.Expression, error) {
	left, e := ctx.logicalANDExpression()
	if e != nil {
		return nil, e
	}
	if ctx.next() == LOR {
		ctx.consume()
		right, e := ctx.logicalORExpression()
		if e != nil {
			return nil, e
		}
		return tree.Or{Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) logicalANDExpression() (tree.Expression, error) {
	left, e := ctx.equalityExpression()
	if e != nil {
		return nil, e
	}
	if ctx.next() == LAND {
		ctx.consume()
		right, e := ctx.logicalANDExpression()
		if e != nil {
			return nil, e
		}
		return tree.And{Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) equalityExpression() (tree.Expression, error) {
	left, e := ctx.relationalExpression()
	if e != nil {
		return nil, e
	}
	switch ctx.next() {
	case EQL, NEQ:
		op, _ := ctx.consume()
		right, e := ctx.equalityExpression()
		if e != nil {
			return nil, e
		}
		return tree.Comparison{Op: comparisonOperator[op], Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) relationalExpression() (tree.Expression, error) {
	left, e := ctx.inclusiveORExpression()
	if e != nil {
		return nil, e
	}
	switch ctx.next() {
	case LT, GT, LTE, GTE:
		op, _ := ctx.consume()
		right, e := ctx.relationalExpression()
		if e != nil {
			return nil, e
		}
		return tree.Comparison{Op: comparisonOperator[op], Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) inclusiveORExpression() (tree.Expression, error) {
	left, e := ctx.exclusiveORExpression()
	if e != nil {
		return nil, e
	}
	if ctx.next() == OR {
		ctx.consume()
		right, e := ctx.inclusiveORExpression()
		if e != nil {
			return nil, e
		}
		return tree.Arithmetic{Op: tree.BINOR, Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) exclusiveORExpression() (tree.Expression, error) {
	left, e := ctx.andExpression()
	if e != nil {
		return nil, e
	}
	if ctx.next() == XOR {
		ctx.consume()
		right, e := ctx.exclusiveORExpression()
		if e != nil {
			return nil, e
		}

		return tree.Arithmetic{Op: tree.BINXOR, Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) andExpression() (tree.Expression, error) {
	left, e := ctx.shiftExpression()
	if e != nil {
		return nil, e
	}
	if ctx.next() == AND {
		ctx.consume()
		right, e := ctx.andExpression()
		if e != nil {
			return nil, e
		}
		return tree.Arithmetic{Op: tree.BINAND, Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) shiftExpression() (tree.Expression, error) {
	left, e := ctx.additiveExpression()
	if e != nil {
		return nil, e
	}
	switch ctx.next() {
	case LSH, RSH:
		op, _ := ctx.consume()
		right, e := ctx.shiftExpression()
		if e != nil {
			return nil, e
		}
		return tree.Arithmetic{Op: shiftOperator[op], Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) additiveExpression() (tree.Expression, error) {
	left, e := ctx.multiplicativeExpression()
	if e != nil {
		return nil, e
	}
	switch ctx.next() {
	case ADD, SUB:
		op, _ := ctx.consume()
		right, e := ctx.additiveExpression()
		if e != nil {
			return nil, e
		}
		return tree.Arithmetic{Op: addOperator[op], Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) multiplicativeExpression() (tree.Expression, error) {
	left, e := ctx.unaryExpression()
	if e != nil {
		return nil, e
	}
	switch ctx.next() {
	case MUL, DIV, MOD:
		op, _ := ctx.consume()
		right, e := ctx.multiplicativeExpression()
		if e != nil {
			return nil, e
		}
		return tree.Arithmetic{Op: multOperator[op], Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) unaryExpression() (tree.Expression, error) {
	switch ctx.next() {
	case INV:
		ctx.consume()
		left, e := ctx.primary()
		if e != nil {
			return nil, e
		}
		return tree.BinaryNegation{left}, nil
	case NOT:
		ctx.consume()
		left, e := ctx.primary()
		if e != nil {
			return nil, e
		}
		return tree.Negation{left}, nil
	}
	return ctx.primary()
}

func (ctx *parseContext) collectArgs() ([]tree.Any, error) {
	args := []tree.Any{}
	for ctx.next() != RPAREN {
		ctx.consume() //consume the opening paren, or last comma
		res, e := ctx.logicalORExpression()
		if e != nil {
			return nil, e
		}
		args = append(args, res)
		switch ctx.next() {
		case RPAREN, COMMA:
			// Do nothing here
		default:
			return nil, ctx.genErr("')' or ','")
		}
	}
	ctx.consume()
	return args, nil
}

func (ctx *parseContext) collectNumerics() ([]tree.Numeric, error) {
	args := []tree.Numeric{}
	for ctx.next() != RPAREN {
		ctx.consume() //consume the opening paren, or last comma
		res, e := ctx.logicalORExpression()
		if e != nil {
			return nil, e
		}
		args = append(args, res)
		switch ctx.next() {
		case RPAREN, COMMA:
			// Do nothing here
		default:
			return nil, ctx.genErr("')' or ','")
		}
	}
	ctx.consume()
	return args, nil
}

func (ctx *parseContext) primary() (tree.Expression, error) {
	switch ctx.next() {
	case LPAREN:
		ctx.consume()
		val, e := ctx.logicalORExpression()
		if e != nil {
			return nil, e
		}
		op, _ := ctx.consume()
		if op != RPAREN {
			return nil, ctx.genErr("')'")
		}
		return val, nil
	case ARG:
		_, data := ctx.consume()
		tp := tree.Full
		pref := strings.TrimPrefix(string(data), "arg")
		if strings.HasPrefix(pref, "H") {
			tp = tree.Hi
			pref = strings.TrimPrefix(pref, "H")
		}
		if strings.HasPrefix(pref, "L") {
			tp = tree.Low
			pref = strings.TrimPrefix(pref, "L")
		}
		val, _ := strconv.Atoi(pref)
		// This should never error out
		return tree.Argument{Index: val, Type: tp}, nil
	case IDENT:
		_, data := ctx.consume()
		if ctx.next() == LPAREN {
			args, e := ctx.collectArgs()
			if e != nil {
				return nil, e
			}
			return tree.Call{Name: string(data), Args: args}, nil
		}
		return tree.Variable{string(data)}, nil
	case IN, NOTIN:
		op, _ := ctx.consume()
		if ctx.next() == LPAREN {
			all, e := ctx.collectNumerics()
			if e != nil {
				return nil, e
			}
			return tree.Inclusion{Positive: op == IN, Left: all[0], Rights: all[1:]}, nil
		}
		return nil, ctx.genErr("'('")
	case INT:
		_, data := ctx.consume()
		val, _ := strconv.ParseUint(string(data), 0, 32)
		return tree.NumericLiteral{uint64(val)}, nil
	case TRUE:
		ctx.consume()
		return tree.BooleanLiteral{true}, nil
	case FALSE:
		ctx.consume()
		return tree.BooleanLiteral{false}, nil
	case EOF:
		return nil, errors.New("unexpected end of line")
	}

	return nil, ctx.genErr("primary expression")
}
<commit_msg>Make sure that the parser actually CAN parse 64 bit numbers<commit_after>package parser

import (
	"errors"
	"fmt"
	"regexp"
	"strconv"
	"strings"

	"github.com/twtiger/gosecco/tree"
)

var (
	allowRE      = regexp.MustCompile(`^ *1$`)
	returnRE     = regexp.MustCompile(`^ *return *([[:word:]]+)$`)
	exprReturnRE = regexp.MustCompile(`; *return *([[:word:]]+)$`)
)

type traceData struct {
	file string
	line int
}

type parser struct {
	forBinding bool
	traceData  *traceData
}

func parseExpressionForBinding(expr string) (tree.Expression, bool, uint16, error) {
	return newParser(true, nil).parseExpression(expr)
}

func parseExpression(expr string) (tree.Expression, bool, uint16, error) {
	return newParser(false, nil).parseExpression(expr)
}

func newParser(forBinding bool, td *traceData) *parser {
	return &parser{
		forBinding,
		td,
	}
}

func (p *parser) parseSpecialCases(expr string) (tree.Expression, bool, uint16, bool, string, error) {
	hasRet := false
	ret := uint16(0)
	newExpr := expr
	if !p.forBinding {
		if match := allowRE.FindStringSubmatch(expr); match != nil {
			return tree.BooleanLiteral{true}, false, 0, true, newExpr, nil
		}

		if match := returnRE.FindStringSubmatch(expr); match != nil {
			errno, err := strconv.ParseUint(match[1], 0, 16)
			if err == nil {
				return nil, true, uint16(errno), true, newExpr, nil
			}
			return nil, false, 0, true, newExpr, err
		}

		if match := exprReturnRE.FindStringSubmatch(expr); match != nil {
			newExpr = strings.TrimSuffix(expr, match[0])
			errno, err := strconv.ParseUint(match[1], 0, 16)
			if err == nil {
				hasRet = true
				ret = uint16(errno)
			} else {
				return nil, false, 0, true, newExpr, err
			}
		}
	}
	return nil, hasRet, ret, false, newExpr, nil
}

func (p *parser) parseExpression(expr string) (tree.Expression, bool, uint16, error) {
	expression, hasRet, ret, done, expr, err := p.parseSpecialCases(expr)
	if done {
		return expression, hasRet, ret, err
	}

	tokens, err := tokenize(expr, func(ts, te int, data []byte) error {
		trace := "<input>:-1"
		if p.traceData != nil {
			trace = fmt.Sprintf("%s:%d", p.traceData.file, p.traceData.line)
		}
		return fmt.Errorf("unexpected token at %s:%d: '%s'", trace, ts, data[ts:te])

	})
	if err != nil {
		return nil, false, 0, err
	}
	ctx := parseContext{0, tokens, len(tokens) == 0, p}

	if ctx.atEnd {
		return nil, hasRet, ret, nil
	}

	expression, err = ctx.logicalORExpression()
	if err != nil {
		return nil, false, 0, err
	}

	if err = ctx.end(); err != nil {
		return nil, false, 0, err
	}

	return expression, hasRet, ret, nil
}

func (ctx *parseContext) genErr(exp string) error {
	found := "EOF"
	if !ctx.atEnd {
		td := ""
		if ctx.tokens[ctx.index].td != nil {
			td = " " + string(ctx.tokens[ctx.index].td)
		}
		found = fmt.Sprintf("'%s'%s", tokens[ctx.tokens[ctx.index].t], td)
	}
	return fmt.Errorf("expression is invalid. unable to parse: expected %s, found %s", exp, found)
}

func (ctx *parseContext) end() error {
	if !ctx.atEnd {
		return ctx.genErr("EOF")
	}
	return nil
}

func (ctx *parseContext) logicalORExpression() (tree.Expression, error) {
	left, e := ctx.logicalANDExpression()
	if e != nil {
		return nil, e
	}
	if ctx.next() == LOR {
		ctx.consume()
		right, e := ctx.logicalORExpression()
		if e != nil {
			return nil, e
		}
		return tree.Or{Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) logicalANDExpression() (tree.Expression, error) {
	left, e := ctx.equalityExpression()
	if e != nil {
		return nil, e
	}
	if ctx.next() == LAND {
		ctx.consume()
		right, e := ctx.logicalANDExpression()
		if e != nil {
			return nil, e
		}
		return tree.And{Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) equalityExpression() (tree.Expression, error) {
	left, e := ctx.relationalExpression()
	if e != nil {
		return nil, e
	}
	switch ctx.next() {
	case EQL, NEQ:
		op, _ := ctx.consume()
		right, e := ctx.equalityExpression()
		if e != nil {
			return nil, e
		}
		return tree.Comparison{Op: comparisonOperator[op], Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) relationalExpression() (tree.Expression, error) {
	left, e := ctx.inclusiveORExpression()
	if e != nil {
		return nil, e
	}
	switch ctx.next() {
	case LT, GT, LTE, GTE:
		op, _ := ctx.consume()
		right, e := ctx.relationalExpression()
		if e != nil {
			return nil, e
		}
		return tree.Comparison{Op: comparisonOperator[op], Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) inclusiveORExpression() (tree.Expression, error) {
	left, e := ctx.exclusiveORExpression()
	if e != nil {
		return nil, e
	}
	if ctx.next() == OR {
		ctx.consume()
		right, e := ctx.inclusiveORExpression()
		if e != nil {
			return nil, e
		}
		return tree.Arithmetic{Op: tree.BINOR, Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) exclusiveORExpression() (tree.Expression, error) {
	left, e := ctx.andExpression()
	if e != nil {
		return nil, e
	}
	if ctx.next() == XOR {
		ctx.consume()
		right, e := ctx.exclusiveORExpression()
		if e != nil {
			return nil, e
		}

		return tree.Arithmetic{Op: tree.BINXOR, Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) andExpression() (tree.Expression, error) {
	left, e := ctx.shiftExpression()
	if e != nil {
		return nil, e
	}
	if ctx.next() == AND {
		ctx.consume()
		right, e := ctx.andExpression()
		if e != nil {
			return nil, e
		}
		return tree.Arithmetic{Op: tree.BINAND, Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) shiftExpression() (tree.Expression, error) {
	left, e := ctx.additiveExpression()
	if e != nil {
		return nil, e
	}
	switch ctx.next() {
	case LSH, RSH:
		op, _ := ctx.consume()
		right, e := ctx.shiftExpression()
		if e != nil {
			return nil, e
		}
		return tree.Arithmetic{Op: shiftOperator[op], Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) additiveExpression() (tree.Expression, error) {
	left, e := ctx.multiplicativeExpression()
	if e != nil {
		return nil, e
	}
	switch ctx.next() {
	case ADD, SUB:
		op, _ := ctx.consume()
		right, e := ctx.additiveExpression()
		if e != nil {
			return nil, e
		}
		return tree.Arithmetic{Op: addOperator[op], Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) multiplicativeExpression() (tree.Expression, error) {
	left, e := ctx.unaryExpression()
	if e != nil {
		return nil, e
	}
	switch ctx.next() {
	case MUL, DIV, MOD:
		op, _ := ctx.consume()
		right, e := ctx.multiplicativeExpression()
		if e != nil {
			return nil, e
		}
		return tree.Arithmetic{Op: multOperator[op], Left: left, Right: right}, nil
	}
	return left, nil
}

func (ctx *parseContext) unaryExpression() (tree.Expression, error) {
	switch ctx.next() {
	case INV:
		ctx.consume()
		left, e := ctx.primary()
		if e != nil {
			return nil, e
		}
		return tree.BinaryNegation{left}, nil
	case NOT:
		ctx.consume()
		left, e := ctx.primary()
		if e != nil {
			return nil, e
		}
		return tree.Negation{left}, nil
	}
	return ctx.primary()
}

func (ctx *parseContext) collectArgs() ([]tree.Any, error) {
	args := []tree.Any{}
	for ctx.next() != RPAREN {
		ctx.consume() //consume the opening paren, or last comma
		res, e := ctx.logicalORExpression()
		if e != nil {
			return nil, e
		}
		args = append(args, res)
		switch ctx.next() {
		case RPAREN, COMMA:
			// Do nothing here
		default:
			return nil, ctx.genErr("')' or ','")
		}
	}
	ctx.consume()
	return args, nil
}

func (ctx *parseContext) collectNumerics() ([]tree.Numeric, error) {
	args := []tree.Numeric{}
	for ctx.next() != RPAREN {
		ctx.consume() //consume the opening paren, or last comma
		res, e := ctx.logicalORExpression()
		if e != nil {
			return nil, e
		}
		args = append(args, res)
		switch ctx.next() {
		case RPAREN, COMMA:
			// Do nothing here
		default:
			return nil, ctx.genErr("')' or ','")
		}
	}
	ctx.consume()
	return args, nil
}

func (ctx *parseContext) primary() (tree.Expression, error) {
	switch ctx.next() {
	case LPAREN:
		ctx.consume()
		val, e := ctx.logicalORExpression()
		if e != nil {
			return nil, e
		}
		op, _ := ctx.consume()
		if op != RPAREN {
			return nil, ctx.genErr("')'")
		}
		return val, nil
	case ARG:
		_, data := ctx.consume()
		tp := tree.Full
		pref := strings.TrimPrefix(string(data), "arg")
		if strings.HasPrefix(pref, "H") {
			tp = tree.Hi
			pref = strings.TrimPrefix(pref, "H")
		}
		if strings.HasPrefix(pref, "L") {
			tp = tree.Low
			pref = strings.TrimPrefix(pref, "L")
		}
		val, _ := strconv.Atoi(pref)
		// This should never error out
		return tree.Argument{Index: val, Type: tp}, nil
	case IDENT:
		_, data := ctx.consume()
		if ctx.next() == LPAREN {
			args, e := ctx.collectArgs()
			if e != nil {
				return nil, e
			}
			return tree.Call{Name: string(data), Args: args}, nil
		}
		return tree.Variable{string(data)}, nil
	case IN, NOTIN:
		op, _ := ctx.consume()
		if ctx.next() == LPAREN {
			all, e := ctx.collectNumerics()
			if e != nil {
				return nil, e
			}
			return tree.Inclusion{Positive: op == IN, Left: all[0], Rights: all[1:]}, nil
		}
		return nil, ctx.genErr("'('")
	case INT:
		_, data := ctx.consume()
		// This error is ignored since unless the tokenizer is completely broken, this should never break
		val, _ := strconv.ParseUint(string(data), 0, 64)
		return tree.NumericLiteral{uint64(val)}, nil
	case TRUE:
		ctx.consume()
		return tree.BooleanLiteral{true}, nil
	case FALSE:
		ctx.consume()
		return tree.BooleanLiteral{false}, nil
	case EOF:
		return nil, errors.New("unexpected end of line")
	}

	return nil, ctx.genErr("primary expression")
}
<|endoftext|>