<commit_before>// Copyright © 2017 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package broker

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"sort"
	"time"

	ttnlog "github.com/TheThingsNetwork/go-utils/log"
	pb "github.com/TheThingsNetwork/ttn/api/broker"
	pb_discovery "github.com/TheThingsNetwork/ttn/api/discovery"
	"github.com/TheThingsNetwork/ttn/api/fields"
	"github.com/TheThingsNetwork/ttn/api/networkserver"
	pb_lorawan "github.com/TheThingsNetwork/ttn/api/protocol/lorawan"
	"github.com/TheThingsNetwork/ttn/api/trace"
	"github.com/TheThingsNetwork/ttn/core/types"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/fcnt"
	"github.com/brocaar/lorawan"
)

const maxFCntGap = 16384

func (b *broker) HandleUplink(uplink *pb.UplinkMessage) (err error) {
	ctx := b.Ctx.WithFields(fields.Get(uplink))
	start := time.Now()
	deduplicatedUplink := new(pb.DeduplicatedUplinkMessage)
	deduplicatedUplink.ServerTime = start.UnixNano()
	defer func() {
		if err != nil {
			if deduplicatedUplink != nil {
				deduplicatedUplink.Trace = deduplicatedUplink.Trace.WithEvent(trace.DropEvent, "reason", err)
			}
			ctx.WithError(err).Warn("Could not handle uplink")
		} else {
			ctx.WithField("Duration", time.Now().Sub(start)).Info("Handled uplink")
		}
		if deduplicatedUplink != nil && b.monitorStream != nil {
			b.monitorStream.Send(deduplicatedUplink)
		}
	}()

	b.status.uplink.Mark(1)

	uplink.Trace = uplink.Trace.WithEvent(trace.ReceiveEvent)

	// De-duplicate uplink messages
	duplicates := b.deduplicateUplink(uplink)
	if len(duplicates) == 0 {
		return nil
	}
	ctx = ctx.WithField("Duplicates", len(duplicates))

	b.status.uplinkUnique.Mark(1)

	deduplicatedUplink.Payload = duplicates[0].Payload
	deduplicatedUplink.ProtocolMetadata = duplicates[0].ProtocolMetadata
	deduplicatedUplink.Trace = deduplicatedUplink.Trace.WithEvent(trace.DeduplicateEvent,
		"duplicates", len(duplicates),
	)
	for _, duplicate := range duplicates {
		if duplicate.Trace != nil {
			deduplicatedUplink.Trace.Parents = append(deduplicatedUplink.Trace.Parents, duplicate.Trace)
		}
	}

	if deduplicatedUplink.ProtocolMetadata.GetLorawan() == nil {
		return errors.NewErrInvalidArgument("Uplink", "does not contain LoRaWAN metadata")
	}

	// LoRaWAN: Unmarshal
	var phyPayload lorawan.PHYPayload
	err = phyPayload.UnmarshalBinary(deduplicatedUplink.Payload)
	if err != nil {
		return err
	}
	macPayload, ok := phyPayload.MACPayload.(*lorawan.MACPayload)
	if !ok {
		return errors.NewErrInvalidArgument("Uplink", "does not contain a MAC payload")
	}

	// Request devices from NS
	devAddr := types.DevAddr(macPayload.FHDR.DevAddr)
	ctx = ctx.WithFields(ttnlog.Fields{
		"DevAddr": devAddr,
		"FCnt":    macPayload.FHDR.FCnt,
	})
	var getDevicesResp *networkserver.DevicesResponse
	getDevicesResp, err = b.ns.GetDevices(b.Component.GetContext(b.nsToken), &networkserver.DevicesRequest{
		DevAddr: &devAddr,
		FCnt:    macPayload.FHDR.FCnt,
	})
	if err != nil {
		return errors.Wrap(errors.FromGRPCError(err), "NetworkServer did not return devices")
	}
	b.status.deduplication.Update(int64(len(getDevicesResp.Results)))
	if len(getDevicesResp.Results) == 0 {
		return errors.NewErrNotFound(fmt.Sprintf("Device with DevAddr %s and FCnt <= %d", devAddr, macPayload.FHDR.FCnt))
	}
	ctx = ctx.WithField("DevAddrResults", len(getDevicesResp.Results))
	deduplicatedUplink.Trace = deduplicatedUplink.Trace.WithEvent("got devices from networkserver",
		"devices", len(getDevicesResp.Results),
	)

	// Sort by FCntUp to optimize the number of MIC checks
	sort.Sort(ByFCntUp(getDevicesResp.Results))

	// Find AppEUI/DevEUI through MIC check
	var device *pb_lorawan.Device
	var micChecks int
	originalFCnt := macPayload.FHDR.FCnt
	for _, candidate := range getDevicesResp.Results {
		nwkSKey := lorawan.AES128Key(*candidate.NwkSKey)

		// First check with the 16 bit counter
		micChecks++
		ok, err = phyPayload.ValidateMIC(nwkSKey)
		if err != nil {
			return err
		}
		if ok {
			device = candidate
			break
		}

		if candidate.Uses32BitFCnt {
			macPayload.FHDR.FCnt = fcnt.GetFull(candidate.FCntUp, uint16(originalFCnt))

			// If 32 bit counter has different value, perform another MIC check
			if macPayload.FHDR.FCnt != originalFCnt {
				micChecks++
				ok, err = phyPayload.ValidateMIC(nwkSKey)
				if err != nil {
					return err
				}
				if ok {
					device = candidate
					break
				}
			}
		}
	}
	if device == nil {
		return errors.NewErrNotFound("device that validates MIC")
	}

	ctx = ctx.WithFields(ttnlog.Fields{
		"MICChecks": micChecks,
		"DevEUI":    device.DevEui,
		"AppEUI":    device.AppEui,
		"AppID":     device.AppId,
		"DevID":     device.DevId,
		"FCnt":      originalFCnt,
	})
	deduplicatedUplink.DevEui = device.DevEui
	deduplicatedUplink.AppEui = device.AppEui
	deduplicatedUplink.AppId = device.AppId
	deduplicatedUplink.DevId = device.DevId
	deduplicatedUplink.Trace = deduplicatedUplink.Trace.WithEvent(trace.CheckMICEvent, "mic checks", micChecks)
	if macPayload.FHDR.FCnt != originalFCnt {
		ctx = ctx.WithField("RealFCnt", macPayload.FHDR.FCnt)
	}

	switch {
	case macPayload.FHDR.FCnt > device.FCntUp && macPayload.FHDR.FCnt-device.FCntUp <= maxFCntGap:
		// FCnt higher than latest and within max FCnt gap (normal case)
	case device.DisableFCntCheck:
		// FCnt Check disabled. Rely on MIC check only
	case device.FCntUp == 0:
		// FCntUp is reset. We don't know where the device will start sending.
	case macPayload.FHDR.FCnt == device.FCntUp:
		if phyPayload.MHDR.MType == lorawan.ConfirmedDataUp {
			// Retry of confirmed uplink
			break
		}
		fallthrough
	case macPayload.FHDR.FCnt <= device.FCntUp:
		return errors.NewErrInvalidArgument("FCnt", "not high enough")
	case macPayload.FHDR.FCnt-device.FCntUp > maxFCntGap:
		return errors.NewErrInvalidArgument("FCnt", "too high")
	default:
		return errors.NewErrInternal("FCnt check failed")
	}

	// Add FCnt to Metadata (because it's not marshaled in lorawan payload)
	deduplicatedUplink.ProtocolMetadata.GetLorawan().FCnt = macPayload.FHDR.FCnt

	// Collect GatewayMetadata and DownlinkOptions
	var downlinkOptions []*pb.DownlinkOption
	for _, duplicate := range duplicates {
		deduplicatedUplink.GatewayMetadata = append(deduplicatedUplink.GatewayMetadata, duplicate.GatewayMetadata)
		downlinkOptions = append(downlinkOptions, duplicate.DownlinkOptions...)
	}

	// Select best DownlinkOption
	if len(downlinkOptions) > 0 {
		deduplicatedUplink.ResponseTemplate = &pb.DownlinkMessage{
			DevEui:         device.DevEui,
			AppEui:         device.AppEui,
			AppId:          device.AppId,
			DevId:          device.DevId,
			DownlinkOption: selectBestDownlink(downlinkOptions),
		}
	}

	// Pass Uplink through NS
	deduplicatedUplink, err = b.ns.Uplink(b.Component.GetContext(b.nsToken), deduplicatedUplink)
	if err != nil {
		return errors.Wrap(errors.FromGRPCError(err), "NetworkServer did not handle uplink")
	}

	var announcements []*pb_discovery.Announcement
	announcements, err = b.Discovery.GetAllHandlersForAppID(device.AppId)
	if err != nil {
		return err
	}
	if len(announcements) == 0 {
		return errors.NewErrNotFound(fmt.Sprintf("Handler for AppID %s", device.AppId))
	}
	if len(announcements) > 1 {
		return errors.NewErrInternal(fmt.Sprintf("Multiple Handlers for AppID %s", device.AppId))
	}

	var handler chan<- *pb.DeduplicatedUplinkMessage
	handler, err = b.getHandlerUplink(announcements[0].Id)
	if err != nil {
		return err
	}

	deduplicatedUplink.Trace = deduplicatedUplink.Trace.WithEvent(trace.ForwardEvent,
		"handler", announcements[0].Id,
	)

	handler <- deduplicatedUplink

	return nil
}

func (b *broker) deduplicateUplink(duplicate *pb.UplinkMessage) (uplinks []*pb.UplinkMessage) {
	sum := md5.Sum(duplicate.Payload)
	key := hex.EncodeToString(sum[:])
	list := b.uplinkDeduplicator.Deduplicate(key, duplicate)
	if len(list) == 0 {
		return
	}
	for _, duplicate := range list {
		uplinks = append(uplinks, duplicate.(*pb.UplinkMessage))
	}
	return
}

func selectBestDownlink(options []*pb.DownlinkOption) *pb.DownlinkOption {
	sort.Sort(ByScore(options))
	return options[0]
}

// ByFCntUp implements sort.Interface for []*pb_lorawan.Device based on FCnt
type ByFCntUp []*pb_lorawan.Device

func (a ByFCntUp) Len() int      { return len(a) }
func (a ByFCntUp) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByFCntUp) Less(i, j int) bool {
	// Devices that disable the FCnt check have low priority
	if a[i].DisableFCntCheck {
		return 2*int(a[i].FCntUp)+100 < int(a[j].FCntUp)
	}
	return int(a[i].FCntUp) < int(a[j].FCntUp)
}
<commit_msg>[hotfix] Restore original FCnt after 32 bit MIC check<commit_after>// Copyright © 2017 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package broker

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"sort"
	"time"

	ttnlog "github.com/TheThingsNetwork/go-utils/log"
	pb "github.com/TheThingsNetwork/ttn/api/broker"
	pb_discovery "github.com/TheThingsNetwork/ttn/api/discovery"
	"github.com/TheThingsNetwork/ttn/api/fields"
	"github.com/TheThingsNetwork/ttn/api/networkserver"
	pb_lorawan "github.com/TheThingsNetwork/ttn/api/protocol/lorawan"
	"github.com/TheThingsNetwork/ttn/api/trace"
	"github.com/TheThingsNetwork/ttn/core/types"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/fcnt"
	"github.com/brocaar/lorawan"
)

const maxFCntGap = 16384

func (b *broker) HandleUplink(uplink *pb.UplinkMessage) (err error) {
	ctx := b.Ctx.WithFields(fields.Get(uplink))
	start := time.Now()
	deduplicatedUplink := new(pb.DeduplicatedUplinkMessage)
	deduplicatedUplink.ServerTime = start.UnixNano()
	defer func() {
		if err != nil {
			if deduplicatedUplink != nil {
				deduplicatedUplink.Trace = deduplicatedUplink.Trace.WithEvent(trace.DropEvent, "reason", err)
			}
			ctx.WithError(err).Warn("Could not handle uplink")
		} else {
			ctx.WithField("Duration", time.Now().Sub(start)).Info("Handled uplink")
		}
		if deduplicatedUplink != nil && b.monitorStream != nil {
			b.monitorStream.Send(deduplicatedUplink)
		}
	}()

	b.status.uplink.Mark(1)

	uplink.Trace = uplink.Trace.WithEvent(trace.ReceiveEvent)

	// De-duplicate uplink messages
	duplicates := b.deduplicateUplink(uplink)
	if len(duplicates) == 0 {
		return nil
	}
	ctx = ctx.WithField("Duplicates", len(duplicates))

	b.status.uplinkUnique.Mark(1)

	deduplicatedUplink.Payload = duplicates[0].Payload
	deduplicatedUplink.ProtocolMetadata = duplicates[0].ProtocolMetadata
	deduplicatedUplink.Trace = deduplicatedUplink.Trace.WithEvent(trace.DeduplicateEvent,
		"duplicates", len(duplicates),
	)
	for _, duplicate := range duplicates {
		if duplicate.Trace != nil {
			deduplicatedUplink.Trace.Parents = append(deduplicatedUplink.Trace.Parents, duplicate.Trace)
		}
	}

	if deduplicatedUplink.ProtocolMetadata.GetLorawan() == nil {
		return errors.NewErrInvalidArgument("Uplink", "does not contain LoRaWAN metadata")
	}

	// LoRaWAN: Unmarshal
	var phyPayload lorawan.PHYPayload
	err = phyPayload.UnmarshalBinary(deduplicatedUplink.Payload)
	if err != nil {
		return err
	}
	macPayload, ok := phyPayload.MACPayload.(*lorawan.MACPayload)
	if !ok {
		return errors.NewErrInvalidArgument("Uplink", "does not contain a MAC payload")
	}

	// Request devices from NS
	devAddr := types.DevAddr(macPayload.FHDR.DevAddr)
	ctx = ctx.WithFields(ttnlog.Fields{
		"DevAddr": devAddr,
		"FCnt":    macPayload.FHDR.FCnt,
	})
	var getDevicesResp *networkserver.DevicesResponse
	getDevicesResp, err = b.ns.GetDevices(b.Component.GetContext(b.nsToken), &networkserver.DevicesRequest{
		DevAddr: &devAddr,
		FCnt:    macPayload.FHDR.FCnt,
	})
	if err != nil {
		return errors.Wrap(errors.FromGRPCError(err), "NetworkServer did not return devices")
	}
	b.status.deduplication.Update(int64(len(getDevicesResp.Results)))
	if len(getDevicesResp.Results) == 0 {
		return errors.NewErrNotFound(fmt.Sprintf("Device with DevAddr %s and FCnt <= %d", devAddr, macPayload.FHDR.FCnt))
	}
	ctx = ctx.WithField("DevAddrResults", len(getDevicesResp.Results))
	deduplicatedUplink.Trace = deduplicatedUplink.Trace.WithEvent("got devices from networkserver",
		"devices", len(getDevicesResp.Results),
	)

	// Sort by FCntUp to optimize the number of MIC checks
	sort.Sort(ByFCntUp(getDevicesResp.Results))

	// Find AppEUI/DevEUI through MIC check
	var device *pb_lorawan.Device
	var micChecks int
	originalFCnt := macPayload.FHDR.FCnt
	for _, candidate := range getDevicesResp.Results {
		nwkSKey := lorawan.AES128Key(*candidate.NwkSKey)

		// First check with the 16 bit counter
		micChecks++
		ok, err = phyPayload.ValidateMIC(nwkSKey)
		if err != nil {
			return err
		}
		if ok {
			device = candidate
			break
		}

		if fullFCnt := fcnt.GetFull(candidate.FCntUp, uint16(originalFCnt)); fullFCnt != originalFCnt && candidate.Uses32BitFCnt {
			macPayload.FHDR.FCnt = fullFCnt

			// Then check again with the 32 bit counter
			if macPayload.FHDR.FCnt != originalFCnt {
				micChecks++
				ok, err = phyPayload.ValidateMIC(nwkSKey)
				if err != nil {
					return err
				}
				if ok {
					device = candidate
					break
				}
			}

			macPayload.FHDR.FCnt = originalFCnt
		}
	}
	if device == nil {
		return errors.NewErrNotFound("device that validates MIC")
	}

	ctx = ctx.WithFields(ttnlog.Fields{
		"MICChecks": micChecks,
		"DevEUI":    device.DevEui,
		"AppEUI":    device.AppEui,
		"AppID":     device.AppId,
		"DevID":     device.DevId,
		"FCnt":      originalFCnt,
	})
	deduplicatedUplink.DevEui = device.DevEui
	deduplicatedUplink.AppEui = device.AppEui
	deduplicatedUplink.AppId = device.AppId
	deduplicatedUplink.DevId = device.DevId
	deduplicatedUplink.Trace = deduplicatedUplink.Trace.WithEvent(trace.CheckMICEvent, "mic checks", micChecks)
	if macPayload.FHDR.FCnt != originalFCnt {
		ctx = ctx.WithField("RealFCnt", macPayload.FHDR.FCnt)
	}

	switch {
	case macPayload.FHDR.FCnt > device.FCntUp && macPayload.FHDR.FCnt-device.FCntUp <= maxFCntGap:
		// FCnt higher than latest and within max FCnt gap (normal case)
	case device.DisableFCntCheck:
		// FCnt Check disabled. Rely on MIC check only
	case device.FCntUp == 0:
		// FCntUp is reset. We don't know where the device will start sending.
	case macPayload.FHDR.FCnt == device.FCntUp:
		if phyPayload.MHDR.MType == lorawan.ConfirmedDataUp {
			// Retry of confirmed uplink
			break
		}
		fallthrough
	case macPayload.FHDR.FCnt <= device.FCntUp:
		return errors.NewErrInvalidArgument("FCnt", "not high enough")
	case macPayload.FHDR.FCnt-device.FCntUp > maxFCntGap:
		return errors.NewErrInvalidArgument("FCnt", "too high")
	default:
		return errors.NewErrInternal("FCnt check failed")
	}

	// Add FCnt to Metadata (because it's not marshaled in lorawan payload)
	deduplicatedUplink.ProtocolMetadata.GetLorawan().FCnt = macPayload.FHDR.FCnt

	// Collect GatewayMetadata and DownlinkOptions
	var downlinkOptions []*pb.DownlinkOption
	for _, duplicate := range duplicates {
		deduplicatedUplink.GatewayMetadata = append(deduplicatedUplink.GatewayMetadata, duplicate.GatewayMetadata)
		downlinkOptions = append(downlinkOptions, duplicate.DownlinkOptions...)
	}

	// Select best DownlinkOption
	if len(downlinkOptions) > 0 {
		deduplicatedUplink.ResponseTemplate = &pb.DownlinkMessage{
			DevEui:         device.DevEui,
			AppEui:         device.AppEui,
			AppId:          device.AppId,
			DevId:          device.DevId,
			DownlinkOption: selectBestDownlink(downlinkOptions),
		}
	}

	// Pass Uplink through NS
	deduplicatedUplink, err = b.ns.Uplink(b.Component.GetContext(b.nsToken), deduplicatedUplink)
	if err != nil {
		return errors.Wrap(errors.FromGRPCError(err), "NetworkServer did not handle uplink")
	}

	var announcements []*pb_discovery.Announcement
	announcements, err = b.Discovery.GetAllHandlersForAppID(device.AppId)
	if err != nil {
		return err
	}
	if len(announcements) == 0 {
		return errors.NewErrNotFound(fmt.Sprintf("Handler for AppID %s", device.AppId))
	}
	if len(announcements) > 1 {
		return errors.NewErrInternal(fmt.Sprintf("Multiple Handlers for AppID %s", device.AppId))
	}

	var handler chan<- *pb.DeduplicatedUplinkMessage
	handler, err = b.getHandlerUplink(announcements[0].Id)
	if err != nil {
		return err
	}

	deduplicatedUplink.Trace = deduplicatedUplink.Trace.WithEvent(trace.ForwardEvent,
		"handler", announcements[0].Id,
	)

	handler <- deduplicatedUplink

	return nil
}

func (b *broker) deduplicateUplink(duplicate *pb.UplinkMessage) (uplinks []*pb.UplinkMessage) {
	sum := md5.Sum(duplicate.Payload)
	key := hex.EncodeToString(sum[:])
	list := b.uplinkDeduplicator.Deduplicate(key, duplicate)
	if len(list) == 0 {
		return
	}
	for _, duplicate := range list {
		uplinks = append(uplinks, duplicate.(*pb.UplinkMessage))
	}
	return
}

func selectBestDownlink(options []*pb.DownlinkOption) *pb.DownlinkOption {
	sort.Sort(ByScore(options))
	return options[0]
}

// ByFCntUp implements sort.Interface for []*pb_lorawan.Device based on FCnt
type ByFCntUp []*pb_lorawan.Device

func (a ByFCntUp) Len() int      { return len(a) }
func (a ByFCntUp) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByFCntUp) Less(i, j int) bool {
	// Devices that disable the FCnt check have low priority
	if a[i].DisableFCntCheck {
		return 2*int(a[i].FCntUp)+100 < int(a[j].FCntUp)
	}
	return int(a[i].FCntUp) < int(a[j].FCntUp)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/moby/hyperkit/go"
	"github.com/satori/go.uuid"
	log "github.com/sirupsen/logrus"
)

const (
	hyperkitNetworkingNone         string = "none"
	hyperkitNetworkingDockerForMac        = "docker-for-mac"
	hyperkitNetworkingVPNKit              = "vpnkit"
	hyperkitNetworkingVMNet               = "vmnet"
	hyperkitNetworkingDefault             = hyperkitNetworkingDockerForMac
)

// Process the run arguments and execute run
func runHyperKit(args []string) {
	flags := flag.NewFlagSet("hyperkit", flag.ExitOnError)
	invoked := filepath.Base(os.Args[0])
	flags.Usage = func() {
		fmt.Printf("USAGE: %s run hyperkit [options] prefix\n\n", invoked)
		fmt.Printf("'prefix' specifies the path to the VM image.\n")
		fmt.Printf("\n")
		fmt.Printf("Options:\n")
		flags.PrintDefaults()
	}
	hyperkitPath := flags.String("hyperkit", "", "Path to hyperkit binary (if not in default location)")
	cpus := flags.Int("cpus", 1, "Number of CPUs")
	mem := flags.Int("mem", 1024, "Amount of memory in MB")
	var disks Disks
	flags.Var(&disks, "disk", "Disk config. [file=]path[,size=1G]")
	data := flags.String("data", "", "Metadata to pass to VM (either a path to a file or a string)")
	ipStr := flags.String("ip", "", "IP address for the VM")
	state := flags.String("state", "", "Path to directory to keep VM state in")
	vsockports := flags.String("vsock-ports", "", "List of vsock ports to forward from the guest on startup (comma separated). A unix domain socket for each port will be created in the state directory")
	networking := flags.String("networking", hyperkitNetworkingDefault, "Networking mode. Valid options are 'default', 'docker-for-mac', 'vpnkit[,socket-path]', 'vmnet' and 'none'. 'docker-for-mac' connects to the network used by Docker for Mac. 'vpnkit' connects to the VPNKit socket specified. If socket-path is omitted a new VPNKit instance will be started and 'vpnkit_eth.sock' will be created in the state directory. 'vmnet' uses the Apple vmnet framework, requires root/sudo. 'none' disables networking.`")

	// Boot type; we try to determine automatically
	uefiBoot := flags.Bool("uefi", false, "Use UEFI boot")
	isoBoot := flags.Bool("iso", false, "Boot image is an ISO")
	kernelBoot := flags.Bool("kernel", false, "Boot image is kernel+initrd+cmdline 'path'-kernel/-initrd/-cmdline")

	// Paths and settings for UEFI firware
	// Note, the default uses the firmware shipped with Docker for Mac
	fw := flags.String("fw", "/Applications/Docker.app/Contents/Resources/uefi/UEFI.fd", "Path to OVMF firmware for UEFI boot")

	if err := flags.Parse(args); err != nil {
		log.Fatal("Unable to parse args")
	}
	remArgs := flags.Args()
	if len(remArgs) == 0 {
		fmt.Println("Please specify the prefix to the image to boot\n")
		flags.Usage()
		os.Exit(1)
	}
	path := remArgs[0]
	prefix := path

	info, err := os.Stat(path)
	stat := err == nil

	/// ignore a directory
	if stat && info.Mode().IsDir() {
		stat = false
	}

	_, err = os.Stat(path + "-kernel")
	statKernel := err == nil

	var isoPaths []string

	// try to autodetect boot type if not specified
	// if the path does not exist, and the kernel does, must be trying to do a kernel boot
	// if the path does exist and ends in ISO, must be trying ISO boot
	if !stat && statKernel && !*isoBoot {
		*kernelBoot = true
	} else if stat && strings.HasSuffix(path, ".iso") && !*kernelBoot {
		*isoBoot = true
	}

	switch {
	case *kernelBoot && *isoBoot:
		log.Fatalf("Cannot specify both kernel and ISO boot together")
	case *kernelBoot:
		if !statKernel {
			log.Fatalf("Cannot find kernel file (%s): %v", path+"-kernel", err)
		}
		_, err = os.Stat(path + "-initrd.img")
		statInitrd := err == nil
		if !statInitrd {
			log.Fatalf("Cannot find initrd file (%s): %v", path+"-initrd.img", err)
		}
	case *isoBoot:
		if !stat {
			log.Fatalf("Cannot find ISO to boot")
		}
		prefix = strings.TrimSuffix(path, ".iso")
		// hyperkit only supports UEFI ISO boot at present
		if !*uefiBoot {
			log.Fatalf("Hyperkit requires --uefi to be set to boot an ISO")
		}
		isoPaths = append(isoPaths, path)
	}

	if *uefiBoot {
		_, err := os.Stat(*fw)
		if err != nil {
			log.Fatalf("Cannot open UEFI firmware file (%s): %v", *fw, err)
		}
	}

	if *state == "" {
		*state = prefix + "-state"
	}
	if err := os.MkdirAll(*state, 0755); err != nil {
		log.Fatalf("Could not create state directory: %v", err)
	}

	if *data != "" {
		var d []byte
		if _, err := os.Stat(*data); os.IsNotExist(err) {
			d = []byte(*data)
		} else {
			d, err = ioutil.ReadFile(*data)
			if err != nil {
				log.Fatalf("Cannot read user data: %v", err)
			}
		}
		isoPath := filepath.Join(*state, "data.iso")
		if err := WriteMetadataISO(isoPath, d); err != nil {
			log.Fatalf("Cannot write user data ISO: %v", err)
		}
		isoPaths = append(isoPaths, isoPath)
	}

	vpnKitKey := ""
	if *ipStr != "" {
		// If an IP address was requested construct a "special" UUID
		// for the VM.
		if ip := net.ParseIP(*ipStr); len(ip) > 0 {
			uuid := make([]byte, 16)
			uuid[12] = ip.To4()[0]
			uuid[13] = ip.To4()[1]
			uuid[14] = ip.To4()[2]
			uuid[15] = ip.To4()[3]
			vpnKitKey = fmt.Sprintf("%x-%x-%x-%x-%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:])
		}
	}

	// Generate new UUID, otherwise /sys/class/dmi/id/product_uuid is identical on all VMs
	vmUUID := uuid.NewV4().String()

	// Run
	var cmdline []byte
	if *kernelBoot {
		cmdline, err = ioutil.ReadFile(prefix + "-cmdline")
		if err != nil {
			log.Fatalf("Cannot open cmdline file: %v", err)
		}
	}

	// Create new HyperKit instance (w/o networking for now)
	h, err := hyperkit.New(*hyperkitPath, "", *state)
	if err != nil {
		log.Fatalln("Error creating hyperkit: ", err)
	}

	for i, d := range disks {
		id := ""
		if i != 0 {
			id = strconv.Itoa(i)
		}
		if d.Size != 0 && d.Path == "" {
			d.Path = filepath.Join(*state, "disk"+id+".img")
		}
		if d.Path == "" {
			log.Fatalf("disk specified with no size or name")
		}
		hd := hyperkit.DiskConfig{Path: d.Path, Size: d.Size}
		h.Disks = append(h.Disks, hd)
	}

	if h.VSockPorts, err = stringToIntArray(*vsockports, ","); err != nil {
		log.Fatalln("Unable to parse vsock-ports: ", err)
	}

	// Select network mode
	var vpnKitProcess *os.Process
	if *networking == "" || *networking == "default" {
		dflt := hyperkitNetworkingDefault
		networking = &dflt
	}
	netMode := strings.SplitN(*networking, ",", 2)

	switch netMode[0] {
	case hyperkitNetworkingDockerForMac:
		h.VPNKitSock = filepath.Join(os.Getenv("HOME"), "Library/Containers/com.docker.docker/Data/s50")
	case hyperkitNetworkingVPNKit:
		if len(netMode) > 1 {
			// Socket path specified, try to use existing VPNKit instance
			h.VPNKitSock = netMode[1]
		} else {
			// Start new VPNKit instance
			h.VPNKitSock = filepath.Join(*state, "vpnkit_eth.sock")
			vpnKitPortSocket := filepath.Join(*state, "vpnkit_port.sock")
			vsockSocket := filepath.Join(*state, "connect")
			vpnKitProcess, err = launchVPNKit(h.VPNKitSock, vsockSocket, vpnKitPortSocket)
			if err != nil {
				log.Fatalln("Unable to start vpnkit: ", err)
			}
			defer func() {
				if vpnKitProcess != nil {
					err := vpnKitProcess.Kill()
					if err != nil {
						log.Println(err)
					}
				}
			}()
			// The guest will use this 9P mount to configure which ports to forward
			h.Sockets9P = []hyperkit.Socket9P{{Path: vpnKitPortSocket, Tag: "port"}}
			// VSOCK port 62373 is used to pass traffic from host->guest
			h.VSockPorts = append(h.VSockPorts, 62373)
		}
	case hyperkitNetworkingVMNet:
		h.VPNKitSock = ""
		h.VMNet = true
	case hyperkitNetworkingNone:
		h.VPNKitSock = ""
	default:
		log.Fatalf("Invalid networking mode: %s", netMode[0])
	}

	if *kernelBoot {
		h.Kernel = prefix + "-kernel"
		h.Initrd = prefix + "-initrd.img"
	} else {
		h.Bootrom = *fw
	}
	h.VPNKitKey = vpnKitKey
	h.UUID = vmUUID
	h.ISOImages = isoPaths
	h.VSock = true
	h.CPUs = *cpus
	h.Memory = *mem

	err = h.Run(string(cmdline))
	if err != nil {
		log.Fatalf("Cannot run hyperkit: %v", err)
	}
}

// createListenSocket creates a new unix domain socket and returns the open file
func createListenSocket(path string) (*os.File, error) {
	os.Remove(path)
	conn, err := net.ListenUnix("unix", &net.UnixAddr{Name: path, Net: "unix"})
	if err != nil {
		return nil, fmt.Errorf("unable to create socket: %v", err)
	}
	f, err := conn.File()
	if err != nil {
		return nil, err
	}
	return f, nil
}

// launchVPNKit starts a new instance of VPNKit. Ethernet socket and port socket
// will be created and passed to VPNKit. The VSOCK socket should be created
// by HyperKit when it starts.
func launchVPNKit(etherSock string, vsockSock string, portSock string) (*os.Process, error) {
	var err error

	vpnKitPath, err := exec.LookPath("vpnkit")
	if err != nil {
		return nil, fmt.Errorf("Unable to find vpnkit binary")
	}

	etherFile, err := createListenSocket(etherSock)
	if err != nil {
		return nil, err
	}

	portFile, err := createListenSocket(portSock)
	if err != nil {
		return nil, err
	}

	cmd := exec.Command(vpnKitPath,
		"--ethernet", "fd:3",
		"--vsock-path", vsockSock,
		"--port", "fd:4")

	cmd.ExtraFiles = append(cmd.ExtraFiles, etherFile)
	cmd.ExtraFiles = append(cmd.ExtraFiles, portFile)

	cmd.Env = os.Environ() // pass env for DEBUG

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

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

	go cmd.Wait() // run in background

	return cmd.Process, nil
}
<commit_msg>Add missing case when running non existent file in hyperkit<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/moby/hyperkit/go"
	"github.com/satori/go.uuid"
	log "github.com/sirupsen/logrus"
)

const (
	hyperkitNetworkingNone         string = "none"
	hyperkitNetworkingDockerForMac        = "docker-for-mac"
	hyperkitNetworkingVPNKit              = "vpnkit"
	hyperkitNetworkingVMNet               = "vmnet"
	hyperkitNetworkingDefault             = hyperkitNetworkingDockerForMac
)

// Process the run arguments and execute run
func runHyperKit(args []string) {
	flags := flag.NewFlagSet("hyperkit", flag.ExitOnError)
	invoked := filepath.Base(os.Args[0])
	flags.Usage = func() {
		fmt.Printf("USAGE: %s run hyperkit [options] prefix\n\n", invoked)
		fmt.Printf("'prefix' specifies the path to the VM image.\n")
		fmt.Printf("\n")
		fmt.Printf("Options:\n")
		flags.PrintDefaults()
	}
	hyperkitPath := flags.String("hyperkit", "", "Path to hyperkit binary (if not in default location)")
	cpus := flags.Int("cpus", 1, "Number of CPUs")
	mem := flags.Int("mem", 1024, "Amount of memory in MB")
	var disks Disks
	flags.Var(&disks, "disk", "Disk config. [file=]path[,size=1G]")
	data := flags.String("data", "", "Metadata to pass to VM (either a path to a file or a string)")
	ipStr := flags.String("ip", "", "IP address for the VM")
	state := flags.String("state", "", "Path to directory to keep VM state in")
	vsockports := flags.String("vsock-ports", "", "List of vsock ports to forward from the guest on startup (comma separated). A unix domain socket for each port will be created in the state directory")
	networking := flags.String("networking", hyperkitNetworkingDefault, "Networking mode. Valid options are 'default', 'docker-for-mac', 'vpnkit[,socket-path]', 'vmnet' and 'none'. 'docker-for-mac' connects to the network used by Docker for Mac. 'vpnkit' connects to the VPNKit socket specified. If socket-path is omitted a new VPNKit instance will be started and 'vpnkit_eth.sock' will be created in the state directory. 'vmnet' uses the Apple vmnet framework, requires root/sudo. 'none' disables networking.`")

	// Boot type; we try to determine automatically
	uefiBoot := flags.Bool("uefi", false, "Use UEFI boot")
	isoBoot := flags.Bool("iso", false, "Boot image is an ISO")
	kernelBoot := flags.Bool("kernel", false, "Boot image is kernel+initrd+cmdline 'path'-kernel/-initrd/-cmdline")

	// Paths and settings for UEFI firware
	// Note, the default uses the firmware shipped with Docker for Mac
	fw := flags.String("fw", "/Applications/Docker.app/Contents/Resources/uefi/UEFI.fd", "Path to OVMF firmware for UEFI boot")

	if err := flags.Parse(args); err != nil {
		log.Fatal("Unable to parse args")
	}
	remArgs := flags.Args()
	if len(remArgs) == 0 {
		fmt.Println("Please specify the prefix to the image to boot\n")
		flags.Usage()
		os.Exit(1)
	}
	path := remArgs[0]
	prefix := path

	info, err := os.Stat(path)
	stat := err == nil

	// ignore a directory
	if stat && info.Mode().IsDir() {
		stat = false
	}

	_, err = os.Stat(path + "-kernel")
	statKernel := err == nil

	var isoPaths []string

	// try to autodetect boot type if not specified
	// if the path does not exist, and the kernel does, must be trying to do a kernel boot
	// if the path does exist and ends in ISO, must be trying ISO boot
	if !stat && statKernel && !*isoBoot {
		*kernelBoot = true
	} else if stat && strings.HasSuffix(path, ".iso") && !*kernelBoot {
		*isoBoot = true
	}

	switch {
	case *kernelBoot && *isoBoot:
		log.Fatalf("Cannot specify both kernel and ISO boot together")
	case *kernelBoot:
		if !statKernel {
			log.Fatalf("Cannot find kernel file (%s): %v", path+"-kernel", err)
		}
		_, err = os.Stat(path + "-initrd.img")
		statInitrd := err == nil
		if !statInitrd {
			log.Fatalf("Cannot find initrd file (%s): %v", path+"-initrd.img", err)
		}
	case *isoBoot:
		if !stat {
			log.Fatalf("Cannot find ISO to boot")
		}
		prefix = strings.TrimSuffix(path, ".iso")
		// hyperkit only supports UEFI ISO boot at present
		if !*uefiBoot {
			log.Fatalf("Hyperkit requires --uefi to be set to boot an ISO")
		}
		isoPaths = append(isoPaths, path)
	default:
		if !stat {
			log.Fatalf("Cannot find file %s to boot", path)
		}
		log.Fatalf("Unrecognised boot type, please specify on command line")
	}

	if *uefiBoot {
		_, err := os.Stat(*fw)
		if err != nil {
			log.Fatalf("Cannot open UEFI firmware file (%s): %v", *fw, err)
		}
	}

	if *state == "" {
		*state = prefix + "-state"
	}
	if err := os.MkdirAll(*state, 0755); err != nil {
		log.Fatalf("Could not create state directory: %v", err)
	}

	if *data != "" {
		var d []byte
		if _, err := os.Stat(*data); os.IsNotExist(err) {
			d = []byte(*data)
		} else {
			d, err = ioutil.ReadFile(*data)
			if err != nil {
				log.Fatalf("Cannot read user data: %v", err)
			}
		}
		isoPath := filepath.Join(*state, "data.iso")
		if err := WriteMetadataISO(isoPath, d); err != nil {
			log.Fatalf("Cannot write user data ISO: %v", err)
		}
		isoPaths = append(isoPaths, isoPath)
	}

	vpnKitKey := ""
	if *ipStr != "" {
		// If an IP address was requested construct a "special" UUID
		// for the VM.
		if ip := net.ParseIP(*ipStr); len(ip) > 0 {
			uuid := make([]byte, 16)
			uuid[12] = ip.To4()[0]
			uuid[13] = ip.To4()[1]
			uuid[14] = ip.To4()[2]
			uuid[15] = ip.To4()[3]
			vpnKitKey = fmt.Sprintf("%x-%x-%x-%x-%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:])
		}
	}

	// Generate new UUID, otherwise /sys/class/dmi/id/product_uuid is identical on all VMs
	vmUUID := uuid.NewV4().String()

	// Run
	var cmdline []byte
	if *kernelBoot {
		cmdline, err = ioutil.ReadFile(prefix + "-cmdline")
		if err != nil {
			log.Fatalf("Cannot open cmdline file: %v", err)
		}
	}

	// Create new HyperKit instance (w/o networking for now)
	h, err := hyperkit.New(*hyperkitPath, "", *state)
	if err != nil {
		log.Fatalln("Error creating hyperkit: ", err)
	}

	for i, d := range disks {
		id := ""
		if i != 0 {
			id = strconv.Itoa(i)
		}
		if d.Size != 0 && d.Path == "" {
			d.Path = filepath.Join(*state, "disk"+id+".img")
		}
		if d.Path == "" {
			log.Fatalf("disk specified with no size or name")
		}
		hd := hyperkit.DiskConfig{Path: d.Path, Size: d.Size}
		h.Disks = append(h.Disks, hd)
	}

	if h.VSockPorts, err = stringToIntArray(*vsockports, ","); err != nil {
		log.Fatalln("Unable to parse vsock-ports: ", err)
	}

	// Select network mode
	var vpnKitProcess *os.Process
	if *networking == "" || *networking == "default" {
		dflt := hyperkitNetworkingDefault
		networking = &dflt
	}
	netMode := strings.SplitN(*networking, ",", 2)

	switch netMode[0] {
	case hyperkitNetworkingDockerForMac:
		h.VPNKitSock = filepath.Join(os.Getenv("HOME"), "Library/Containers/com.docker.docker/Data/s50")
	case hyperkitNetworkingVPNKit:
		if len(netMode) > 1 {
			// Socket path specified, try to use existing VPNKit instance
			h.VPNKitSock = netMode[1]
		} else {
			// Start new VPNKit instance
			h.VPNKitSock = filepath.Join(*state, "vpnkit_eth.sock")
			vpnKitPortSocket := filepath.Join(*state, "vpnkit_port.sock")
			vsockSocket := filepath.Join(*state, "connect")
			vpnKitProcess, err = launchVPNKit(h.VPNKitSock, vsockSocket, vpnKitPortSocket)
			if err != nil {
				log.Fatalln("Unable to start vpnkit: ", err)
			}
			defer func() {
				if vpnKitProcess != nil {
					err := vpnKitProcess.Kill()
					if err != nil {
						log.Println(err)
					}
				}
			}()
			// The guest will use this 9P mount to configure which ports to forward
			h.Sockets9P = []hyperkit.Socket9P{{Path: vpnKitPortSocket, Tag: "port"}}
			// VSOCK port 62373 is used to pass traffic from host->guest
			h.VSockPorts = append(h.VSockPorts, 62373)
		}
	case hyperkitNetworkingVMNet:
		h.VPNKitSock = ""
		h.VMNet = true
	case hyperkitNetworkingNone:
		h.VPNKitSock = ""
	default:
		log.Fatalf("Invalid networking mode: %s", netMode[0])
	}

	if *kernelBoot {
		h.Kernel = prefix + "-kernel"
		h.Initrd = prefix + "-initrd.img"
	} else {
		h.Bootrom = *fw
	}
	h.VPNKitKey = vpnKitKey
	h.UUID = vmUUID
	h.ISOImages = isoPaths
	h.VSock = true
	h.CPUs = *cpus
	h.Memory = *mem

	err = h.Run(string(cmdline))
	if err != nil {
		log.Fatalf("Cannot run hyperkit: %v", err)
	}
}

// createListenSocket creates a new unix domain socket and returns the open file
func createListenSocket(path string) (*os.File, error) {
	os.Remove(path)
	conn, err := net.ListenUnix("unix", &net.UnixAddr{Name: path, Net: "unix"})
	if err != nil {
		return nil, fmt.Errorf("unable to create socket: %v", err)
	}
	f, err := conn.File()
	if err != nil {
		return nil, err
	}
	return f, nil
}

// launchVPNKit starts a new instance of VPNKit. Ethernet socket and port socket
// will be created and passed to VPNKit. The VSOCK socket should be created
// by HyperKit when it starts.
func launchVPNKit(etherSock string, vsockSock string, portSock string) (*os.Process, error) {
	var err error

	vpnKitPath, err := exec.LookPath("vpnkit")
	if err != nil {
		return nil, fmt.Errorf("Unable to find vpnkit binary")
	}

	etherFile, err := createListenSocket(etherSock)
	if err != nil {
		return nil, err
	}

	portFile, err := createListenSocket(portSock)
	if err != nil {
		return nil, err
	}

	cmd := exec.Command(vpnKitPath,
		"--ethernet", "fd:3",
		"--vsock-path", vsockSock,
		"--port", "fd:4")

	cmd.ExtraFiles = append(cmd.ExtraFiles, etherFile)
	cmd.ExtraFiles = append(cmd.ExtraFiles, portFile)

	cmd.Env = os.Environ() // pass env for DEBUG

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

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

	go cmd.Wait() // run in background

	return cmd.Process, nil
}
<|endoftext|>


<commit_before>// api.go: Queries db, exposes JSON API:
// 1. GET /monkeys.json: lists all entries
// 2. POST /monkeys.json: create entity of the type
// 3. GET /monkeys/[enc id].json retrieves a specific entity
// 4. PUT /monkeys/[enc id].json updates a specific entity
// 5. DELETE /monkey/[enc id].json: deletes that entity
package api

import (
	"database/sql"

	_ "github.com/go-sql-driver/mysql"

	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/gorilla/mux"
)

var (
	dbAddr = flag.String("db_addr", "", "If set, TCP host for the DB. If not set, address is read from etcd")
	// Note that we always bind to the same port inside the container; the
	// .service file can map it to any external port that's desired
	// based on which stage we're running.
	bindAddr                        = ":9100"
	stage                           = ""      // prod|staging|testN|dev|unittest
	maxRequestSize            int64 = 1048576 // largest allowed request, in bytes
	statusUnprocessableEntity       = 422
)

// Monkey is an entity we deal with in the API.
type (
	Monkey struct {
		Id        int       `json:"id"`
		Name      string    `json:"name"`
		Birthdate time.Time `json:"birthdate"`
	}
	// Monkeys are a collection of monkey.
	Monkeys []*Monkey
	// MonkeyAPI defines the interface on how we interact with monkeys.
	MonkeyAPI interface {
		GetMonkey(int) (*Monkey, error)
		GetMonkeys() (*Monkeys, error)
		AddMonkey(Monkey) error
		// TODO: add UpdateMonkey, DeleteMonkey.
	}
)

// String returns a human-readable description of the monkey.
func (m Monkey) String() string {
	return fmt.Sprintf("%s (%d) was born on %v", m.Name, m.Id, m.Birthdate.Format("Mon, 02 Jan 2006"))
}

// String returns a human-readable description of the monkeys.
func (ms Monkeys) String() string {
	r := ""
	for i, m := range ms {
		if i > 0 {
			r += ", "
		}
		r += m.String()
	}
	return r
}

// Serve blocks forever, serving the API on bindAddr.
func Serve() {
	flag.Parse()
	stage = os.Getenv("STAGE")
	if stage == "" {
		log.Fatalf("FATAL: no STAGE set as environment variable")
	}
	fmt.Printf("api layer for %s stage starting..\n", stage)
	log.Fatal(http.ListenAndServe(bindAddr, newRouter(apiHandler{jsonAPI{}})))
}

func getDB() (*sql.DB, error) {
	// TODO: read DB info from etcd.
	if *dbAddr == "" {
		return nil, fmt.Errorf("TODO: implement reading DB from etcd")
	} else {
		sqlSource := fmt.Sprintf(
			"%s:%s@tcp(%s)/%s",
			// Note: Obviously not secure, in production we'd have an encrypted
			// config.
			"dbuser",   // user
			"dbsecret", // pass
			*dbAddr,    // db address + port
			"monkeydb") // db
		db, err := sql.Open("mysql", sqlSource)
		if err != nil {
			return nil, err
		}
		return db, db.Ping()
	}
}

type apiHandler struct {
	api MonkeyAPI
}

type jsonAPI struct{}

func (api jsonAPI) GetMonkey(id int) (*Monkey, error) {
	// TODO: query MySQL db here.
	return nil, fmt.Errorf("TODO: implement getMonkey")
}

func (api jsonAPI) GetMonkeys() (*Monkeys, error) {
	// TODO: get DB from /services/db/%stage% from etcd, get "host" and "port" keys.
	// TODO: query MySQL db here.
	db, err := getDB()
	if err != nil {
		return nil, fmt.Errorf("failed to contact DB: %v", err)
	}
	rows, err := db.Query(`
      SELECT monkeyId, monkeyName, birthDate
      FROM monkeys
      LIMIT 1000;`)
	if err != nil {
		return nil, fmt.Errorf("failed to query DB: %v", err)
	}
	defer rows.Close()
	monkeys := Monkeys{}
	for rows.Next() {
		id := 0
		name := ""
		sec := int64(0)

		if err = rows.Scan(&id, &name, &sec); err != nil {
			return nil, fmt.Errorf("failed to scan: %v", err)
		}
		// Note: If this was exposed to users, we'd need to display it in
		// their own timezone (explicitly selected).
		birthdate := time.Unix(sec, 0).UTC()
		monkeys = append(monkeys, &Monkey{id, name, birthdate})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("row error: %v", err)
	}
	return &monkeys, nil
}

func (api jsonAPI) AddMonkey(m Monkey) error {
	// TODO: insert data into MySQL db here.
	return fmt.Errorf("TODO: implement addMonkey")
}

// newRouter returns a new HTTP router for the endpoints of the API.
func newRouter(h apiHandler) *mux.Router {
	r := mux.NewRouter().StrictSlash(true)
	r.HandleFunc("/monkeys", h.getMonkeys).Methods("GET")
	r.HandleFunc("/monkeys", h.createMonkey).Methods("POST")
	r.HandleFunc("/monkeys/{key}", h.getMonkey).Methods("GET")
	r.HandleFunc("/monkey/{key}", h.updateMonkey).Methods("PUT")
	r.HandleFunc("/monkeys/{key}", h.deleteMonkey).Methods("DELETE")
	return r
}

// getMonkey fetches all monkeys.
func (h apiHandler) getMonkeys(w http.ResponseWriter, r *http.Request) {
	m, err := h.api.GetMonkeys()
	if err != nil {
		log.Printf("failed to fetch monkeys: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
	err = json.NewEncoder(w).Encode(m)
	if err != nil {
		log.Printf("failed to encode monkeys: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
}

// createMonkey creates a new monkey.
func (h apiHandler) createMonkey(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(io.LimitReader(r.Body, maxRequestSize))
	if err != nil {
		log.Printf("failed to read monkey: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
	if err := r.Body.Close(); err != nil {
		log.Printf("failed to close request: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
	m := Monkey{}
	if err := json.Unmarshal(body, &m); err != nil {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(statusUnprocessableEntity)
		if err := json.NewEncoder(w).Encode(err); err != nil {
			log.Printf("failed to write encoding error: %v", err)
			http.Error(w, "Internal server error.", http.StatusInternalServerError)
			return
		}
		return
	}
	err = h.api.AddMonkey(m)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(m); err != nil {
		log.Printf("failed to write encoding error: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
}

// getMonkey fetches a specific monkey.
func (h apiHandler) getMonkey(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")

	vars := mux.Vars(r)
	// Note: In a production environment, we likely should expose hashes
	// of database ids, not the raw ids.
	id, err := strconv.Atoi(vars["key"])
	if err != nil {
		log.Printf("bad monkey id %q: %v", vars["key"], err)
		http.Error(w, fmt.Sprintf("No such id %q.", vars["key"]), http.StatusBadRequest)
		return
	}

	m, err := h.api.GetMonkey(id)
	if err != nil {
		log.Printf("failed to fetch monkey: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
	if m == nil {
		log.Printf("no monkey with id %d\n", id)
		w.WriteHeader(http.StatusNotFound)
		return
	}
	err = json.NewEncoder(w).Encode(m)
	if err != nil {
		log.Printf("failed to encode monkey: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
}

// updateMonkey updates a monkey.
func (h apiHandler) updateMonkey(w http.ResponseWriter, r *http.Request) {
	msg := "TODO: implement updateMonkey\n"
	log.Printf(msg)
	http.Error(w, msg, http.StatusInternalServerError)
}

// deleteMonkey deletes a monkey.
func (h apiHandler) deleteMonkey(w http.ResponseWriter, r *http.Request) {
	msg := "TODO: implement deleteMonkey\n"
	log.Printf(msg)
	http.Error(w, msg, http.StatusInternalServerError)
}
<commit_msg>implements api.GetMonkey, also reading DB addr from etcd<commit_after>// api.go: Queries db, exposes JSON API:
// 1. GET /monkeys.json: lists all entries
// 2. POST /monkeys.json: create entity of the type
// 3. GET /monkeys/[enc id].json retrieves a specific entity
// 4. PUT /monkeys/[enc id].json updates a specific entity
// 5. DELETE /monkey/[enc id].json: deletes that entity
package api

import (
	"database/sql"

	_ "github.com/go-sql-driver/mysql"
	"github.com/golang/glog"

	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/coreos/go-etcd/etcd"
	"github.com/gorilla/mux"
)

var (
	dbAddr = flag.String("db_addr", "", "If set, TCP host for the DB. If not set, address is read from etcd")
	// Note that we always bind to the same port inside the container; the
	// .service file can map it to any external port that's desired
	// based on which stage we're running.
	bindAddr                        = ":9100"
	stage                           = ""      // prod|staging|testN|dev|unittest
	maxRequestSize            int64 = 1048576 // largest allowed request, in bytes
	statusUnprocessableEntity       = 422
)

// Monkey is an entity we deal with in the API.
type (
	Monkey struct {
		Id        int       `json:"id"`
		Name      string    `json:"name"`
		Birthdate time.Time `json:"birthdate"`
	}
	// Monkeys are a collection of monkey.
	Monkeys []*Monkey
	// MonkeyAPI defines the interface on how we interact with monkeys.
	MonkeyAPI interface {
		GetMonkey(int) (*Monkey, error)
		GetMonkeys() (*Monkeys, error)
		AddMonkey(Monkey) error
		// TODO: add UpdateMonkey, DeleteMonkey.
	}
)

// String returns a human-readable description of the monkey.
func (m Monkey) String() string {
	return fmt.Sprintf("%s (%d) was born on %v", m.Name, m.Id, m.Birthdate.Format("Mon, 02 Jan 2006"))
}

// String returns a human-readable description of the monkeys.
func (ms Monkeys) String() string {
	r := ""
	for i, m := range ms {
		if i > 0 {
			r += ", "
		}
		r += m.String()
	}
	return r
}

// Serve blocks forever, serving the API on bindAddr.
func Serve() {
	flag.Parse()
	stage = os.Getenv("STAGE")
	if stage == "" {
		log.Fatalf("FATAL: no STAGE set as environment variable")
	}
	fmt.Printf("api layer for %s stage starting..\n", stage)
	log.Fatal(http.ListenAndServe(bindAddr, newRouter(apiHandler{jsonAPI{}})))
}

// getEtcdHost returns the Host info from etcd.
func getEtcdHost() (string, error) {
	path := fmt.Sprintf("/services/db/%s", stage)
	c := etcd.NewClient([]string{"http://127.0.0.1:4001"})
	r, err := c.Get(path, false, false)
	if err != nil {
		return "", fmt.Errorf("failed to read %s: %v", path, err)
	}
	v := r.Node.Value
	glog.V(1).Infof("read value %q from %s\n", v, path)
	addr := struct {
		Host    string `json:"host"`
		Port    int    `json:"port"`
		Version string `json:"version"`
	}{}
	err = json.Unmarshal([]byte(v), &addr)
	if err != nil {
		return "", fmt.Errorf("failed to interpret etcd value %q: %v", v, err)
	}
	return fmt.Sprintf("%s:%s", addr.Host, addr.Port), nil
}

func getDB() (*sql.DB, error) {
	addr := *dbAddr
	if *dbAddr == "" {
		var err error
		addr, err = getEtcdHost()
		if err != nil {
			return nil, fmt.Errorf("failed to read from etcd: %v", err)
		}
	}
	sqlSource := fmt.Sprintf(
		"%s:%s@tcp(%s)/%s",
		// Note: Obviously not secure, in production we'd have an encrypted
		// config.
		"dbuser",   // user
		"dbsecret", // pass
		addr,       // db address + port
		"monkeydb") // db
	glog.V(1).Infof("connecting to MySQL at %s..\n", sqlSource)
	db, err := sql.Open("mysql", sqlSource)
	if err != nil {
		return nil, err
	}
	return db, db.Ping()
}

type apiHandler struct {
	api MonkeyAPI
}

type jsonAPI struct{}

func (api jsonAPI) GetMonkey(id int) (*Monkey, error) {
	db, err := getDB()
	if err != nil {
		return nil, fmt.Errorf("failed to contact DB: %v", err)
	}
	row := db.QueryRow(`
      SELECT monkeyName, birthDate
      FROM monkeys
      WHERE monkeyId=?`, id)
	name := ""
	sec := int64(0)
	if err = row.Scan(&name, &sec); err != nil {
		return nil, fmt.Errorf("failed to scan: %v", err)
	}
	// Note: If this was exposed to users, we'd need to display it in
	// their own timezone (explicitly selected).
	birthdate := time.Unix(sec, 0).UTC()
	return &Monkey{id, name, birthdate}, nil
}

func (api jsonAPI) GetMonkeys() (*Monkeys, error) {
	db, err := getDB()
	if err != nil {
		return nil, fmt.Errorf("failed to contact DB: %v", err)
	}
	rows, err := db.Query(`
      SELECT monkeyId, monkeyName, birthDate
      FROM monkeys
      LIMIT 1000;`)
	if err != nil {
		return nil, fmt.Errorf("failed to query DB: %v", err)
	}
	defer rows.Close()
	monkeys := Monkeys{}
	for rows.Next() {
		id := 0
		name := ""
		sec := int64(0)

		if err = rows.Scan(&id, &name, &sec); err != nil {
			return nil, fmt.Errorf("failed to scan: %v", err)
		}
		// Note: If this was exposed to users, we'd need to display it in
		// their own timezone (explicitly selected).
		birthdate := time.Unix(sec, 0).UTC()
		monkeys = append(monkeys, &Monkey{id, name, birthdate})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("row error: %v", err)
	}
	return &monkeys, nil
}

func (api jsonAPI) AddMonkey(m Monkey) error {
	// TODO: insert data into MySQL db here.
	return fmt.Errorf("TODO: implement addMonkey")
}

// newRouter returns a new HTTP router for the endpoints of the API.
func newRouter(h apiHandler) *mux.Router {
	r := mux.NewRouter().StrictSlash(true)
	r.HandleFunc("/monkeys", h.getMonkeys).Methods("GET")
	r.HandleFunc("/monkeys", h.createMonkey).Methods("POST")
	r.HandleFunc("/monkeys/{key}", h.getMonkey).Methods("GET")
	r.HandleFunc("/monkey/{key}", h.updateMonkey).Methods("PUT")
	r.HandleFunc("/monkeys/{key}", h.deleteMonkey).Methods("DELETE")
	return r
}

// getMonkey fetches all monkeys.
func (h apiHandler) getMonkeys(w http.ResponseWriter, r *http.Request) {
	m, err := h.api.GetMonkeys()
	if err != nil {
		glog.Errorf("failed to fetch monkeys: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
	err = json.NewEncoder(w).Encode(m)
	if err != nil {
		glog.Errorf("failed to encode monkeys: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
}

// createMonkey creates a new monkey.
func (h apiHandler) createMonkey(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(io.LimitReader(r.Body, maxRequestSize))
	if err != nil {
		glog.Errorf("failed to read monkey: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
	if err := r.Body.Close(); err != nil {
		glog.Errorf("failed to close request: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
	m := Monkey{}
	if err := json.Unmarshal(body, &m); err != nil {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(statusUnprocessableEntity)
		if err := json.NewEncoder(w).Encode(err); err != nil {
			glog.Errorf("failed to write encoding error: %v", err)
			http.Error(w, "Internal server error.", http.StatusInternalServerError)
			return
		}
		return
	}
	err = h.api.AddMonkey(m)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusCreated)
	if err := json.NewEncoder(w).Encode(m); err != nil {
		glog.Errorf("failed to write encoding error: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
}

// getMonkey fetches a specific monkey.
func (h apiHandler) getMonkey(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")

	vars := mux.Vars(r)
	// Note: In a production environment, we likely should expose hashes
	// of database ids, not the raw ids.
	id, err := strconv.Atoi(vars["key"])
	if err != nil {
		glog.Errorf("bad monkey id %q: %v", vars["key"], err)
		http.Error(w, fmt.Sprintf("No such id %q.", vars["key"]), http.StatusBadRequest)
		return
	}

	m, err := h.api.GetMonkey(id)
	if err != nil {
		glog.Errorf("failed to fetch monkey: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
	if m == nil {
		glog.Errorf("no monkey with id %d\n", id)
		w.WriteHeader(http.StatusNotFound)
		return
	}
	err = json.NewEncoder(w).Encode(m)
	if err != nil {
		glog.Errorf("failed to encode monkey: %v", err)
		http.Error(w, "Internal server error.", http.StatusInternalServerError)
		return
	}
}

// updateMonkey updates a monkey.
func (h apiHandler) updateMonkey(w http.ResponseWriter, r *http.Request) {
	msg := "TODO: implement updateMonkey\n"
	glog.Errorf(msg)
	http.Error(w, msg, http.StatusInternalServerError)
}

// deleteMonkey deletes a monkey.
func (h apiHandler) deleteMonkey(w http.ResponseWriter, r *http.Request) {
	msg := "TODO: implement deleteMonkey\n"
	glog.Errorf(msg)
	http.Error(w, msg, http.StatusInternalServerError)
}
<|endoftext|>


<commit_before>package daemon

import (
	"errors"
	"fmt"
	"github.com/aws/aws-sdk-go/service/ec2"
	"testing"
)

type FakeMetadataFetcher struct {
	FAvailable bool
	Meta       map[string]string
}

func (m FakeMetadataFetcher) Available() bool {
	return m.FAvailable
}

func (m FakeMetadataFetcher) GetMetadata(key string) (string, error) {
	v, ok := m.Meta[key]
	if ok {
		return v, nil
	}
	return v, errors.New(fmt.Sprintf("Key %s unknown"))
}

type FakeRouteTableFetcher struct{}

func (r FakeRouteTableFetcher) GetRouteTables() ([]*ec2.RouteTable, error) {
	return []*ec2.RouteTable{}, nil
}

func getD(a bool) Daemon {
	d := Daemon{
		ConfigFile: "../tests/awsnycast.yaml",
	}
	fakeM := FakeMetadataFetcher{
		FAvailable: a,
	}
	fakeR := FakeRouteTableFetcher{}
	fakeM.Meta = make(map[string]string)
	d.MetadataFetcher = fakeM
	d.RouteTableFetcher = fakeR
	return d
}

func TestSetupUnavailable(t *testing.T) {
	d := getD(false)
	err := d.Setup()
	if err != nil {
		t.Log(err)
		t.Fail()
	}
	if d.MetadataFetcher.Available() {
		t.Fail()
	}
}

func TestSetupAvailable(t *testing.T) {
	d := getD(true)
	err := d.Setup()
	if err != nil {
		t.Fail()
	}
	if !d.MetadataFetcher.Available() {
		t.Fail()
	}
}

func TestSetupHealthChecks(t *testing.T) {
	d := getD(true)
	err := d.Setup()
	if err != nil {
		t.Fail()
	}
	if d.Config.Healthchecks["public"].IsRunning() {
		t.Log("HealthChecks already running")
		t.Fail()
	}
	d.runHealthChecks()
	if !d.Config.Healthchecks["public"].IsRunning() {
		t.Log("HealthChecks did not start running")
		t.Fail()
	}
}

func TestGetSubnetIdMacFail(t *testing.T) {
	d := getD(true)
	_, err := d.GetSubnetId()
	if err == nil {
		t.Fail()
	}
}

func TestGetSubnetIdMacFail2(t *testing.T) {
	d := getD(true)
	d.MetadataFetcher.(FakeMetadataFetcher).Meta["mac"] = "06:1d:ea:6f:8c:6e"
	_, err := d.GetSubnetId()
	if err == nil {
		t.Fail()
	}
}

func TestGetSubnetIdMacOk(t *testing.T) {
	d := getD(true)
	d.MetadataFetcher.(FakeMetadataFetcher).Meta["mac"] = "06:1d:ea:6f:8c:6e"
	d.MetadataFetcher.(FakeMetadataFetcher).Meta["network/interfaces/macs/06:1d:ea:6f:8c:6e/subnet-id"] = "subnet-28b0e940"
	val, err := d.GetSubnetId()
	if err != nil {
		t.Fail()
	}
	if val != "subnet-28b0e940" {
		t.Fail()
	}
}
<commit_msg>Fix the tests<commit_after>package daemon

import (
	"errors"
	"fmt"
	"github.com/aws/aws-sdk-go/service/ec2"
	"testing"
)

type FakeMetadataFetcher struct {
	FAvailable bool
	Meta       map[string]string
}

func (m FakeMetadataFetcher) Available() bool {
	return m.FAvailable
}

func (m FakeMetadataFetcher) GetMetadata(key string) (string, error) {
	v, ok := m.Meta[key]
	if ok {
		return v, nil
	}
	return v, errors.New(fmt.Sprintf("Key %s unknown"))
}

type FakeRouteTableFetcher struct{}

func (r FakeRouteTableFetcher) GetRouteTables() ([]*ec2.RouteTable, error) {
	return []*ec2.RouteTable{}, nil
}

func getD(a bool) Daemon {
	d := Daemon{
		ConfigFile: "../tests/awsnycast.yaml",
	}
	fakeM := FakeMetadataFetcher{
		FAvailable: a,
	}
	fakeR := FakeRouteTableFetcher{}
	fakeM.Meta = make(map[string]string)
	fakeM.Meta["placement/availability-zone"] = "us-west-1a"
	d.MetadataFetcher = fakeM
	d.RouteTableFetcher = fakeR
	return d
}

func TestSetupUnavailable(t *testing.T) {
	d := getD(false)
	err := d.Setup()
	if err != nil {
		t.Log(err)
		t.Fail()
	}
	if d.MetadataFetcher.Available() {
		t.Fail()
	}
}

func TestSetupAvailable(t *testing.T) {
	d := getD(true)
	err := d.Setup()
	if err != nil {
		t.Fail()
	}
	if !d.MetadataFetcher.Available() {
		t.Fail()
	}
}

func TestSetupHealthChecks(t *testing.T) {
	d := getD(true)
	err := d.Setup()
	if err != nil {
		t.Fail()
	}
	if d.Config.Healthchecks["public"].IsRunning() {
		t.Log("HealthChecks already running")
		t.Fail()
	}
	d.runHealthChecks()
	if !d.Config.Healthchecks["public"].IsRunning() {
		t.Log("HealthChecks did not start running")
		t.Fail()
	}
}

func TestGetSubnetIdMacFail(t *testing.T) {
	d := getD(true)
	_, err := d.GetSubnetId()
	if err == nil {
		t.Fail()
	}
}

func TestGetSubnetIdMacFail2(t *testing.T) {
	d := getD(true)
	d.MetadataFetcher.(FakeMetadataFetcher).Meta["mac"] = "06:1d:ea:6f:8c:6e"
	_, err := d.GetSubnetId()
	if err == nil {
		t.Fail()
	}
}

func TestGetSubnetIdMacOk(t *testing.T) {
	d := getD(true)
	d.MetadataFetcher.(FakeMetadataFetcher).Meta["mac"] = "06:1d:ea:6f:8c:6e"
	d.MetadataFetcher.(FakeMetadataFetcher).Meta["network/interfaces/macs/06:1d:ea:6f:8c:6e/subnet-id"] = "subnet-28b0e940"
	val, err := d.GetSubnetId()
	if err != nil {
		t.Fail()
	}
	if val != "subnet-28b0e940" {
		t.Fail()
	}
}
<|endoftext|>


<commit_before>package render

import (
	"encoding/json"
	"net/http"

	"github.com/Sirupsen/logrus"
	"github.com/anthonynsimon/parrot/errors"
)

var jsonContentType = "application/json; charset=utf-8"

type apiResponseBody struct {
	responseMeta `json:"meta,omitempty"`
	Payload      interface{} `json:"payload,omitempty"`
}

type responseMeta struct {
	Status int   `json:"status,omitempty"`
	Error  error `json:"error,omitempty"`
}

func Error(w http.ResponseWriter, err error) {
	status := http.StatusInternalServerError
	if e, ok := err.(*errors.Error); ok {
		status = e.Status
	}
	ErrorWithStatus(w, status, err)
}

func ErrorWithStatus(w http.ResponseWriter, status int, errs error) {
	w.Header().Set("Content-Type", jsonContentType)
	w.WriteHeader(status)

	body := apiResponseBody{
		responseMeta: responseMeta{
			Status: status,
			Error:  errs},
		Payload: nil}

	handleBody(w, body)
}

func JSON(w http.ResponseWriter, status int, payload interface{}) {
	w.Header().Set("Content-Type", jsonContentType)
	w.WriteHeader(status)

	body := apiResponseBody{
		responseMeta: responseMeta{
			Status: status},
		Payload: payload}

	handleBody(w, body)
}

func JSONWithHeaders(w http.ResponseWriter, status int, headers map[string]string, payload interface{}) {
	h := w.Header()
	h.Set("Content-Type", jsonContentType)
	for k, v := range headers {
		h.Set(k, v)
	}
	w.WriteHeader(status)

	body := apiResponseBody{
		responseMeta: responseMeta{
			Status: status},
		Payload: payload}

	handleBody(w, body)
}

func handleBody(w http.ResponseWriter, body apiResponseBody) {
	encoded, err := json.Marshal(body)
	if err != nil {
		logrus.Error(err)
	}

	w.Write(encoded)
}
<commit_msg>Fix param name<commit_after>package render

import (
	"encoding/json"
	"net/http"

	"github.com/Sirupsen/logrus"
	"github.com/anthonynsimon/parrot/errors"
)

var jsonContentType = "application/json; charset=utf-8"

type apiResponseBody struct {
	responseMeta `json:"meta,omitempty"`
	Payload      interface{} `json:"payload,omitempty"`
}

type responseMeta struct {
	Status int   `json:"status,omitempty"`
	Error  error `json:"error,omitempty"`
}

func Error(w http.ResponseWriter, err error) {
	status := http.StatusInternalServerError
	if e, ok := err.(*errors.Error); ok {
		status = e.Status
	}
	ErrorWithStatus(w, status, err)
}

func ErrorWithStatus(w http.ResponseWriter, status int, err error) {
	w.Header().Set("Content-Type", jsonContentType)
	w.WriteHeader(status)

	body := apiResponseBody{
		responseMeta: responseMeta{
			Status: status,
			Error:  err},
		Payload: nil}

	handleBody(w, body)
}

func JSON(w http.ResponseWriter, status int, payload interface{}) {
	w.Header().Set("Content-Type", jsonContentType)
	w.WriteHeader(status)

	body := apiResponseBody{
		responseMeta: responseMeta{
			Status: status},
		Payload: payload}

	handleBody(w, body)
}

func JSONWithHeaders(w http.ResponseWriter, status int, headers map[string]string, payload interface{}) {
	h := w.Header()
	h.Set("Content-Type", jsonContentType)
	for k, v := range headers {
		h.Set(k, v)
	}
	w.WriteHeader(status)

	body := apiResponseBody{
		responseMeta: responseMeta{
			Status: status},
		Payload: payload}

	handleBody(w, body)
}

func handleBody(w http.ResponseWriter, body apiResponseBody) {
	encoded, err := json.Marshal(body)
	if err != nil {
		logrus.Error(err)
	}

	w.Write(encoded)
}
<|endoftext|>


<commit_before>package reqlimit

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func TestRequestLimitHandler(t *testing.T) {
	maxRequests := uint64(10)
	duration := time.Second * 5

	limiter := New(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprintln(w, "success")
	}), maxRequests, duration)

	ts := httptest.NewServer(limiter)
	defer ts.Close()

	for i := 0; uint64(i) < maxRequests; i++ {
		_, err := http.Get(ts.URL)
		if err != nil {
			t.Fatal(err)
		}
	}

	res, err := http.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	if res.StatusCode != http.StatusTooManyRequests {
		t.Fatal("limiter did not return TooManyRequests after exceeding maxrequests")
	}

	time.Sleep(duration)

	res, err = http.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	if res.StatusCode != http.StatusOK {
		t.Fatal("limiter did not return StatusOK after waiting for the duration to expire")
	}
}

func BenchmarkLimitedHandler(b *testing.B) {
	maxRequests := uint64(10000)
	duration := time.Second

	r, _ := http.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})
	limitedHandler := New(testHandler, maxRequests, duration)

	for i := 0; i < b.N; i++ {
		limitedHandler.ServeHTTP(w, r)
	}
}

func BenchmarkHandler(b *testing.B) {
	r, _ := http.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})

	for i := 0; i < b.N; i++ {
		testHandler.ServeHTTP(w, r)
	}
}
<commit_msg>add TestRequestLimitHandlerDiffferentIPs<commit_after>package reqlimit

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func TestRequestLimitHandler(t *testing.T) {
	maxRequests := uint64(10)
	duration := time.Second * 5

	limiter := New(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprintln(w, "success")
	}), maxRequests, duration)

	ts := httptest.NewServer(limiter)
	defer ts.Close()

	for i := 0; uint64(i) < maxRequests; i++ {
		_, err := http.Get(ts.URL)
		if err != nil {
			t.Fatal(err)
		}
	}

	res, err := http.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	if res.StatusCode != http.StatusTooManyRequests {
		t.Fatal("limiter did not return TooManyRequests after exceeding maxrequests")
	}

	time.Sleep(duration)

	res, err = http.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	if res.StatusCode != http.StatusOK {
		t.Fatal("limiter did not return StatusOK after waiting for the duration to expire")
	}
}

func testRequest(remoteAddr string) (*httptest.ResponseRecorder, *http.Request) {
	r, _ := http.NewRequest("GET", "/", nil)
	r.RemoteAddr = remoteAddr
	w := httptest.NewRecorder()
	return w, r
}

func TestRequestLimitHandlerDifferentIPs(t *testing.T) {
	maxRequests := uint64(10)
	duration := time.Minute

	testHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})
	limitedHandler := New(testHandler, maxRequests, duration)

	// fill up the limit
	for i := 0; uint64(i) < maxRequests; i++ {
		limitedHandler.ServeHTTP(testRequest("3.4.5.6:7483"))
	}

	// verify that the limit is hit
	w, r := testRequest("3.4.5.6:8080")
	limitedHandler.ServeHTTP(w, r)
	if w.Result().StatusCode != http.StatusTooManyRequests {
		t.Fatal("expected the limit to be hit")
	}

	// verify that we can successfully request with a different ip
	w, r = testRequest("1.2.3.4:1234")
	limitedHandler.ServeHTTP(w, r)
	if w.Result().StatusCode != http.StatusOK {
		t.Fatal("expected to get StatusOK after changing IP")
	}
}

func BenchmarkLimitedHandler(b *testing.B) {
	b.ReportAllocs()

	maxRequests := uint64(10000)
	duration := time.Second

	r, _ := http.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})
	limitedHandler := New(testHandler, maxRequests, duration)

	for i := 0; i < b.N; i++ {
		limitedHandler.ServeHTTP(w, r)
	}
}

func BenchmarkHandler(b *testing.B) {
	b.ReportAllocs()

	r, _ := http.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})

	for i := 0; i < b.N; i++ {
		testHandler.ServeHTTP(w, r)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Package opt provides sets of options used by LevelDB.
package opt

import (
	"errors"
	"sync"

	"github.com/syndtr/goleveldb/leveldb/cache"
	"github.com/syndtr/goleveldb/leveldb/comparer"
	"github.com/syndtr/goleveldb/leveldb/filter"
)

var (
	ErrInvalid    = errors.New("invalid value")
	ErrNotSet     = errors.New("not set")
	ErrNotAllowed = errors.New("not allowed")
)

const (
	DefaultWriteBuffer          = 4 << 20
	DefaultMaxOpenFiles         = 1000
	DefaultBlockCacheSize       = 8 << 20
	DefaultBlockSize            = 4096
	DefaultBlockRestartInterval = 16
	DefaultCompressionType      = SnappyCompression
)

type OptionsFlag uint

const (
	// If set, the database will be created if it is missing.
	OFCreateIfMissing OptionsFlag = 1 << iota

	// If set, an error is raised if the database already exists.
	OFErrorIfExist

	// If set, the implementation will do aggressive checking of the
	// data it is processing and will stop early if it detects any
	// errors.  This may have unforeseen ramifications: for example, a
	// corruption of one DB entry may cause a large number of entries to
	// become unreadable or for the entire DB to become unopenable.
	OFParanoidCheck
)

// Database compression type
type Compression uint

func (c Compression) String() string {
	switch c {
	case DefaultCompression:
		return "default"
	case NoCompression:
		return "none"
	case SnappyCompression:
		return "snappy"
	}
	return "unknown"
}

const (
	DefaultCompression Compression = iota
	NoCompression
	SnappyCompression
	nCompression
)

// Options represent sets of LevelDB options.
type Options struct {
	// Comparer used to define the order of keys in the table.
	// Default: a comparer that uses lexicographic byte-wise ordering
	//
	// REQUIRES: The client must ensure that the comparer supplied
	// here has the same name and orders keys *exactly* the same as the
	// comparer provided to previous open calls on the same DB.
	Comparer comparer.Comparer

	// Specify the database flag.
	Flag OptionsFlag

	// Amount of data to build up in memory (backed by an unsorted journal
	// on disk) before converting to a sorted on-disk file.
	//
	// Larger values increase performance, especially during bulk loads.
	// Up to two write buffers may be held in memory at the same time,
	// so you may wish to adjust this parameter to control memory usage.
	// Also, a larger write buffer will result in a longer recovery time
	// the next time the database is opened.
	//
	// Default: 4MB
	WriteBuffer int

	// Number of open files that can be used by the DB.  You may need to
	// increase this if your database has a large working set (budget
	// one open file per 2MB of working set).
	//
	// Default: 1000
	MaxOpenFiles int

	// Control over blocks (user data is stored in a set of blocks, and
	// a block is the unit of reading from disk).

	// If non-NULL, use the specified cache for blocks.
	// If NULL, leveldb will automatically create and use an 8MB internal cache.
	// Default: NULL
	BlockCache cache.Cache

	// Approximate size of user data packed per block.  Note that the
	// block size specified here corresponds to uncompressed data.  The
	// actual size of the unit read from disk may be smaller if
	// compression is enabled.  This parameter can be changed dynamically.
	//
	// Default: 4K
	BlockSize int

	// Number of keys between restart points for delta encoding of keys.
	// This parameter can be changed dynamically.  Most clients should
	// leave this parameter alone.
	//
	// Default: 16
	BlockRestartInterval int

	// Compress blocks using the specified compression algorithm.  This
	// parameter can be changed dynamically.
	//
	// Default: kSnappyCompression, which gives lightweight but fast
	// compression.
	//
	// Typical speeds of kSnappyCompression on an Intel(R) Core(TM)2 2.4GHz:
	//    ~200-500MB/s compression
	//    ~400-800MB/s decompression
	// Note that these speeds are significantly faster than most
	// persistent storage speeds, and therefore it is typically never
	// worth switching to kNoCompression.  Even if the input data is
	// incompressible, the kSnappyCompression implementation will
	// efficiently detect that and will switch to uncompressed mode.
	CompressionType Compression

	// If non-NULL, use the specified filter policy to reduce disk reads.
	// Many applications will benefit from passing the result of
	// NewBloomFilter() here.
	//
	// As long as the same filter (name) was used as last time the
	// database was opened, the previous filter is reused. That is,
	// the filter does not need to be rebuilt. This is made possible
	// since each filter is persisted to disk on a per sstable
	// basis.
	//
	// As opposed to the comparer, a filter can be replaced after a
	// database has been created. If this is done, the previous
	// persisted filter will be ignored for every old sstable.
	// Every new table will use the newly introduced filter. This
	// means that all/some sstables will lack a filter during a
	// transition period. Note that this might have an impact on
	// performance. Rewriting every single key/value will force
	// introduction of the new filter.
	//
	// Default: NULL
	Filter filter.Filter

	mu sync.RWMutex
}

// OptionsGetter wraps methods used to get sanitized options.
type OptionsGetter interface {
	GetComparer() comparer.Comparer
	HasFlag(flag OptionsFlag) bool
	GetWriteBuffer() int
	GetMaxOpenFiles() int
	GetBlockCache() cache.Cache
	GetBlockSize() int
	GetBlockRestartInterval() int
	GetCompressionType() Compression
	GetFilter() filter.Filter
}

// OptionsSetter wraps methods used to set options.
type OptionsSetter interface {
	SetComparer(cmp comparer.Comparer) error
	SetFlag(flag OptionsFlag) error
	ClearFlag(flag OptionsFlag) error
	SetWriteBuffer(size int) error
	SetMaxOpenFiles(max int) error
	SetBlockCache(cache cache.Cache) error
	SetBlockCacheCapacity(capacity int) error
	SetBlockSize(size int) error
	SetBlockRestartInterval(interval int) error
	SetCompressionType(compression Compression) error
	SetFilter(p filter.Filter) error
}

// Getter

func (o *Options) GetComparer() comparer.Comparer {
	if o == nil {
		return comparer.DefaultComparer
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.Comparer == nil {
		return comparer.DefaultComparer
	}
	return o.Comparer
}

func (o *Options) HasFlag(flag OptionsFlag) bool {
	if o == nil {
		return false
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	return (o.Flag & flag) != 0
}

func (o *Options) GetWriteBuffer() int {
	if o == nil {
		return DefaultWriteBuffer
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.WriteBuffer <= 0 {
		return DefaultWriteBuffer
	}
	return o.WriteBuffer
}

func (o *Options) GetMaxOpenFiles() int {
	if o == nil {
		return DefaultMaxOpenFiles
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.MaxOpenFiles <= 0 {
		return DefaultMaxOpenFiles
	}
	return o.MaxOpenFiles
}

func (o *Options) GetBlockCache() cache.Cache {
	if o == nil {
		return nil
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	return o.BlockCache
}

func (o *Options) GetBlockSize() int {
	if o == nil {
		return DefaultBlockSize
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.BlockSize <= 0 {
		return DefaultBlockSize
	}
	return o.BlockSize
}

func (o *Options) GetBlockRestartInterval() int {
	if o == nil {
		return DefaultBlockRestartInterval
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.BlockRestartInterval <= 0 {
		return DefaultBlockRestartInterval
	}
	return o.BlockRestartInterval
}

func (o *Options) GetCompressionType() Compression {
	if o == nil {
		return DefaultCompressionType
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.CompressionType <= DefaultCompression || o.CompressionType >= nCompression {
		return DefaultCompressionType
	}
	return o.CompressionType
}

func (o *Options) GetFilter() filter.Filter {
	if o == nil {
		return nil
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	return o.Filter
}

// Setter

func (o *Options) SetComparer(cmp comparer.Comparer) error {
	if o == nil {
		return ErrNotSet
	}
	if cmp == nil {
		return ErrInvalid
	}
	o.mu.Lock()
	o.Comparer = cmp
	o.mu.Unlock()
	return nil
}

func (o *Options) SetFlag(flag OptionsFlag) error {
	if o == nil {
		return ErrNotSet
	}
	o.mu.Lock()
	o.Flag |= flag
	o.mu.Unlock()
	return nil
}

func (o *Options) ClearFlag(flag OptionsFlag) error {
	if o == nil {
		return ErrNotSet
	}
	o.mu.Lock()
	o.Flag &= ^flag
	o.mu.Unlock()
	return nil
}

func (o *Options) SetWriteBuffer(size int) error {
	if o == nil {
		return ErrNotSet
	}
	if size <= 0 {
		return ErrInvalid
	}
	o.mu.Lock()
	o.WriteBuffer = size
	o.mu.Unlock()
	return nil
}

func (o *Options) SetMaxOpenFiles(max int) error {
	if o == nil {
		return ErrNotSet
	}
	if max <= 0 {
		return ErrInvalid
	}
	o.mu.Lock()
	o.MaxOpenFiles = max
	o.mu.Unlock()
	return nil
}

func (o *Options) SetBlockCache(cache cache.Cache) error {
	if o == nil {
		return ErrNotSet
	}
	o.mu.Lock()
	o.BlockCache = cache
	o.mu.Unlock()
	return nil
}

func (o *Options) SetBlockCacheCapacity(capacity int) error {
	if o == nil {
		return ErrNotSet
	}
	o.mu.Lock()
	if o.BlockCache == nil {
		return ErrNotSet
	}
	o.BlockCache.SetCapacity(capacity)
	o.mu.Unlock()
	return nil
}

func (o *Options) SetBlockSize(size int) error {
	if o == nil {
		return ErrNotSet
	}
	if size <= 0 {
		return ErrInvalid
	}
	o.mu.Lock()
	o.BlockSize = size
	o.mu.Unlock()
	return nil
}

func (o *Options) SetBlockRestartInterval(interval int) error {
	if o == nil {
		return ErrNotSet
	}
	if interval <= 0 {
		return ErrInvalid
	}
	o.mu.Lock()
	o.BlockRestartInterval = interval
	o.mu.Unlock()
	return nil
}

func (o *Options) SetCompressionType(compression Compression) error {
	if o == nil {
		return ErrNotSet
	}
	if o.CompressionType >= nCompression {
		return ErrInvalid
	}
	o.mu.Lock()
	o.CompressionType = compression
	o.mu.Unlock()
	return nil
}

func (o *Options) SetFilter(p filter.Filter) error {
	if o == nil {
		return ErrNotSet
	}
	o.mu.Lock()
	o.Filter = p
	o.mu.Unlock()
	return nil
}

type ReadOptionsFlag uint

const (
	// If true, all data read from underlying storage will be
	// verified against corresponding checksums.
	RFVerifyChecksums ReadOptionsFlag = 1 << iota

	// Should the data read for this iteration be cached in memory?
	// If set iteration chaching will be disabled.
	// Callers may wish to set this flag for bulk scans.
	RFDontFillCache
)

// ReadOptions represent sets of options used by LevelDB during read
// operations.
type ReadOptions struct {
	// Specify the read flag.
	Flag ReadOptionsFlag
}

type ReadOptionsGetter interface {
	HasFlag(flag ReadOptionsFlag) bool
}

func (o *ReadOptions) HasFlag(flag ReadOptionsFlag) bool {
	if o == nil {
		return false
	}
	return (o.Flag & flag) != 0
}

type WriteOptionsFlag uint

const (
	// If set, the write will be flushed from the operating system
	// buffer cache (by calling WritableFile::Sync()) before the write
	// is considered complete.  If this flag is true, writes will be
	// slower.
	//
	// If this flag is false, and the machine crashes, some recent
	// writes may be lost.  Note that if it is just the process that
	// crashes (i.e., the machine does not reboot), no writes will be
	// lost even if sync==false.
	//
	// In other words, a DB write with sync==false has similar
	// crash semantics as the "write()" system call.  A DB write
	// with sync==true has similar crash semantics to a "write()"
	// system call followed by "fsync()".
	WFSync WriteOptionsFlag = 1 << iota
)

// WriteOptions represent sets of options used by LevelDB during write
// operations.
type WriteOptions struct {
	// Specify the write flag.
	Flag WriteOptionsFlag
}

type WriteOptionsGetter interface {
	HasFlag(flag WriteOptionsFlag) bool
}

func (o *WriteOptions) HasFlag(flag WriteOptionsFlag) bool {
	if o == nil {
		return false
	}
	return (o.Flag & flag) != 0
}
<commit_msg>Documenting opt.Options.Comparer<commit_after>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Package opt provides sets of options used by LevelDB.
package opt

import (
	"errors"
	"sync"

	"github.com/syndtr/goleveldb/leveldb/cache"
	"github.com/syndtr/goleveldb/leveldb/comparer"
	"github.com/syndtr/goleveldb/leveldb/filter"
)

var (
	ErrInvalid    = errors.New("invalid value")
	ErrNotSet     = errors.New("not set")
	ErrNotAllowed = errors.New("not allowed")
)

const (
	DefaultWriteBuffer          = 4 << 20
	DefaultMaxOpenFiles         = 1000
	DefaultBlockCacheSize       = 8 << 20
	DefaultBlockSize            = 4096
	DefaultBlockRestartInterval = 16
	DefaultCompressionType      = SnappyCompression
)

type OptionsFlag uint

const (
	// If set, the database will be created if it is missing.
	OFCreateIfMissing OptionsFlag = 1 << iota

	// If set, an error is raised if the database already exists.
	OFErrorIfExist

	// If set, the implementation will do aggressive checking of the
	// data it is processing and will stop early if it detects any
	// errors.  This may have unforeseen ramifications: for example, a
	// corruption of one DB entry may cause a large number of entries to
	// become unreadable or for the entire DB to become unopenable.
	OFParanoidCheck
)

// Database compression type
type Compression uint

func (c Compression) String() string {
	switch c {
	case DefaultCompression:
		return "default"
	case NoCompression:
		return "none"
	case SnappyCompression:
		return "snappy"
	}
	return "unknown"
}

const (
	DefaultCompression Compression = iota
	NoCompression
	SnappyCompression
	nCompression
)

// Options represent sets of LevelDB options.
type Options struct {
	// Comparer used to define the order of keys in the table.
	// Default: a comparer that uses lexicographic byte-wise
	// ordering.
	//
	// REQUIRES: The client must ensure that the comparer supplied
	// here has the same name and orders keys *exactly* the same as the
	// comparer provided to previous open calls on the same DB.
	// Additionally, the client must also make sure that the
	// supplied comparer retains the same name. Otherwise, an error
	// will be returned on reopening the database.
	Comparer comparer.Comparer

	// Specify the database flag.
	Flag OptionsFlag

	// Amount of data to build up in memory (backed by an unsorted journal
	// on disk) before converting to a sorted on-disk file.
	//
	// Larger values increase performance, especially during bulk loads.
	// Up to two write buffers may be held in memory at the same time,
	// so you may wish to adjust this parameter to control memory usage.
	// Also, a larger write buffer will result in a longer recovery time
	// the next time the database is opened.
	//
	// Default: 4MB
	WriteBuffer int

	// Number of open files that can be used by the DB.  You may need to
	// increase this if your database has a large working set (budget
	// one open file per 2MB of working set).
	//
	// Default: 1000
	MaxOpenFiles int

	// Control over blocks (user data is stored in a set of blocks, and
	// a block is the unit of reading from disk).

	// If non-NULL, use the specified cache for blocks.
	// If NULL, leveldb will automatically create and use an 8MB internal cache.
	// Default: NULL
	BlockCache cache.Cache

	// Approximate size of user data packed per block.  Note that the
	// block size specified here corresponds to uncompressed data.  The
	// actual size of the unit read from disk may be smaller if
	// compression is enabled.  This parameter can be changed dynamically.
	//
	// Default: 4K
	BlockSize int

	// Number of keys between restart points for delta encoding of keys.
	// This parameter can be changed dynamically.  Most clients should
	// leave this parameter alone.
	//
	// Default: 16
	BlockRestartInterval int

	// Compress blocks using the specified compression algorithm.  This
	// parameter can be changed dynamically.
	//
	// Default: kSnappyCompression, which gives lightweight but fast
	// compression.
	//
	// Typical speeds of kSnappyCompression on an Intel(R) Core(TM)2 2.4GHz:
	//    ~200-500MB/s compression
	//    ~400-800MB/s decompression
	// Note that these speeds are significantly faster than most
	// persistent storage speeds, and therefore it is typically never
	// worth switching to kNoCompression.  Even if the input data is
	// incompressible, the kSnappyCompression implementation will
	// efficiently detect that and will switch to uncompressed mode.
	CompressionType Compression

	// If non-NULL, use the specified filter policy to reduce disk reads.
	// Many applications will benefit from passing the result of
	// NewBloomFilter() here.
	//
	// As long as the same filter (name) was used as last time the
	// database was opened, the previous filter is reused. That is,
	// the filter does not need to be rebuilt. This is made possible
	// since each filter is persisted to disk on a per sstable
	// basis.
	//
	// As opposed to the comparer, a filter can be replaced after a
	// database has been created. If this is done, the previous
	// persisted filter will be ignored for every old sstable.
	// Every new table will use the newly introduced filter. This
	// means that all/some sstables will lack a filter during a
	// transition period. Note that this might have an impact on
	// performance. Rewriting every single key/value will force
	// introduction of the new filter.
	//
	// Default: NULL
	Filter filter.Filter

	mu sync.RWMutex
}

// OptionsGetter wraps methods used to get sanitized options.
type OptionsGetter interface {
	GetComparer() comparer.Comparer
	HasFlag(flag OptionsFlag) bool
	GetWriteBuffer() int
	GetMaxOpenFiles() int
	GetBlockCache() cache.Cache
	GetBlockSize() int
	GetBlockRestartInterval() int
	GetCompressionType() Compression
	GetFilter() filter.Filter
}

// OptionsSetter wraps methods used to set options.
type OptionsSetter interface {
	SetComparer(cmp comparer.Comparer) error
	SetFlag(flag OptionsFlag) error
	ClearFlag(flag OptionsFlag) error
	SetWriteBuffer(size int) error
	SetMaxOpenFiles(max int) error
	SetBlockCache(cache cache.Cache) error
	SetBlockCacheCapacity(capacity int) error
	SetBlockSize(size int) error
	SetBlockRestartInterval(interval int) error
	SetCompressionType(compression Compression) error
	SetFilter(p filter.Filter) error
}

// Getter

func (o *Options) GetComparer() comparer.Comparer {
	if o == nil {
		return comparer.DefaultComparer
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.Comparer == nil {
		return comparer.DefaultComparer
	}
	return o.Comparer
}

func (o *Options) HasFlag(flag OptionsFlag) bool {
	if o == nil {
		return false
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	return (o.Flag & flag) != 0
}

func (o *Options) GetWriteBuffer() int {
	if o == nil {
		return DefaultWriteBuffer
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.WriteBuffer <= 0 {
		return DefaultWriteBuffer
	}
	return o.WriteBuffer
}

func (o *Options) GetMaxOpenFiles() int {
	if o == nil {
		return DefaultMaxOpenFiles
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.MaxOpenFiles <= 0 {
		return DefaultMaxOpenFiles
	}
	return o.MaxOpenFiles
}

func (o *Options) GetBlockCache() cache.Cache {
	if o == nil {
		return nil
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	return o.BlockCache
}

func (o *Options) GetBlockSize() int {
	if o == nil {
		return DefaultBlockSize
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.BlockSize <= 0 {
		return DefaultBlockSize
	}
	return o.BlockSize
}

func (o *Options) GetBlockRestartInterval() int {
	if o == nil {
		return DefaultBlockRestartInterval
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.BlockRestartInterval <= 0 {
		return DefaultBlockRestartInterval
	}
	return o.BlockRestartInterval
}

func (o *Options) GetCompressionType() Compression {
	if o == nil {
		return DefaultCompressionType
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	if o.CompressionType <= DefaultCompression || o.CompressionType >= nCompression {
		return DefaultCompressionType
	}
	return o.CompressionType
}

func (o *Options) GetFilter() filter.Filter {
	if o == nil {
		return nil
	}
	o.mu.RLock()
	defer o.mu.RUnlock()
	return o.Filter
}

// Setter

func (o *Options) SetComparer(cmp comparer.Comparer) error {
	if o == nil {
		return ErrNotSet
	}
	if cmp == nil {
		return ErrInvalid
	}
	o.mu.Lock()
	o.Comparer = cmp
	o.mu.Unlock()
	return nil
}

func (o *Options) SetFlag(flag OptionsFlag) error {
	if o == nil {
		return ErrNotSet
	}
	o.mu.Lock()
	o.Flag |= flag
	o.mu.Unlock()
	return nil
}

func (o *Options) ClearFlag(flag OptionsFlag) error {
	if o == nil {
		return ErrNotSet
	}
	o.mu.Lock()
	o.Flag &= ^flag
	o.mu.Unlock()
	return nil
}

func (o *Options) SetWriteBuffer(size int) error {
	if o == nil {
		return ErrNotSet
	}
	if size <= 0 {
		return ErrInvalid
	}
	o.mu.Lock()
	o.WriteBuffer = size
	o.mu.Unlock()
	return nil
}

func (o *Options) SetMaxOpenFiles(max int) error {
	if o == nil {
		return ErrNotSet
	}
	if max <= 0 {
		return ErrInvalid
	}
	o.mu.Lock()
	o.MaxOpenFiles = max
	o.mu.Unlock()
	return nil
}

func (o *Options) SetBlockCache(cache cache.Cache) error {
	if o == nil {
		return ErrNotSet
	}
	o.mu.Lock()
	o.BlockCache = cache
	o.mu.Unlock()
	return nil
}

func (o *Options) SetBlockCacheCapacity(capacity int) error {
	if o == nil {
		return ErrNotSet
	}
	o.mu.Lock()
	if o.BlockCache == nil {
		return ErrNotSet
	}
	o.BlockCache.SetCapacity(capacity)
	o.mu.Unlock()
	return nil
}

func (o *Options) SetBlockSize(size int) error {
	if o == nil {
		return ErrNotSet
	}
	if size <= 0 {
		return ErrInvalid
	}
	o.mu.Lock()
	o.BlockSize = size
	o.mu.Unlock()
	return nil
}

func (o *Options) SetBlockRestartInterval(interval int) error {
	if o == nil {
		return ErrNotSet
	}
	if interval <= 0 {
		return ErrInvalid
	}
	o.mu.Lock()
	o.BlockRestartInterval = interval
	o.mu.Unlock()
	return nil
}

func (o *Options) SetCompressionType(compression Compression) error {
	if o == nil {
		return ErrNotSet
	}
	if o.CompressionType >= nCompression {
		return ErrInvalid
	}
	o.mu.Lock()
	o.CompressionType = compression
	o.mu.Unlock()
	return nil
}

func (o *Options) SetFilter(p filter.Filter) error {
	if o == nil {
		return ErrNotSet
	}
	o.mu.Lock()
	o.Filter = p
	o.mu.Unlock()
	return nil
}

type ReadOptionsFlag uint

const (
	// If true, all data read from underlying storage will be
	// verified against corresponding checksums.
	RFVerifyChecksums ReadOptionsFlag = 1 << iota

	// Should the data read for this iteration be cached in memory?
	// If set iteration chaching will be disabled.
	// Callers may wish to set this flag for bulk scans.
	RFDontFillCache
)

// ReadOptions represent sets of options used by LevelDB during read
// operations.
type ReadOptions struct {
	// Specify the read flag.
	Flag ReadOptionsFlag
}

type ReadOptionsGetter interface {
	HasFlag(flag ReadOptionsFlag) bool
}

func (o *ReadOptions) HasFlag(flag ReadOptionsFlag) bool {
	if o == nil {
		return false
	}
	return (o.Flag & flag) != 0
}

type WriteOptionsFlag uint

const (
	// If set, the write will be flushed from the operating system
	// buffer cache (by calling WritableFile::Sync()) before the write
	// is considered complete.  If this flag is true, writes will be
	// slower.
	//
	// If this flag is false, and the machine crashes, some recent
	// writes may be lost.  Note that if it is just the process that
	// crashes (i.e., the machine does not reboot), no writes will be
	// lost even if sync==false.
	//
	// In other words, a DB write with sync==false has similar
	// crash semantics as the "write()" system call.  A DB write
	// with sync==true has similar crash semantics to a "write()"
	// system call followed by "fsync()".
	WFSync WriteOptionsFlag = 1 << iota
)

// WriteOptions represent sets of options used by LevelDB during write
// operations.
type WriteOptions struct {
	// Specify the write flag.
	Flag WriteOptionsFlag
}

type WriteOptionsGetter interface {
	HasFlag(flag WriteOptionsFlag) bool
}

func (o *WriteOptions) HasFlag(flag WriteOptionsFlag) bool {
	if o == nil {
		return false
	}
	return (o.Flag & flag) != 0
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/urfave/cli"

	"github.com/dirk/quickhook/context"
	"github.com/dirk/quickhook/hooks"
)

const VERSION = "1.2.0"

func main() {
	context, err := setupContextInWd()
	if err != nil {
		panic(err)
	}

	app := cli.NewApp()
	app.Name = "quickhook"
	app.Version = VERSION
	app.Usage = "Git hook runner"

	app.Commands = []cli.Command{
		{
			Name:  "hook",
			Usage: "Run a hook",
			Action: func(c *cli.Context) error {
				return cli.ShowSubcommandHelp(c)
			},
			Subcommands: []cli.Command{
				cli.Command{
					Name: "pre-commit",
					Flags: []cli.Flag{
						allFlag(),
						filesFlag(),
						noColorFlag(),
					},
					Action: func(c *cli.Context) error {
						files := []string{}
						if c.Bool("files") {
							files = c.Args()
						}

						err := hooks.PreCommit(context, &hooks.PreCommitOpts{
							All:     c.Bool("all"),
							Files:   files,
							NoColor: c.Bool("no-color"),
						})
						if err != nil {
							panic(err)
						}
						return nil
					},
				},
				cli.Command{
					Name: "commit-msg",
					Flags: []cli.Flag{
						noColorFlag(),
					},
					Action: func(c *cli.Context) error {
						messageTempFile := c.Args().Get(0)
						if messageTempFile == "" {
							fmt.Println("Missing message temp file argument")
							os.Exit(1)
						}

						err := hooks.CommitMsg(context, &hooks.CommitMsgOpts{
							NoColor:         c.Bool("no-color"),
							MessageTempFile: messageTempFile,
						})
						if err != nil {
							panic(err)
						}
						return nil
					},
				},
			},
		},
		{
			Name:      "install",
			Usage:     "Install Quickhook shims into .git/hooks",
			ArgsUsage: " ", // Don't show "[arguments...]"
			Flags: []cli.Flag{
				yesFlag(),
			},
			Action: func(c *cli.Context) error {
				prompt := c.Bool("yes") != true

				err := Install(context, prompt)
				if err != nil {
					panic(err)
				}
				return nil
			},
		},
	}

	app.Run(os.Args)
}

// Set up `Context` in current working directory
func setupContextInWd() (*context.Context, error) {
	wd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	return context.NewContext(wd)
}

func noColorFlag() cli.Flag {
	return cli.BoolFlag{
		Name:   "no-color",
		EnvVar: "NO_COLOR,QUICKHOOK_NO_COLOR",
		Usage:  "Don't colorize output",
	}
}

func allFlag() cli.Flag {
	return cli.BoolFlag{
		Name:  "all, a",
		Usage: "Run on all Git-tracked files",
	}
}

func filesFlag() cli.Flag {
	return cli.BoolFlag{
		Name:  "files, F",
		Usage: "Run on the given comma-separated list of files",
	}
}

func yesFlag() cli.Flag {
	return cli.BoolFlag{
		Name:  "yes, y",
		Usage: "Assume yes for all prompts",
	}
}
<commit_msg>Bump version to v1.3.0<commit_after>package main

import (
	"fmt"
	"os"

	"github.com/urfave/cli"

	"github.com/dirk/quickhook/context"
	"github.com/dirk/quickhook/hooks"
)

const VERSION = "1.3.0"

func main() {
	context, err := setupContextInWd()
	if err != nil {
		panic(err)
	}

	app := cli.NewApp()
	app.Name = "quickhook"
	app.Version = VERSION
	app.Usage = "Git hook runner"

	app.Commands = []cli.Command{
		{
			Name:  "hook",
			Usage: "Run a hook",
			Action: func(c *cli.Context) error {
				return cli.ShowSubcommandHelp(c)
			},
			Subcommands: []cli.Command{
				cli.Command{
					Name: "pre-commit",
					Flags: []cli.Flag{
						allFlag(),
						filesFlag(),
						noColorFlag(),
					},
					Action: func(c *cli.Context) error {
						files := []string{}
						if c.Bool("files") {
							files = c.Args()
						}

						err := hooks.PreCommit(context, &hooks.PreCommitOpts{
							All:     c.Bool("all"),
							Files:   files,
							NoColor: c.Bool("no-color"),
						})
						if err != nil {
							panic(err)
						}
						return nil
					},
				},
				cli.Command{
					Name: "commit-msg",
					Flags: []cli.Flag{
						noColorFlag(),
					},
					Action: func(c *cli.Context) error {
						messageTempFile := c.Args().Get(0)
						if messageTempFile == "" {
							fmt.Println("Missing message temp file argument")
							os.Exit(1)
						}

						err := hooks.CommitMsg(context, &hooks.CommitMsgOpts{
							NoColor:         c.Bool("no-color"),
							MessageTempFile: messageTempFile,
						})
						if err != nil {
							panic(err)
						}
						return nil
					},
				},
			},
		},
		{
			Name:      "install",
			Usage:     "Install Quickhook shims into .git/hooks",
			ArgsUsage: " ", // Don't show "[arguments...]"
			Flags: []cli.Flag{
				yesFlag(),
			},
			Action: func(c *cli.Context) error {
				prompt := c.Bool("yes") != true

				err := Install(context, prompt)
				if err != nil {
					panic(err)
				}
				return nil
			},
		},
	}

	app.Run(os.Args)
}

// Set up `Context` in current working directory
func setupContextInWd() (*context.Context, error) {
	wd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	return context.NewContext(wd)
}

func noColorFlag() cli.Flag {
	return cli.BoolFlag{
		Name:   "no-color",
		EnvVar: "NO_COLOR,QUICKHOOK_NO_COLOR",
		Usage:  "Don't colorize output",
	}
}

func allFlag() cli.Flag {
	return cli.BoolFlag{
		Name:  "all, a",
		Usage: "Run on all Git-tracked files",
	}
}

func filesFlag() cli.Flag {
	return cli.BoolFlag{
		Name:  "files, F",
		Usage: "Run on the given comma-separated list of files",
	}
}

func yesFlag() cli.Flag {
	return cli.BoolFlag{
		Name:  "yes, y",
		Usage: "Assume yes for all prompts",
	}
}
<|endoftext|>