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

package handler

import (
	"errors"

	"github.com/TheThingsNetwork/ttn/api"
	pb_broker "github.com/TheThingsNetwork/ttn/api/broker"
	pb "github.com/TheThingsNetwork/ttn/api/handler"
	pb_lorawan "github.com/TheThingsNetwork/ttn/api/protocol/lorawan"
	"github.com/TheThingsNetwork/ttn/core/handler/application"
	"github.com/TheThingsNetwork/ttn/core/handler/device"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

type handlerManager struct {
	*handler
}

var errf = grpc.Errorf

func (h *handlerManager) getDevice(ctx context.Context, in *pb.DeviceIdentifier) (*device.Device, error) {
	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Device Identifier")
	}
	claims, err := h.Component.ValidateContext(ctx)
	if err != nil {
		return nil, err
	}
	if !claims.CanEditApp(in.AppId) {
		return nil, errf(codes.Unauthenticated, "No access to this device")
	}
	dev, err := h.devices.Get(in.AppId, in.DevId)
	if err != nil {
		return nil, err
	}
	if !claims.CanEditApp(dev.AppID) {
		return nil, errf(codes.Unauthenticated, "No access to this device")
	}
	return dev, nil
}

func (h *handlerManager) GetDevice(ctx context.Context, in *pb.DeviceIdentifier) (*pb.Device, error) {
	dev, err := h.getDevice(ctx, in)
	if err != nil {
		return nil, err
	}

	nsDev, err := h.ttnDeviceManager.GetDevice(ctx, &pb_lorawan.DeviceIdentifier{
		AppEui: &dev.AppEUI,
		DevEui: &dev.DevEUI,
	})
	if err != nil {
		return nil, err
	}

	return &pb.Device{
		AppId: dev.AppID,
		DevId: dev.DevID,
		Device: &pb.Device_LorawanDevice{LorawanDevice: &pb_lorawan.Device{
			AppId:            dev.AppID,
			AppEui:           nsDev.AppEui,
			DevId:            dev.DevID,
			DevEui:           nsDev.DevEui,
			DevAddr:          nsDev.DevAddr,
			NwkSKey:          nsDev.NwkSKey,
			AppSKey:          &dev.AppSKey,
			AppKey:           &dev.AppKey,
			FCntUp:           nsDev.FCntUp,
			FCntDown:         nsDev.FCntDown,
			DisableFCntCheck: nsDev.DisableFCntCheck,
			Uses32BitFCnt:    nsDev.Uses32BitFCnt,
			LastSeen:         nsDev.LastSeen,
		}},
	}, nil
}

func (h *handlerManager) SetDevice(ctx context.Context, in *pb.Device) (*api.Ack, error) {
	_, err := h.getDevice(ctx, &pb.DeviceIdentifier{AppId: in.AppId, DevId: in.DevId})
	if err != nil && err != device.ErrNotFound {
		return nil, err
	}

	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Device")
	}

	updated := &device.Device{
		AppID: in.AppId,
		DevID: in.DevId,
	}

	lorawan := in.GetLorawanDevice()
	if lorawan == nil {
		err = grpcErrf(codes.InvalidArgument, "No LoRaWAN Device")
	}

	_, err = h.ttnDeviceManager.SetDevice(ctx, &pb_lorawan.Device{
		AppId:            in.AppId,
		AppEui:           lorawan.AppEui,
		DevEui:           lorawan.DevEui,
		DevAddr:          lorawan.DevAddr,
		NwkSKey:          lorawan.NwkSKey,
		FCntUp:           lorawan.FCntUp,
		FCntDown:         lorawan.FCntDown,
		DisableFCntCheck: lorawan.DisableFCntCheck,
		Uses32BitFCnt:    lorawan.Uses32BitFCnt,
	})
	if err != nil {
		return nil, err
	}

	updated.AppEUI = *lorawan.AppEui
	updated.DevEUI = *lorawan.DevEui

	if lorawan.DevAddr != nil && lorawan.NwkSKey != nil && lorawan.AppSKey != nil {
		updated.DevAddr = *lorawan.DevAddr
		updated.NwkSKey = *lorawan.NwkSKey
		updated.AppSKey = *lorawan.AppSKey
	}

	if lorawan.AppKey != nil {
		updated.AppKey = *lorawan.AppKey
	}

	err = h.devices.Set(updated)
	if err != nil {
		return nil, err
	}

	return &api.Ack{}, nil
}

func (h *handlerManager) DeleteDevice(ctx context.Context, in *pb.DeviceIdentifier) (*api.Ack, error) {
	dev, err := h.getDevice(ctx, in)
	if err != nil {
		return nil, err
	}
	_, err = h.ttnDeviceManager.DeleteDevice(ctx, &pb_lorawan.DeviceIdentifier{AppEui: &dev.AppEUI, DevEui: &dev.DevEUI})
	if err != nil {
		return nil, err
	}
	err = h.devices.Delete(in.AppId, in.DevId)
	if err != nil {
		return nil, err
	}
	return &api.Ack{}, nil
}

func (h *handlerManager) GetDevicesForApplication(ctx context.Context, in *pb.ApplicationIdentifier) (*pb.DeviceList, error) {
	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Application Identifier")
	}
	claims, err := h.Component.ValidateContext(ctx)
	if err != nil {
		return nil, err
	}
	if !claims.CanEditApp(in.AppId) {
		return nil, errf(codes.Unauthenticated, "No access to this application")
	}
	devices, err := h.devices.ListForApp(in.AppId)
	if err != nil {
		return nil, err
	}
	res := &pb.DeviceList{Devices: []*pb.Device{}}
	for _, dev := range devices {
		res.Devices = append(res.Devices, &pb.Device{
			AppId: dev.AppID,
			DevId: dev.DevID,
			Device: &pb.Device_LorawanDevice{LorawanDevice: &pb_lorawan.Device{
				AppId:   dev.AppID,
				AppEui:  &dev.AppEUI,
				DevId:   dev.DevID,
				DevEui:  &dev.DevEUI,
				DevAddr: &dev.DevAddr,
				NwkSKey: &dev.NwkSKey,
				AppSKey: &dev.AppSKey,
				AppKey:  &dev.AppKey,
			}},
		})
	}
	return res, nil
}

func (h *handlerManager) getApplication(ctx context.Context, in *pb.ApplicationIdentifier) (*application.Application, error) {
	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Application Identifier")
	}
	claims, err := h.Component.ValidateContext(ctx)
	if err != nil {
		return nil, err
	}
	if !claims.CanEditApp(in.AppId) {
		return nil, errf(codes.Unauthenticated, "No access to this application")
	}
	app, err := h.applications.Get(in.AppId)
	if err != nil {
		return nil, err
	}
	if !claims.CanEditApp(app.AppID) {
		return nil, errf(codes.Unauthenticated, "No access to this application")
	}
	return app, nil
}

func (h *handlerManager) GetApplication(ctx context.Context, in *pb.ApplicationIdentifier) (*pb.Application, error) {
	app, err := h.getApplication(ctx, in)
	if err != nil {
		return nil, err
	}

	return &pb.Application{
		AppId:     app.AppID,
		Decoder:   app.Decoder,
		Converter: app.Converter,
		Validator: app.Validator,
	}, nil
}

func (h *handlerManager) SetApplication(ctx context.Context, in *pb.Application) (*api.Ack, error) {
	app, err := h.getApplication(ctx, &pb.ApplicationIdentifier{AppId: in.AppId})
	if err != nil && err != application.ErrNotFound {
		return nil, err
	}

	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Application")
	}

	err = h.applications.Set(&application.Application{
		AppID:     in.AppId,
		Decoder:   in.Decoder,
		Converter: in.Converter,
		Validator: in.Validator,
	})
	if err != nil {
		return nil, err
	}

	if app == nil {
		// Add this application ID to the cache if needed
		h.applicationIdsLock.Lock()
		var alreadyInCache bool
		for _, id := range h.applicationIds {
			if id == in.AppId {
				alreadyInCache = true
				break
			}
		}
		if !alreadyInCache {
			h.applicationIds = append(h.applicationIds, in.AppId)
		}
		h.applicationIdsLock.Unlock()

		// If we had to add it, we also have to announce it to the Discovery and Broker
		if !alreadyInCache {
			h.announce()
			_, err := h.ttnBrokerManager.RegisterApplicationHandler(ctx, &pb_broker.ApplicationHandlerRegistration{
				AppId:     in.AppId,
				HandlerId: h.Identity.Id,
			})
			if err != nil {
				h.Ctx.WithField("AppID", in.AppId).WithError(err).Warn("Could not register Application with Broker")
			}
		}
	}

	return &api.Ack{}, nil
}

func (h *handlerManager) DeleteApplication(ctx context.Context, in *pb.ApplicationIdentifier) (*api.Ack, error) {
	_, err := h.getApplication(ctx, in)
	if err != nil {
		return nil, err
	}

	err = h.applications.Delete(in.AppId)
	if err != nil {
		return nil, err
	}

	// Remove this application ID from the cache
	h.applicationIdsLock.Lock()
	newApplicationIDList := make([]string, 0, len(h.applicationIds))
	for _, id := range h.applicationIds {
		if id != in.AppId {
			newApplicationIDList = append(newApplicationIDList, id)
		}
	}
	h.applicationIds = newApplicationIDList
	h.applicationIdsLock.Unlock()

	return &api.Ack{}, nil
}

func (h *handlerManager) GetStatus(ctx context.Context, in *pb.StatusRequest) (*pb.Status, error) {
	return nil, errors.New("Not Implemented")
}

func (b *handler) RegisterManager(s *grpc.Server) {
	server := &handlerManager{b}
	pb.RegisterHandlerManagerServer(s, server)
	pb.RegisterApplicationManagerServer(s, server)
}
<commit_msg>Changes in Device management (handler)<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package handler

import (
	"errors"

	"github.com/TheThingsNetwork/ttn/api"
	pb_broker "github.com/TheThingsNetwork/ttn/api/broker"
	pb "github.com/TheThingsNetwork/ttn/api/handler"
	pb_lorawan "github.com/TheThingsNetwork/ttn/api/protocol/lorawan"
	"github.com/TheThingsNetwork/ttn/core/handler/application"
	"github.com/TheThingsNetwork/ttn/core/handler/device"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

type handlerManager struct {
	*handler
}

var errf = grpc.Errorf

func (h *handlerManager) getDevice(ctx context.Context, in *pb.DeviceIdentifier) (*device.Device, error) {
	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Device Identifier")
	}
	claims, err := h.Component.ValidateContext(ctx)
	if err != nil {
		return nil, err
	}
	if !claims.CanEditApp(in.AppId) {
		return nil, errf(codes.Unauthenticated, "No access to this device")
	}
	dev, err := h.devices.Get(in.AppId, in.DevId)
	if err != nil {
		return nil, err
	}
	if !claims.CanEditApp(dev.AppID) {
		return nil, errf(codes.Unauthenticated, "No access to this device")
	}
	return dev, nil
}

func (h *handlerManager) GetDevice(ctx context.Context, in *pb.DeviceIdentifier) (*pb.Device, error) {
	dev, err := h.getDevice(ctx, in)
	if err != nil {
		return nil, err
	}

	nsDev, err := h.ttnDeviceManager.GetDevice(ctx, &pb_lorawan.DeviceIdentifier{
		AppEui: &dev.AppEUI,
		DevEui: &dev.DevEUI,
	})
	if err != nil {
		return nil, err
	}

	return &pb.Device{
		AppId: dev.AppID,
		DevId: dev.DevID,
		Device: &pb.Device_LorawanDevice{LorawanDevice: &pb_lorawan.Device{
			AppId:            dev.AppID,
			AppEui:           nsDev.AppEui,
			DevId:            dev.DevID,
			DevEui:           nsDev.DevEui,
			DevAddr:          nsDev.DevAddr,
			NwkSKey:          nsDev.NwkSKey,
			AppSKey:          &dev.AppSKey,
			AppKey:           &dev.AppKey,
			FCntUp:           nsDev.FCntUp,
			FCntDown:         nsDev.FCntDown,
			DisableFCntCheck: nsDev.DisableFCntCheck,
			Uses32BitFCnt:    nsDev.Uses32BitFCnt,
			LastSeen:         nsDev.LastSeen,
		}},
	}, nil
}

func (h *handlerManager) SetDevice(ctx context.Context, in *pb.Device) (*api.Ack, error) {
	dev, err := h.getDevice(ctx, &pb.DeviceIdentifier{AppId: in.AppId, DevId: in.DevId})
	if err != nil && err != device.ErrNotFound {
		return nil, err
	}

	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Device")
	}

	lorawan := in.GetLorawanDevice()
	if lorawan == nil {
		return nil, grpcErrf(codes.InvalidArgument, "No LoRaWAN Device")
	}

	if dev != nil { // When this is an update
		if dev.AppEUI != *lorawan.AppEui || dev.DevEUI != *lorawan.DevEui {
			// If the AppEUI or DevEUI is changed, we should remove the device from the NetworkServer and re-add it later
			_, err = h.ttnDeviceManager.DeleteDevice(ctx, &pb_lorawan.DeviceIdentifier{
				AppEui: &dev.AppEUI,
				DevEui: &dev.DevEUI,
			})
			if err != nil {
				return nil, err
			}
		}
	} else { // When this is a create

	}

	updated := &device.Device{
		AppID:  in.AppId,
		DevID:  in.DevId,
		AppEUI: *lorawan.AppEui,
		DevEUI: *lorawan.DevEui,
	}

	if lorawan.DevAddr != nil {
		updated.DevAddr = *lorawan.DevAddr
	}
	if lorawan.NwkSKey != nil {
		updated.NwkSKey = *lorawan.NwkSKey
	}
	if lorawan.AppSKey != nil {
		updated.AppSKey = *lorawan.AppSKey
	}
	if lorawan.AppKey != nil {
		updated.AppKey = *lorawan.AppKey
	}

	nsUpdated := &pb_lorawan.Device{
		AppId:            in.AppId,
		DevId:            in.DevId,
		AppEui:           lorawan.AppEui,
		DevEui:           lorawan.DevEui,
		DevAddr:          lorawan.DevAddr,
		NwkSKey:          lorawan.NwkSKey,
		FCntUp:           lorawan.FCntUp,
		FCntDown:         lorawan.FCntDown,
		DisableFCntCheck: lorawan.DisableFCntCheck,
		Uses32BitFCnt:    lorawan.Uses32BitFCnt,
	}

	_, err = h.ttnDeviceManager.SetDevice(ctx, nsUpdated)
	if err != nil {
		return nil, err
	}

	err = h.devices.Set(updated)
	if err != nil {
		return nil, err
	}

	return &api.Ack{}, nil
}

func (h *handlerManager) DeleteDevice(ctx context.Context, in *pb.DeviceIdentifier) (*api.Ack, error) {
	dev, err := h.getDevice(ctx, in)
	if err != nil {
		return nil, err
	}
	_, err = h.ttnDeviceManager.DeleteDevice(ctx, &pb_lorawan.DeviceIdentifier{AppEui: &dev.AppEUI, DevEui: &dev.DevEUI})
	if err != nil {
		return nil, err
	}
	err = h.devices.Delete(in.AppId, in.DevId)
	if err != nil {
		return nil, err
	}
	return &api.Ack{}, nil
}

func (h *handlerManager) GetDevicesForApplication(ctx context.Context, in *pb.ApplicationIdentifier) (*pb.DeviceList, error) {
	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Application Identifier")
	}
	claims, err := h.Component.ValidateContext(ctx)
	if err != nil {
		return nil, err
	}
	if !claims.CanEditApp(in.AppId) {
		return nil, errf(codes.Unauthenticated, "No access to this application")
	}
	devices, err := h.devices.ListForApp(in.AppId)
	if err != nil {
		return nil, err
	}
	res := &pb.DeviceList{Devices: []*pb.Device{}}
	for _, dev := range devices {
		res.Devices = append(res.Devices, &pb.Device{
			AppId: dev.AppID,
			DevId: dev.DevID,
			Device: &pb.Device_LorawanDevice{LorawanDevice: &pb_lorawan.Device{
				AppId:   dev.AppID,
				AppEui:  &dev.AppEUI,
				DevId:   dev.DevID,
				DevEui:  &dev.DevEUI,
				DevAddr: &dev.DevAddr,
				NwkSKey: &dev.NwkSKey,
				AppSKey: &dev.AppSKey,
				AppKey:  &dev.AppKey,
			}},
		})
	}
	return res, nil
}

func (h *handlerManager) getApplication(ctx context.Context, in *pb.ApplicationIdentifier) (*application.Application, error) {
	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Application Identifier")
	}
	claims, err := h.Component.ValidateContext(ctx)
	if err != nil {
		return nil, err
	}
	if !claims.CanEditApp(in.AppId) {
		return nil, errf(codes.Unauthenticated, "No access to this application")
	}
	app, err := h.applications.Get(in.AppId)
	if err != nil {
		return nil, err
	}
	if !claims.CanEditApp(app.AppID) {
		return nil, errf(codes.Unauthenticated, "No access to this application")
	}
	return app, nil
}

func (h *handlerManager) GetApplication(ctx context.Context, in *pb.ApplicationIdentifier) (*pb.Application, error) {
	app, err := h.getApplication(ctx, in)
	if err != nil {
		return nil, err
	}

	return &pb.Application{
		AppId:     app.AppID,
		Decoder:   app.Decoder,
		Converter: app.Converter,
		Validator: app.Validator,
	}, nil
}

func (h *handlerManager) SetApplication(ctx context.Context, in *pb.Application) (*api.Ack, error) {
	app, err := h.getApplication(ctx, &pb.ApplicationIdentifier{AppId: in.AppId})
	if err != nil && err != application.ErrNotFound {
		return nil, err
	}

	if !in.Validate() {
		return nil, grpcErrf(codes.InvalidArgument, "Invalid Application")
	}

	err = h.applications.Set(&application.Application{
		AppID:     in.AppId,
		Decoder:   in.Decoder,
		Converter: in.Converter,
		Validator: in.Validator,
	})
	if err != nil {
		return nil, err
	}

	if app == nil {
		// Add this application ID to the cache if needed
		h.applicationIdsLock.Lock()
		var alreadyInCache bool
		for _, id := range h.applicationIds {
			if id == in.AppId {
				alreadyInCache = true
				break
			}
		}
		if !alreadyInCache {
			h.applicationIds = append(h.applicationIds, in.AppId)
		}
		h.applicationIdsLock.Unlock()

		// If we had to add it, we also have to announce it to the Discovery and Broker
		if !alreadyInCache {
			h.announce()
			_, err := h.ttnBrokerManager.RegisterApplicationHandler(ctx, &pb_broker.ApplicationHandlerRegistration{
				AppId:     in.AppId,
				HandlerId: h.Identity.Id,
			})
			if err != nil {
				h.Ctx.WithField("AppID", in.AppId).WithError(err).Warn("Could not register Application with Broker")
			}
		}
	}

	return &api.Ack{}, nil
}

func (h *handlerManager) DeleteApplication(ctx context.Context, in *pb.ApplicationIdentifier) (*api.Ack, error) {
	_, err := h.getApplication(ctx, in)
	if err != nil {
		return nil, err
	}

	// Get and delete all devices for this application
	devices, err := h.devices.ListForApp(in.AppId)
	if err != nil {
		return nil, err
	}
	for _, dev := range devices {
		_, err = h.ttnDeviceManager.DeleteDevice(ctx, &pb_lorawan.DeviceIdentifier{AppEui: &dev.AppEUI, DevEui: &dev.DevEUI})
		if err != nil {
			return nil, err
		}
		err = h.devices.Delete(dev.AppID, dev.DevID)
		if err != nil {
			return nil, err
		}
	}

	// Delete the Application
	err = h.applications.Delete(in.AppId)
	if err != nil {
		return nil, err
	}

	// Remove this application ID from the cache
	h.applicationIdsLock.Lock()
	newApplicationIDList := make([]string, 0, len(h.applicationIds))
	for _, id := range h.applicationIds {
		if id != in.AppId {
			newApplicationIDList = append(newApplicationIDList, id)
		}
	}
	h.applicationIds = newApplicationIDList
	h.applicationIdsLock.Unlock()

	return &api.Ack{}, nil
}

func (h *handlerManager) GetStatus(ctx context.Context, in *pb.StatusRequest) (*pb.Status, error) {
	return nil, errors.New("Not Implemented")
}

func (b *handler) RegisterManager(s *grpc.Server) {
	server := &handlerManager{b}
	pb.RegisterHandlerManagerServer(s, server)
	pb.RegisterApplicationManagerServer(s, server)
}
<|endoftext|>


<commit_before>package metadata

import (
	"context"
	"encoding/binary"
	"fmt"
	"strings"
	"time"

	"github.com/boltdb/bolt"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/filters"
	"github.com/containerd/containerd/images"
	"github.com/containerd/containerd/labels"
	"github.com/containerd/containerd/metadata/boltutil"
	"github.com/containerd/containerd/namespaces"
	digest "github.com/opencontainers/go-digest"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
)

type imageStore struct {
	tx *bolt.Tx
}

// NewImageStore returns a store backed by a bolt DB
func NewImageStore(tx *bolt.Tx) images.Store {
	return &imageStore{tx: tx}
}

func (s *imageStore) Get(ctx context.Context, name string) (images.Image, error) {
	var image images.Image

	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return images.Image{}, err
	}

	bkt := getImagesBucket(s.tx, namespace)
	if bkt == nil {
		return images.Image{}, errors.Wrapf(errdefs.ErrNotFound, "image %q", name)
	}

	ibkt := bkt.Bucket([]byte(name))
	if ibkt == nil {
		return images.Image{}, errors.Wrapf(errdefs.ErrNotFound, "image %q", name)
	}

	image.Name = name
	if err := readImage(&image, ibkt); err != nil {
		return images.Image{}, errors.Wrapf(err, "image %q", name)
	}

	return image, nil
}

func (s *imageStore) List(ctx context.Context, fs ...string) ([]images.Image, error) {
	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return nil, err
	}

	filter, err := filters.ParseAll(fs...)
	if err != nil {
		return nil, errors.Wrapf(errdefs.ErrInvalidArgument, err.Error())
	}

	bkt := getImagesBucket(s.tx, namespace)
	if bkt == nil {
		return nil, nil // empty store
	}

	var m []images.Image
	if err := bkt.ForEach(func(k, v []byte) error {
		var (
			image = images.Image{
				Name: string(k),
			}
			kbkt = bkt.Bucket(k)
		)

		if err := readImage(&image, kbkt); err != nil {
			return err
		}

		if filter.Match(adaptImage(image)) {
			m = append(m, image)
		}
		return nil
	}); err != nil {
		return nil, err
	}

	return m, nil
}

func (s *imageStore) Create(ctx context.Context, image images.Image) (images.Image, error) {
	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return images.Image{}, err
	}

	if image.Name == "" {
		return images.Image{}, errors.Wrapf(errdefs.ErrInvalidArgument, "image name is required for create")
	}

	if err := validateImage(&image); err != nil {
		return images.Image{}, err
	}

	return image, withImagesBucket(s.tx, namespace, func(bkt *bolt.Bucket) error {
		ibkt, err := bkt.CreateBucket([]byte(image.Name))
		if err != nil {
			if err != bolt.ErrBucketExists {
				return err
			}

			return errors.Wrapf(errdefs.ErrAlreadyExists, "image %q", image.Name)
		}

		image.CreatedAt = time.Now().UTC()
		image.UpdatedAt = image.CreatedAt
		return writeImage(ibkt, &image)
	})
}

func (s *imageStore) Update(ctx context.Context, image images.Image, fieldpaths ...string) (images.Image, error) {
	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return images.Image{}, err
	}

	if image.Name == "" {
		return images.Image{}, errors.Wrapf(errdefs.ErrInvalidArgument, "image name is required for update")
	}

	var updated images.Image
	return updated, withImagesBucket(s.tx, namespace, func(bkt *bolt.Bucket) error {
		ibkt := bkt.Bucket([]byte(image.Name))
		if ibkt == nil {
			return errors.Wrapf(errdefs.ErrNotFound, "image %q", image.Name)
		}

		if err := readImage(&updated, ibkt); err != nil {
			return errors.Wrapf(err, "image %q", image.Name)
		}
		createdat := updated.CreatedAt
		updated.Name = image.Name

		if len(fieldpaths) > 0 {
			for _, path := range fieldpaths {
				if strings.HasPrefix(path, "labels.") {
					if updated.Labels == nil {
						updated.Labels = map[string]string{}
					}

					key := strings.TrimPrefix(path, "labels.")
					updated.Labels[key] = image.Labels[key]
					continue
				}

				switch path {
				case "labels":
					updated.Labels = image.Labels
				case "target":
					// NOTE(stevvooe): While we allow setting individual labels, we
					// only support replacing the target as a unit, since that is
					// commonly pulled as a unit from other sources. It often doesn't
					// make sense to modify the size or digest without touching the
					// mediatype, as well, for example.
					updated.Target = image.Target
				default:
					return errors.Wrapf(errdefs.ErrInvalidArgument, "cannot update %q field on image %q", path, image.Name)
				}
			}
		} else {
			updated = image
		}

		if err := validateImage(&image); err != nil {
			return err
		}

		updated.CreatedAt = createdat
		updated.UpdatedAt = time.Now().UTC()
		return writeImage(ibkt, &updated)
	})
}

func (s *imageStore) Delete(ctx context.Context, name string) error {
	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return err
	}

	return withImagesBucket(s.tx, namespace, func(bkt *bolt.Bucket) error {
		err := bkt.DeleteBucket([]byte(name))
		if err == bolt.ErrBucketNotFound {
			return errors.Wrapf(errdefs.ErrNotFound, "image %q", name)
		}
		return err
	})
}

func validateImage(image *images.Image) error {
	if image.Name == "" {
		return errors.Wrapf(errdefs.ErrInvalidArgument, "image name must not be empty")
	}

	for k, v := range image.Labels {
		if err := labels.Validate(k, v); err != nil {
			return errors.Wrapf(err, "image.Labels")
		}
	}

	return validateTarget(&image.Target)
}

func validateTarget(target *ocispec.Descriptor) error {
	// NOTE(stevvooe): Only validate fields we actually store.

	if err := target.Digest.Validate(); err != nil {
		return errors.Wrapf(errdefs.ErrInvalidArgument, "Target.Digest %q invalid: %v", target.Digest, err)
	}

	if target.Size <= 0 {
		return errors.Wrapf(errdefs.ErrInvalidArgument, "Target.Size must be greater than zero")
	}

	if target.MediaType == "" {
		return errors.Wrapf(errdefs.ErrInvalidArgument, "Target.MediaType must be set")
	}

	return nil
}

func readImage(image *images.Image, bkt *bolt.Bucket) error {
	if err := boltutil.ReadTimestamps(bkt, &image.CreatedAt, &image.UpdatedAt); err != nil {
		return err
	}

	labels, err := boltutil.ReadLabels(bkt)
	if err != nil {
		return err
	}
	image.Labels = labels

	tbkt := bkt.Bucket(bucketKeyTarget)
	if tbkt == nil {
		return errors.New("unable to read target bucket")
	}
	return tbkt.ForEach(func(k, v []byte) error {
		if v == nil {
			return nil // skip it? a bkt maybe?
		}

		// TODO(stevvooe): This is why we need to use byte values for
		// keys, rather than full arrays.
		switch string(k) {
		case string(bucketKeyDigest):
			image.Target.Digest = digest.Digest(v)
		case string(bucketKeyMediaType):
			image.Target.MediaType = string(v)
		case string(bucketKeySize):
			image.Target.Size, _ = binary.Varint(v)
		}

		return nil
	})
}

func writeImage(bkt *bolt.Bucket, image *images.Image) error {
	if err := boltutil.WriteTimestamps(bkt, image.CreatedAt, image.UpdatedAt); err != nil {
		return err
	}

	if err := boltutil.WriteLabels(bkt, image.Labels); err != nil {
		return errors.Wrapf(err, "writing labels for image %v", image.Name)
	}

	// write the target bucket
	tbkt, err := bkt.CreateBucketIfNotExists([]byte(bucketKeyTarget))
	if err != nil {
		return err
	}

	sizeEncoded, err := encodeInt(image.Target.Size)
	if err != nil {
		return err
	}

	for _, v := range [][2][]byte{
		{bucketKeyDigest, []byte(image.Target.Digest)},
		{bucketKeyMediaType, []byte(image.Target.MediaType)},
		{bucketKeySize, sizeEncoded},
	} {
		if err := tbkt.Put(v[0], v[1]); err != nil {
			return err
		}
	}

	return nil
}

func encodeInt(i int64) ([]byte, error) {
	var (
		buf      [binary.MaxVarintLen64]byte
		iEncoded = buf[:]
	)
	iEncoded = iEncoded[:binary.PutVarint(iEncoded, i)]

	if len(iEncoded) == 0 {
		return nil, fmt.Errorf("failed encoding integer = %v", i)
	}
	return iEncoded, nil
}
<commit_msg>remove redundant check for imagestore create<commit_after>package metadata

import (
	"context"
	"encoding/binary"
	"fmt"
	"strings"
	"time"

	"github.com/boltdb/bolt"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/filters"
	"github.com/containerd/containerd/images"
	"github.com/containerd/containerd/labels"
	"github.com/containerd/containerd/metadata/boltutil"
	"github.com/containerd/containerd/namespaces"
	digest "github.com/opencontainers/go-digest"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
)

type imageStore struct {
	tx *bolt.Tx
}

// NewImageStore returns a store backed by a bolt DB
func NewImageStore(tx *bolt.Tx) images.Store {
	return &imageStore{tx: tx}
}

func (s *imageStore) Get(ctx context.Context, name string) (images.Image, error) {
	var image images.Image

	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return images.Image{}, err
	}

	bkt := getImagesBucket(s.tx, namespace)
	if bkt == nil {
		return images.Image{}, errors.Wrapf(errdefs.ErrNotFound, "image %q", name)
	}

	ibkt := bkt.Bucket([]byte(name))
	if ibkt == nil {
		return images.Image{}, errors.Wrapf(errdefs.ErrNotFound, "image %q", name)
	}

	image.Name = name
	if err := readImage(&image, ibkt); err != nil {
		return images.Image{}, errors.Wrapf(err, "image %q", name)
	}

	return image, nil
}

func (s *imageStore) List(ctx context.Context, fs ...string) ([]images.Image, error) {
	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return nil, err
	}

	filter, err := filters.ParseAll(fs...)
	if err != nil {
		return nil, errors.Wrapf(errdefs.ErrInvalidArgument, err.Error())
	}

	bkt := getImagesBucket(s.tx, namespace)
	if bkt == nil {
		return nil, nil // empty store
	}

	var m []images.Image
	if err := bkt.ForEach(func(k, v []byte) error {
		var (
			image = images.Image{
				Name: string(k),
			}
			kbkt = bkt.Bucket(k)
		)

		if err := readImage(&image, kbkt); err != nil {
			return err
		}

		if filter.Match(adaptImage(image)) {
			m = append(m, image)
		}
		return nil
	}); err != nil {
		return nil, err
	}

	return m, nil
}

func (s *imageStore) Create(ctx context.Context, image images.Image) (images.Image, error) {
	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return images.Image{}, err
	}

	if err := validateImage(&image); err != nil {
		return images.Image{}, err
	}

	return image, withImagesBucket(s.tx, namespace, func(bkt *bolt.Bucket) error {
		ibkt, err := bkt.CreateBucket([]byte(image.Name))
		if err != nil {
			if err != bolt.ErrBucketExists {
				return err
			}

			return errors.Wrapf(errdefs.ErrAlreadyExists, "image %q", image.Name)
		}

		image.CreatedAt = time.Now().UTC()
		image.UpdatedAt = image.CreatedAt
		return writeImage(ibkt, &image)
	})
}

func (s *imageStore) Update(ctx context.Context, image images.Image, fieldpaths ...string) (images.Image, error) {
	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return images.Image{}, err
	}

	if image.Name == "" {
		return images.Image{}, errors.Wrapf(errdefs.ErrInvalidArgument, "image name is required for update")
	}

	var updated images.Image
	return updated, withImagesBucket(s.tx, namespace, func(bkt *bolt.Bucket) error {
		ibkt := bkt.Bucket([]byte(image.Name))
		if ibkt == nil {
			return errors.Wrapf(errdefs.ErrNotFound, "image %q", image.Name)
		}

		if err := readImage(&updated, ibkt); err != nil {
			return errors.Wrapf(err, "image %q", image.Name)
		}
		createdat := updated.CreatedAt
		updated.Name = image.Name

		if len(fieldpaths) > 0 {
			for _, path := range fieldpaths {
				if strings.HasPrefix(path, "labels.") {
					if updated.Labels == nil {
						updated.Labels = map[string]string{}
					}

					key := strings.TrimPrefix(path, "labels.")
					updated.Labels[key] = image.Labels[key]
					continue
				}

				switch path {
				case "labels":
					updated.Labels = image.Labels
				case "target":
					// NOTE(stevvooe): While we allow setting individual labels, we
					// only support replacing the target as a unit, since that is
					// commonly pulled as a unit from other sources. It often doesn't
					// make sense to modify the size or digest without touching the
					// mediatype, as well, for example.
					updated.Target = image.Target
				default:
					return errors.Wrapf(errdefs.ErrInvalidArgument, "cannot update %q field on image %q", path, image.Name)
				}
			}
		} else {
			updated = image
		}

		if err := validateImage(&updated); err != nil {
			return err
		}

		updated.CreatedAt = createdat
		updated.UpdatedAt = time.Now().UTC()
		return writeImage(ibkt, &updated)
	})
}

func (s *imageStore) Delete(ctx context.Context, name string) error {
	namespace, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return err
	}

	return withImagesBucket(s.tx, namespace, func(bkt *bolt.Bucket) error {
		err := bkt.DeleteBucket([]byte(name))
		if err == bolt.ErrBucketNotFound {
			return errors.Wrapf(errdefs.ErrNotFound, "image %q", name)
		}
		return err
	})
}

func validateImage(image *images.Image) error {
	if image.Name == "" {
		return errors.Wrapf(errdefs.ErrInvalidArgument, "image name must not be empty")
	}

	for k, v := range image.Labels {
		if err := labels.Validate(k, v); err != nil {
			return errors.Wrapf(err, "image.Labels")
		}
	}

	return validateTarget(&image.Target)
}

func validateTarget(target *ocispec.Descriptor) error {
	// NOTE(stevvooe): Only validate fields we actually store.

	if err := target.Digest.Validate(); err != nil {
		return errors.Wrapf(errdefs.ErrInvalidArgument, "Target.Digest %q invalid: %v", target.Digest, err)
	}

	if target.Size <= 0 {
		return errors.Wrapf(errdefs.ErrInvalidArgument, "Target.Size must be greater than zero")
	}

	if target.MediaType == "" {
		return errors.Wrapf(errdefs.ErrInvalidArgument, "Target.MediaType must be set")
	}

	return nil
}

func readImage(image *images.Image, bkt *bolt.Bucket) error {
	if err := boltutil.ReadTimestamps(bkt, &image.CreatedAt, &image.UpdatedAt); err != nil {
		return err
	}

	labels, err := boltutil.ReadLabels(bkt)
	if err != nil {
		return err
	}
	image.Labels = labels

	tbkt := bkt.Bucket(bucketKeyTarget)
	if tbkt == nil {
		return errors.New("unable to read target bucket")
	}
	return tbkt.ForEach(func(k, v []byte) error {
		if v == nil {
			return nil // skip it? a bkt maybe?
		}

		// TODO(stevvooe): This is why we need to use byte values for
		// keys, rather than full arrays.
		switch string(k) {
		case string(bucketKeyDigest):
			image.Target.Digest = digest.Digest(v)
		case string(bucketKeyMediaType):
			image.Target.MediaType = string(v)
		case string(bucketKeySize):
			image.Target.Size, _ = binary.Varint(v)
		}

		return nil
	})
}

func writeImage(bkt *bolt.Bucket, image *images.Image) error {
	if err := boltutil.WriteTimestamps(bkt, image.CreatedAt, image.UpdatedAt); err != nil {
		return err
	}

	if err := boltutil.WriteLabels(bkt, image.Labels); err != nil {
		return errors.Wrapf(err, "writing labels for image %v", image.Name)
	}

	// write the target bucket
	tbkt, err := bkt.CreateBucketIfNotExists([]byte(bucketKeyTarget))
	if err != nil {
		return err
	}

	sizeEncoded, err := encodeInt(image.Target.Size)
	if err != nil {
		return err
	}

	for _, v := range [][2][]byte{
		{bucketKeyDigest, []byte(image.Target.Digest)},
		{bucketKeyMediaType, []byte(image.Target.MediaType)},
		{bucketKeySize, sizeEncoded},
	} {
		if err := tbkt.Put(v[0], v[1]); err != nil {
			return err
		}
	}

	return nil
}

func encodeInt(i int64) ([]byte, error) {
	var (
		buf      [binary.MaxVarintLen64]byte
		iEncoded = buf[:]
	)
	iEncoded = iEncoded[:binary.PutVarint(iEncoded, i)]

	if len(iEncoded) == 0 {
		return nil, fmt.Errorf("failed encoding integer = %v", i)
	}
	return iEncoded, nil
}
<|endoftext|>


<commit_before>package reform_test

import (
	"errors"
	"time"

	"github.com/AlekSi/pointer"
	"github.com/enodata/faker"

	"gopkg.in/reform.v1"
	"gopkg.in/reform.v1/dialects/postgresql"
	. "gopkg.in/reform.v1/internal/test/models"
)

func (s *ReformSuite) TestInsert() {
	newEmail := faker.Internet().Email()
	person := &Person{Email: &newEmail}
	err := s.q.Insert(person)
	s.NoError(err)
	s.NotEqual(int32(0), person.ID)
	s.Equal("", person.Name)
	s.Equal(&newEmail, person.Email)
	s.WithinDuration(time.Now(), person.CreatedAt, 2*time.Second)
	s.Nil(person.UpdatedAt)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.NoError(err)
	s.Equal(person, person2)

	err = s.q.Insert(person)
	s.Error(err)
}

func (s *ReformSuite) TestInsertWithValues() {
	t := time.Now()
	newEmail := faker.Internet().Email()
	person := &Person{Email: &newEmail, CreatedAt: t, UpdatedAt: &t}
	err := s.q.Insert(person)
	s.NoError(err)
	s.NotEqual(int32(0), person.ID)
	s.Equal("", person.Name)
	s.Equal(&newEmail, person.Email)
	s.WithinDuration(t, person.CreatedAt, 2*time.Second)
	s.WithinDuration(t, *person.UpdatedAt, 2*time.Second)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.NoError(err)
	s.Equal(person, person2)

	err = s.q.Insert(person)
	s.Error(err)
}

func (s *ReformSuite) TestInsertWithPrimaryKey() {
	setIdentityInsert(s.T(), s.q, "people", true)

	newEmail := faker.Internet().Email()
	person := &Person{ID: 50, Email: &newEmail}
	err := s.q.Insert(person)
	s.NoError(err)
	s.Equal(int32(50), person.ID)
	s.Equal("", person.Name)
	s.Equal(&newEmail, person.Email)
	s.WithinDuration(time.Now(), person.CreatedAt, 2*time.Second)
	s.Nil(person.UpdatedAt)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.NoError(err)
	s.Equal(person, person2)

	err = s.q.Insert(person)
	s.Error(err)
}

func (s *ReformSuite) TestInsertWithStringPrimaryKey() {
	project := &Project{ID: "new", End: pointer.ToTime(time.Now().Truncate(24 * time.Hour))}
	err := s.q.Insert(project)
	s.NoError(err)
	s.Equal("new", project.ID)

	project2, err := s.q.FindByPrimaryKeyFrom(ProjectTable, project.ID)
	s.NoError(err)
	s.Equal(project, project2)

	err = s.q.Insert(project)
	s.Error(err)
}

func (s *ReformSuite) TestInsertIntoView() {
	pp := &PersonProject{PersonID: 1, ProjectID: "baron"}
	err := s.q.Insert(pp)
	s.NoError(err)

	err = s.q.Insert(pp)
	s.Error(err)

	s.RestartTransaction()

	pp = &PersonProject{PersonID: 1, ProjectID: "no_such_project"}
	err = s.q.Insert(pp)
	s.Error(err)
}

func (s *ReformSuite) TestInsertColumns() {
	t := time.Now()
	newEmail := faker.Internet().Email()
	person := &Person{Email: &newEmail, CreatedAt: t, UpdatedAt: &t}
	err := s.q.InsertColumns(person, "name", "email", "created_at", "updated_at")
	s.NoError(err)
	s.NotEqual(int32(0), person.ID)
	s.Equal("", person.Name)
	s.Equal((*int32)(nil), person.GroupID)
	s.Equal(&newEmail, person.Email)
	s.WithinDuration(t, person.CreatedAt, 2*time.Second)
	s.WithinDuration(t, *person.UpdatedAt, 2*time.Second)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.NoError(err)
	person.GroupID = pointer.ToInt32(65534)
	s.Equal(person, person2)

	err = s.q.Insert(person)
	s.Error(err)
}

func (s *ReformSuite) TestInsertColumnsIntoView() {
	pp := &PersonProject{PersonID: 1, ProjectID: "baron"}
	err := s.q.InsertColumns(pp, "person_id", "project_id")
	s.NoError(err)

	err = s.q.InsertColumns(pp, "person_id", "project_id")
	s.Error(err)

	s.RestartTransaction()

	pp = &PersonProject{PersonID: 1, ProjectID: "no_such_project"}
	err = s.q.InsertColumns(pp, "person_id", "project_id")
	s.Error(err)
}

func (s *ReformSuite) TestInsertMulti() {
	newEmail := faker.Internet().Email()
	newName := faker.Name().Name()
	person1, person2 := &Person{Email: &newEmail}, &Person{Name: newName}
	err := s.q.InsertMulti(person1, person2)
	s.NoError(err)

	s.Equal(int32(0), person1.ID)
	s.Equal("", person1.Name)
	s.Equal(&newEmail, person1.Email)
	s.WithinDuration(time.Now(), person1.CreatedAt, 2*time.Second)
	s.Nil(person1.UpdatedAt)

	s.Equal(int32(0), person2.ID)
	s.Equal(newName, person2.Name)
	s.Nil(person2.Email)
	s.WithinDuration(time.Now(), person2.CreatedAt, 2*time.Second)
	s.Nil(person2.UpdatedAt)
}

func (s *ReformSuite) TestInsertMultiWithPrimaryKeys() {
	setIdentityInsert(s.T(), s.q, "people", true)

	newEmail := faker.Internet().Email()
	newName := faker.Name().Name()
	person1, person2 := &Person{ID: 50, Email: &newEmail}, &Person{ID: 51, Name: newName}
	err := s.q.InsertMulti(person1, person2)
	s.NoError(err)

	s.Equal(int32(50), person1.ID)
	s.Equal("", person1.Name)
	s.Equal(&newEmail, person1.Email)
	s.WithinDuration(time.Now(), person1.CreatedAt, 2*time.Second)
	s.Nil(person1.UpdatedAt)

	s.Equal(int32(51), person2.ID)
	s.Equal(newName, person2.Name)
	s.Nil(person2.Email)
	s.WithinDuration(time.Now(), person2.CreatedAt, 2*time.Second)
	s.Nil(person2.UpdatedAt)

	person, err := s.q.FindByPrimaryKeyFrom(PersonTable, person1.ID)
	s.NoError(err)
	s.Equal(person1, person)

	person, err = s.q.FindByPrimaryKeyFrom(PersonTable, person2.ID)
	s.NoError(err)
	s.Equal(person2, person)
}

func (s *ReformSuite) TestInsertMultiMixes() {
	err := s.q.InsertMulti()
	s.NoError(err)

	err = s.q.InsertMulti(&Person{}, &Project{})
	s.Error(err)

	err = s.q.InsertMulti(&Person{ID: 1}, &Person{})
	s.Error(err)
}

func (s *ReformSuite) TestInsertIDOnly() {
	id := &IDOnly{}
	err := s.q.Insert(id)
	s.NoError(err)
	s.Equal(int32(1), id.ID)
}

func (s *ReformSuite) TestUpdate() {
	var person Person
	err := s.q.Update(&person)
	s.Equal(reform.ErrNoPK, err)

	person.ID = 99
	err = s.q.Update(&person)
	s.Equal(reform.ErrNoRows, err)

	err = s.q.FindByPrimaryKeyTo(&person, 102)
	s.NoError(err)

	person.Email = pointer.ToString(faker.Internet().Email())
	err = s.q.Update(&person)
	s.NoError(err)
	s.Equal(personCreated, person.CreatedAt)
	s.Require().NotNil(person.UpdatedAt)
	s.WithinDuration(time.Now(), *person.UpdatedAt, 2*time.Second)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.NoError(err)
	s.Equal(&person, person2)
}

func (s *ReformSuite) TestUpdateOverwrite() {
	newEmail := faker.Internet().Email()
	person := Person{ID: 102, Email: pointer.ToString(newEmail)}
	err := s.q.Update(&person)
	s.NoError(err)

	var person2 Person
	err = s.q.FindByPrimaryKeyTo(&person2, person.ID)
	s.NoError(err)
	s.Equal("", person2.Name)
	s.Equal(&newEmail, person2.Email)
	s.WithinDuration(time.Now(), person2.CreatedAt, 2*time.Second)
	s.Require().NotNil(person2.UpdatedAt)
	s.WithinDuration(time.Now(), *person2.UpdatedAt, 2*time.Second)
}

func (s *ReformSuite) TestUpdateColumns() {
	newName := faker.Name().Name()
	newEmail := faker.Internet().Email()

	for p, columns := range map[*Person][]string{
		&Person{Name: "Elfrieda Abbott", Email: &newEmail}:                             {"email", "updated_at"},
		&Person{Name: newName, Email: pointer.ToString("elfrieda_abbott@example.org")}: {"name", "name", "updated_at"},
		&Person{Name: newName, Email: &newEmail}:                                       {"name", "email", "updated_at"},
	} {
		var person Person
		err := s.q.FindByPrimaryKeyTo(&person, 102)
		s.NoError(err)

		person.Name = p.Name
		person.Email = p.Email
		err = s.q.UpdateColumns(&person, columns...)
		s.NoError(err)
		s.Equal(personCreated, person.CreatedAt)
		s.Require().NotNil(person.UpdatedAt)
		s.WithinDuration(time.Now(), *person.UpdatedAt, 2*time.Second)

		person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
		s.NoError(err)
		s.Equal(&person, person2)

		s.RestartTransaction()
	}

	person := &Person{ID: 102, Name: newName, Email: &newEmail, CreatedAt: personCreated}
	for e, columns := range map[error][]string{
		errors.New("reform: unexpected columns: [foo]"):     {"foo"},
		errors.New("reform: will not update PK column: id"): {"id"},
		errors.New("reform: nothing to update"):             {},
	} {
		err := s.q.UpdateColumns(person, columns...)
		s.Error(err)
		s.Equal(e, err)
	}
}

func (s *ReformSuite) TestSave() {
	newName := faker.Name().Name()
	person := &Person{Name: newName}
	err := s.q.Save(person)
	s.NoError(err)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.Require().NoError(err)
	s.Equal(newName, person2.(*Person).Name)
	s.Nil(person2.(*Person).Email)
	s.Equal(person, person2)

	newEmail := faker.Internet().Email()
	person.Email = &newEmail
	err = s.q.Save(person)
	s.NoError(err)

	person2, err = s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.Require().NoError(err)
	s.Equal(newName, person2.(*Person).Name)
	s.Equal(&newEmail, person2.(*Person).Email)
	s.Equal(person, person2)
}

func (s *ReformSuite) TestDelete() {
	person := &Person{ID: 1}
	err := s.q.Delete(person)
	s.NoError(err)
	err = s.q.Reload(person)
	s.Equal(reform.ErrNoRows, err)

	project := &Project{ID: "baron"}
	err = s.q.Delete(project)
	s.NoError(err)
	err = s.q.Reload(project)
	s.Equal(reform.ErrNoRows, err)

	project = &Project{}
	err = s.q.Delete(project)
	s.Equal(reform.ErrNoPK, err)

	project = &Project{ID: "no_such_project"}
	err = s.q.Delete(project)
	s.Equal(reform.ErrNoRows, err)
}

func (s *ReformSuite) TestDeleteFrom() {
	ra, err := s.q.DeleteFrom(PersonTable, "WHERE email IS NULL")
	s.NoError(err)
	s.Equal(uint(3), ra)

	ra, err = s.q.DeleteFrom(PersonTable, "WHERE email IS NULL")
	s.NoError(err)
	s.Equal(uint(0), ra)

	// -1 second for SQLite3, otherwise it also deletes queen itself ¯\_(ツ)_/¯
	ra, err = s.q.DeleteFrom(ProjectTable, "WHERE start < "+s.q.Placeholder(1), queenStart.Add(-time.Second))
	s.NoError(err)
	s.Equal(uint(3), ra)

	ra, err = s.q.DeleteFrom(ProjectTable, "")
	s.NoError(err)
	s.Equal(uint(2), ra)

	ra, err = s.q.DeleteFrom(ProjectTable, "WHERE invalid_tail")
	s.Error(err)
	s.Equal(uint(0), ra)
}

func (s *ReformSuite) TestCommandsSchema() {
	if s.q.Dialect != postgresql.Dialect {
		s.T().Skip("only PostgreSQL supports schemas")
	}

	legacyPerson := &LegacyPerson{Name: pointer.ToString(faker.Name().Name())}
	err := s.q.Save(legacyPerson)
	s.NoError(err)
	err = s.q.Save(legacyPerson)
	s.NoError(err)
	err = s.q.Delete(legacyPerson)
	s.NoError(err)
}
<commit_msg>Fix bogus test.<commit_after>package reform_test

import (
	"errors"
	"time"

	"github.com/AlekSi/pointer"
	"github.com/enodata/faker"

	"gopkg.in/reform.v1"
	"gopkg.in/reform.v1/dialects/postgresql"
	. "gopkg.in/reform.v1/internal/test/models"
)

func (s *ReformSuite) TestInsert() {
	newEmail := faker.Internet().Email()
	person := &Person{Email: &newEmail}
	err := s.q.Insert(person)
	s.NoError(err)
	s.NotEqual(int32(0), person.ID)
	s.Equal("", person.Name)
	s.Equal(&newEmail, person.Email)
	s.WithinDuration(time.Now(), person.CreatedAt, 2*time.Second)
	s.Nil(person.UpdatedAt)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.NoError(err)
	s.Equal(person, person2)

	err = s.q.Insert(person)
	s.Error(err)
}

func (s *ReformSuite) TestInsertWithValues() {
	t := time.Now()
	newEmail := faker.Internet().Email()
	person := &Person{Email: &newEmail, CreatedAt: t, UpdatedAt: &t}
	err := s.q.Insert(person)
	s.NoError(err)
	s.NotEqual(int32(0), person.ID)
	s.Equal("", person.Name)
	s.Equal(&newEmail, person.Email)
	s.WithinDuration(t, person.CreatedAt, 2*time.Second)
	s.WithinDuration(t, *person.UpdatedAt, 2*time.Second)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.NoError(err)
	s.Equal(person, person2)

	err = s.q.Insert(person)
	s.Error(err)
}

func (s *ReformSuite) TestInsertWithPrimaryKey() {
	setIdentityInsert(s.T(), s.q, "people", true)

	newEmail := faker.Internet().Email()
	person := &Person{ID: 50, Email: &newEmail}
	err := s.q.Insert(person)
	s.NoError(err)
	s.Equal(int32(50), person.ID)
	s.Equal("", person.Name)
	s.Equal(&newEmail, person.Email)
	s.WithinDuration(time.Now(), person.CreatedAt, 2*time.Second)
	s.Nil(person.UpdatedAt)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.NoError(err)
	s.Equal(person, person2)

	err = s.q.Insert(person)
	s.Error(err)
}

func (s *ReformSuite) TestInsertWithStringPrimaryKey() {
	today := time.Now().Truncate(24 * time.Hour)
	project := &Project{ID: "new", Start: today, End: pointer.ToTime(today.AddDate(0, 0, 1))}
	err := s.q.Insert(project)
	s.NoError(err)
	s.Equal("new", project.ID)

	project2, err := s.q.FindByPrimaryKeyFrom(ProjectTable, project.ID)
	s.NoError(err)
	s.Equal(project, project2)

	err = s.q.Insert(project)
	s.Error(err)
}

func (s *ReformSuite) TestInsertIntoView() {
	pp := &PersonProject{PersonID: 1, ProjectID: "baron"}
	err := s.q.Insert(pp)
	s.NoError(err)

	err = s.q.Insert(pp)
	s.Error(err)

	s.RestartTransaction()

	pp = &PersonProject{PersonID: 1, ProjectID: "no_such_project"}
	err = s.q.Insert(pp)
	s.Error(err)
}

func (s *ReformSuite) TestInsertColumns() {
	t := time.Now()
	newEmail := faker.Internet().Email()
	person := &Person{Email: &newEmail, CreatedAt: t, UpdatedAt: &t}
	err := s.q.InsertColumns(person, "name", "email", "created_at", "updated_at")
	s.NoError(err)
	s.NotEqual(int32(0), person.ID)
	s.Equal("", person.Name)
	s.Equal((*int32)(nil), person.GroupID)
	s.Equal(&newEmail, person.Email)
	s.WithinDuration(t, person.CreatedAt, 2*time.Second)
	s.WithinDuration(t, *person.UpdatedAt, 2*time.Second)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.NoError(err)
	person.GroupID = pointer.ToInt32(65534)
	s.Equal(person, person2)

	err = s.q.Insert(person)
	s.Error(err)
}

func (s *ReformSuite) TestInsertColumnsIntoView() {
	pp := &PersonProject{PersonID: 1, ProjectID: "baron"}
	err := s.q.InsertColumns(pp, "person_id", "project_id")
	s.NoError(err)

	err = s.q.InsertColumns(pp, "person_id", "project_id")
	s.Error(err)

	s.RestartTransaction()

	pp = &PersonProject{PersonID: 1, ProjectID: "no_such_project"}
	err = s.q.InsertColumns(pp, "person_id", "project_id")
	s.Error(err)
}

func (s *ReformSuite) TestInsertMulti() {
	newEmail := faker.Internet().Email()
	newName := faker.Name().Name()
	person1, person2 := &Person{Email: &newEmail}, &Person{Name: newName}
	err := s.q.InsertMulti(person1, person2)
	s.NoError(err)

	s.Equal(int32(0), person1.ID)
	s.Equal("", person1.Name)
	s.Equal(&newEmail, person1.Email)
	s.WithinDuration(time.Now(), person1.CreatedAt, 2*time.Second)
	s.Nil(person1.UpdatedAt)

	s.Equal(int32(0), person2.ID)
	s.Equal(newName, person2.Name)
	s.Nil(person2.Email)
	s.WithinDuration(time.Now(), person2.CreatedAt, 2*time.Second)
	s.Nil(person2.UpdatedAt)
}

func (s *ReformSuite) TestInsertMultiWithPrimaryKeys() {
	setIdentityInsert(s.T(), s.q, "people", true)

	newEmail := faker.Internet().Email()
	newName := faker.Name().Name()
	person1, person2 := &Person{ID: 50, Email: &newEmail}, &Person{ID: 51, Name: newName}
	err := s.q.InsertMulti(person1, person2)
	s.NoError(err)

	s.Equal(int32(50), person1.ID)
	s.Equal("", person1.Name)
	s.Equal(&newEmail, person1.Email)
	s.WithinDuration(time.Now(), person1.CreatedAt, 2*time.Second)
	s.Nil(person1.UpdatedAt)

	s.Equal(int32(51), person2.ID)
	s.Equal(newName, person2.Name)
	s.Nil(person2.Email)
	s.WithinDuration(time.Now(), person2.CreatedAt, 2*time.Second)
	s.Nil(person2.UpdatedAt)

	person, err := s.q.FindByPrimaryKeyFrom(PersonTable, person1.ID)
	s.NoError(err)
	s.Equal(person1, person)

	person, err = s.q.FindByPrimaryKeyFrom(PersonTable, person2.ID)
	s.NoError(err)
	s.Equal(person2, person)
}

func (s *ReformSuite) TestInsertMultiMixes() {
	err := s.q.InsertMulti()
	s.NoError(err)

	err = s.q.InsertMulti(&Person{}, &Project{})
	s.Error(err)

	err = s.q.InsertMulti(&Person{ID: 1}, &Person{})
	s.Error(err)
}

func (s *ReformSuite) TestInsertIDOnly() {
	id := &IDOnly{}
	err := s.q.Insert(id)
	s.NoError(err)
	s.Equal(int32(1), id.ID)
}

func (s *ReformSuite) TestUpdate() {
	var person Person
	err := s.q.Update(&person)
	s.Equal(reform.ErrNoPK, err)

	person.ID = 99
	err = s.q.Update(&person)
	s.Equal(reform.ErrNoRows, err)

	err = s.q.FindByPrimaryKeyTo(&person, 102)
	s.NoError(err)

	person.Email = pointer.ToString(faker.Internet().Email())
	err = s.q.Update(&person)
	s.NoError(err)
	s.Equal(personCreated, person.CreatedAt)
	s.Require().NotNil(person.UpdatedAt)
	s.WithinDuration(time.Now(), *person.UpdatedAt, 2*time.Second)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.NoError(err)
	s.Equal(&person, person2)
}

func (s *ReformSuite) TestUpdateOverwrite() {
	newEmail := faker.Internet().Email()
	person := Person{ID: 102, Email: pointer.ToString(newEmail)}
	err := s.q.Update(&person)
	s.NoError(err)

	var person2 Person
	err = s.q.FindByPrimaryKeyTo(&person2, person.ID)
	s.NoError(err)
	s.Equal("", person2.Name)
	s.Equal(&newEmail, person2.Email)
	s.WithinDuration(time.Now(), person2.CreatedAt, 2*time.Second)
	s.Require().NotNil(person2.UpdatedAt)
	s.WithinDuration(time.Now(), *person2.UpdatedAt, 2*time.Second)
}

func (s *ReformSuite) TestUpdateColumns() {
	newName := faker.Name().Name()
	newEmail := faker.Internet().Email()

	for p, columns := range map[*Person][]string{
		&Person{Name: "Elfrieda Abbott", Email: &newEmail}:                             {"email", "updated_at"},
		&Person{Name: newName, Email: pointer.ToString("elfrieda_abbott@example.org")}: {"name", "name", "updated_at"},
		&Person{Name: newName, Email: &newEmail}:                                       {"name", "email", "updated_at"},
	} {
		var person Person
		err := s.q.FindByPrimaryKeyTo(&person, 102)
		s.NoError(err)

		person.Name = p.Name
		person.Email = p.Email
		err = s.q.UpdateColumns(&person, columns...)
		s.NoError(err)
		s.Equal(personCreated, person.CreatedAt)
		s.Require().NotNil(person.UpdatedAt)
		s.WithinDuration(time.Now(), *person.UpdatedAt, 2*time.Second)

		person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
		s.NoError(err)
		s.Equal(&person, person2)

		s.RestartTransaction()
	}

	person := &Person{ID: 102, Name: newName, Email: &newEmail, CreatedAt: personCreated}
	for e, columns := range map[error][]string{
		errors.New("reform: unexpected columns: [foo]"):     {"foo"},
		errors.New("reform: will not update PK column: id"): {"id"},
		errors.New("reform: nothing to update"):             {},
	} {
		err := s.q.UpdateColumns(person, columns...)
		s.Error(err)
		s.Equal(e, err)
	}
}

func (s *ReformSuite) TestSave() {
	newName := faker.Name().Name()
	person := &Person{Name: newName}
	err := s.q.Save(person)
	s.NoError(err)

	person2, err := s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.Require().NoError(err)
	s.Equal(newName, person2.(*Person).Name)
	s.Nil(person2.(*Person).Email)
	s.Equal(person, person2)

	newEmail := faker.Internet().Email()
	person.Email = &newEmail
	err = s.q.Save(person)
	s.NoError(err)

	person2, err = s.q.FindByPrimaryKeyFrom(PersonTable, person.ID)
	s.Require().NoError(err)
	s.Equal(newName, person2.(*Person).Name)
	s.Equal(&newEmail, person2.(*Person).Email)
	s.Equal(person, person2)
}

func (s *ReformSuite) TestDelete() {
	person := &Person{ID: 1}
	err := s.q.Delete(person)
	s.NoError(err)
	err = s.q.Reload(person)
	s.Equal(reform.ErrNoRows, err)

	project := &Project{ID: "baron"}
	err = s.q.Delete(project)
	s.NoError(err)
	err = s.q.Reload(project)
	s.Equal(reform.ErrNoRows, err)

	project = &Project{}
	err = s.q.Delete(project)
	s.Equal(reform.ErrNoPK, err)

	project = &Project{ID: "no_such_project"}
	err = s.q.Delete(project)
	s.Equal(reform.ErrNoRows, err)
}

func (s *ReformSuite) TestDeleteFrom() {
	ra, err := s.q.DeleteFrom(PersonTable, "WHERE email IS NULL")
	s.NoError(err)
	s.Equal(uint(3), ra)

	ra, err = s.q.DeleteFrom(PersonTable, "WHERE email IS NULL")
	s.NoError(err)
	s.Equal(uint(0), ra)

	// -1 second for SQLite3, otherwise it also deletes queen itself ¯\_(ツ)_/¯
	ra, err = s.q.DeleteFrom(ProjectTable, "WHERE start < "+s.q.Placeholder(1), queenStart.Add(-time.Second))
	s.NoError(err)
	s.Equal(uint(3), ra)

	ra, err = s.q.DeleteFrom(ProjectTable, "")
	s.NoError(err)
	s.Equal(uint(2), ra)

	ra, err = s.q.DeleteFrom(ProjectTable, "WHERE invalid_tail")
	s.Error(err)
	s.Equal(uint(0), ra)
}

func (s *ReformSuite) TestCommandsSchema() {
	if s.q.Dialect != postgresql.Dialect {
		s.T().Skip("only PostgreSQL supports schemas")
	}

	legacyPerson := &LegacyPerson{Name: pointer.ToString(faker.Name().Name())}
	err := s.q.Save(legacyPerson)
	s.NoError(err)
	err = s.q.Save(legacyPerson)
	s.NoError(err)
	err = s.q.Delete(legacyPerson)
	s.NoError(err)
}
<|endoftext|>


<commit_before>package plugins

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"runtime"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/hashicorp/go-multierror"
	//"github.com/opencontainers/runc/libcontainer/cgroups"
	//cgroupFs "github.com/opencontainers/runc/libcontainer/cgroups/fs"
	//"github.com/opencontainers/runc/libcontainer/cgroups/systemd"
	//cgroupConfig "github.com/opencontainers/runc/libcontainer/configs"

	"github.com/hashicorp/nomad/client/allocdir"
)

var (
	// A mapping of directories on the host OS to attempt to embed inside each
	// task's chroot.
	chrootEnv = map[string]string{
		"/bin":       "/bin",
		"/etc":       "/etc",
		"/lib":       "/lib",
		"/lib32":     "/lib32",
		"/lib64":     "/lib64",
		"/usr/bin":   "/usr/bin",
		"/usr/lib":   "/usr/lib",
		"/usr/share": "/usr/share",
	}
)

type LinuxExecutor struct {
	cmd exec.Cmd
	ctx *ExecutorContext

	//groups  *cgroupConfig.Cgroup
	taskDir string

	logger *log.Logger
	lock   sync.Mutex
}

func NewExecutor(logger *log.Logger) Executor {
	return &LinuxExecutor{logger: logger}
}

func (e *LinuxExecutor) LaunchCmd(command *ExecCommand, ctx *ExecutorContext) (*ProcessState, error) {
	e.ctx = ctx
	e.cmd.Path = command.Cmd
	e.cmd.Args = append([]string{command.Cmd}, command.Args...)
	if filepath.Base(command.Cmd) == command.Cmd {
		if lp, err := exec.LookPath(command.Cmd); err != nil {
		} else {
			e.cmd.Path = lp
		}
	}
	if err := e.configureTaskDir(); err != nil {
		return nil, err
	}
	if err := e.runAs("nobody"); err != nil {
		return nil, err
	}
	e.cmd.Env = ctx.TaskEnv.EnvList()

	stdoPath := filepath.Join(e.taskDir, allocdir.TaskLocal, fmt.Sprintf("%v.stdout", ctx.Task.Name))
	stdo, err := os.OpenFile(stdoPath, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, err
	}
	e.cmd.Stdout = stdo

	stdePath := filepath.Join(e.taskDir, allocdir.TaskLocal, fmt.Sprintf("%v.stderr", ctx.Task.Name))
	stde, err := os.OpenFile(stdePath, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, err
	}
	e.cmd.Stderr = stde

	e.configureChroot()

	if err := e.cmd.Start(); err != nil {
		return nil, fmt.Errorf("error starting command: %v", err)
	}

	return &ProcessState{Pid: e.cmd.Process.Pid, ExitCode: -1, Time: time.Now()}, nil
}

// ConfigureTaskDir creates the necessary directory structure for a proper
// chroot. cleanTaskDir should be called after.
func (e *LinuxExecutor) configureTaskDir() error {
	taskName := e.ctx.Task.Name
	allocDir := e.ctx.AllocDir
	taskDir, ok := allocDir.TaskDirs[taskName]
	if !ok {
		fmt.Errorf("Couldn't find task directory for task %v", taskName)
	}
	e.taskDir = taskDir

	if err := allocDir.MountSharedDir(taskName); err != nil {
		return err
	}

	if err := allocDir.Embed(taskName, chrootEnv); err != nil {
		return err
	}

	// Mount dev
	dev := filepath.Join(taskDir, "dev")
	if !e.pathExists(dev) {
		if err := os.Mkdir(dev, 0777); err != nil {
			return fmt.Errorf("Mkdir(%v) failed: %v", dev, err)
		}

		if err := syscall.Mount("none", dev, "devtmpfs", syscall.MS_RDONLY, ""); err != nil {
			return fmt.Errorf("Couldn't mount /dev to %v: %v", dev, err)
		}
	}

	// Mount proc
	proc := filepath.Join(taskDir, "proc")
	if !e.pathExists(proc) {
		if err := os.Mkdir(proc, 0777); err != nil {
			return fmt.Errorf("Mkdir(%v) failed: %v", proc, err)
		}

		if err := syscall.Mount("none", proc, "proc", syscall.MS_RDONLY, ""); err != nil {
			return fmt.Errorf("Couldn't mount /proc to %v: %v", proc, err)
		}
	}

	// Set the tasks AllocDir environment variable.
	e.ctx.TaskEnv.SetAllocDir(filepath.Join("/", allocdir.SharedAllocName)).SetTaskLocalDir(filepath.Join("/", allocdir.TaskLocal)).Build()
	return nil
}

// runAs takes a user id as a string and looks up the user, and sets the command
// to execute as that user.
func (e *LinuxExecutor) runAs(userid string) error {
	u, err := user.Lookup(userid)
	if err != nil {
		return fmt.Errorf("Failed to identify user %v: %v", userid, err)
	}

	// Convert the uid and gid
	uid, err := strconv.ParseUint(u.Uid, 10, 32)
	if err != nil {
		return fmt.Errorf("Unable to convert userid to uint32: %s", err)
	}
	gid, err := strconv.ParseUint(u.Gid, 10, 32)
	if err != nil {
		return fmt.Errorf("Unable to convert groupid to uint32: %s", err)
	}

	// Set the command to run as that user and group.
	if e.cmd.SysProcAttr == nil {
		e.cmd.SysProcAttr = &syscall.SysProcAttr{}
	}
	if e.cmd.SysProcAttr.Credential == nil {
		e.cmd.SysProcAttr.Credential = &syscall.Credential{}
	}
	e.cmd.SysProcAttr.Credential.Uid = uint32(uid)
	e.cmd.SysProcAttr.Credential.Gid = uint32(gid)

	return nil
}

// pathExists is a helper function to check if the path exists.
func (e *LinuxExecutor) pathExists(path string) bool {
	if _, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}

// configureChroot enters the user command into a chroot if specified in the
// config and on an OS that supports Chroots.
func (e *LinuxExecutor) configureChroot() {
	if !e.ctx.Chroot {
		return
	}
	if e.cmd.SysProcAttr == nil {
		e.cmd.SysProcAttr = &syscall.SysProcAttr{}
	}

	e.cmd.SysProcAttr.Chroot = e.taskDir
	e.cmd.Dir = "/"
}

// cleanTaskDir is an idempotent operation to clean the task directory and
// should be called when tearing down the task.
func (e *LinuxExecutor) cleanTaskDir() error {
	// Prevent a race between Wait/ForceStop
	e.lock.Lock()
	defer e.lock.Unlock()

	// Unmount dev.
	errs := new(multierror.Error)
	dev := filepath.Join(e.taskDir, "dev")
	if e.pathExists(dev) {
		if err := syscall.Unmount(dev, 0); err != nil {
			errs = multierror.Append(errs, fmt.Errorf("Failed to unmount dev (%v): %v", dev, err))
		}

		if err := os.RemoveAll(dev); err != nil {
			errs = multierror.Append(errs, fmt.Errorf("Failed to delete dev directory (%v): %v", dev, err))
		}
	}

	// Unmount
	// proc.
	proc := filepath.Join(e.taskDir, "proc")
	if e.pathExists(proc) {
		if err := syscall.Unmount(proc, 0); err != nil {
			errs = multierror.Append(errs, fmt.Errorf("Failed to unmount proc (%v): %v", proc, err))
		}

		if err := os.RemoveAll(proc); err != nil {
			errs = multierror.Append(errs, fmt.Errorf("Failed to delete proc directory (%v): %v", dev, err))
		}
	}

	return errs.ErrorOrNil()
}

func (e *LinuxExecutor) Wait() (*ProcessState, error) {
	err := e.cmd.Wait()
	if err == nil {
		return &ProcessState{Pid: 0, ExitCode: 0, Time: time.Now()}, nil
	}
	exitCode := 1
	if exitErr, ok := err.(*exec.ExitError); ok {
		if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
			exitCode = status.ExitStatus()
		}
	}
	e.cleanTaskDir()
	return &ProcessState{Pid: 0, ExitCode: exitCode, Time: time.Now()}, nil
}

func (e *LinuxExecutor) Exit() error {
	proc, err := os.FindProcess(e.cmd.Process.Pid)
	if err != nil {
		return fmt.Errorf("failied to find user process %v: %v", e.cmd.Process.Pid, err)
	}
	e.cleanTaskDir()
	return proc.Kill()
}

func (e *LinuxExecutor) ShutDown() error {
	proc, err := os.FindProcess(e.cmd.Process.Pid)
	if err != nil {
		return err
	}
	if runtime.GOOS == "windows" {
		return proc.Kill()
	}
	return proc.Signal(os.Interrupt)
}
<commit_msg>Limiting resources on a process<commit_after>package plugins

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"runtime"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/hashicorp/go-multierror"
	"github.com/opencontainers/runc/libcontainer/cgroups"
	cgroupFs "github.com/opencontainers/runc/libcontainer/cgroups/fs"
	"github.com/opencontainers/runc/libcontainer/cgroups/systemd"
	cgroupConfig "github.com/opencontainers/runc/libcontainer/configs"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/nomad/structs"
)

var (
	// A mapping of directories on the host OS to attempt to embed inside each
	// task's chroot.
	chrootEnv = map[string]string{
		"/bin":       "/bin",
		"/etc":       "/etc",
		"/lib":       "/lib",
		"/lib32":     "/lib32",
		"/lib64":     "/lib64",
		"/usr/bin":   "/usr/bin",
		"/usr/lib":   "/usr/lib",
		"/usr/share": "/usr/share",
	}
)

type LinuxExecutor struct {
	cmd exec.Cmd
	ctx *ExecutorContext

	groups  *cgroupConfig.Cgroup
	taskDir string

	logger *log.Logger
	lock   sync.Mutex
}

func NewExecutor(logger *log.Logger) Executor {
	return &LinuxExecutor{logger: logger}
}

func (e *LinuxExecutor) LaunchCmd(command *ExecCommand, ctx *ExecutorContext) (*ProcessState, error) {
	e.ctx = ctx
	e.cmd.Path = command.Cmd
	e.cmd.Args = append([]string{command.Cmd}, command.Args...)
	if filepath.Base(command.Cmd) == command.Cmd {
		if lp, err := exec.LookPath(command.Cmd); err != nil {
		} else {
			e.cmd.Path = lp
		}
	}
	if err := e.configureTaskDir(); err != nil {
		return nil, err
	}
	if err := e.runAs("nobody"); err != nil {
		return nil, err
	}
	e.cmd.Env = ctx.TaskEnv.EnvList()

	stdoPath := filepath.Join(e.taskDir, allocdir.TaskLocal, fmt.Sprintf("%v.stdout", ctx.Task.Name))
	stdo, err := os.OpenFile(stdoPath, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, err
	}
	e.cmd.Stdout = stdo

	stdePath := filepath.Join(e.taskDir, allocdir.TaskLocal, fmt.Sprintf("%v.stderr", ctx.Task.Name))
	stde, err := os.OpenFile(stdePath, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return nil, err
	}
	e.cmd.Stderr = stde

	e.configureChroot()

	if err := e.configureCgroups(e.ctx.Task.Resources); err != nil {
		return nil, fmt.Errorf("error creating cgroups: %v", err)
	}

	if err := e.cmd.Start(); err != nil {
		return nil, fmt.Errorf("error starting command: %v", err)
	}

	manager := e.getCgroupManager(e.groups)
	if err := manager.Apply(e.cmd.Process.Pid); err != nil {
		e.logger.Printf("[ERROR] unable to join cgroup: %v", err)
		if err := e.Exit(); err != nil {
			e.logger.Printf("[ERROR] unable to kill process: %v", err)
		}
		return nil, err
	}

	return &ProcessState{Pid: e.cmd.Process.Pid, ExitCode: -1, Time: time.Now()}, nil
}

// ConfigureTaskDir creates the necessary directory structure for a proper
// chroot. cleanTaskDir should be called after.
func (e *LinuxExecutor) configureTaskDir() error {
	taskName := e.ctx.Task.Name
	allocDir := e.ctx.AllocDir
	taskDir, ok := allocDir.TaskDirs[taskName]
	if !ok {
		fmt.Errorf("Couldn't find task directory for task %v", taskName)
	}
	e.taskDir = taskDir

	if err := allocDir.MountSharedDir(taskName); err != nil {
		return err
	}

	if err := allocDir.Embed(taskName, chrootEnv); err != nil {
		return err
	}

	// Mount dev
	dev := filepath.Join(taskDir, "dev")
	if !e.pathExists(dev) {
		if err := os.Mkdir(dev, 0777); err != nil {
			return fmt.Errorf("Mkdir(%v) failed: %v", dev, err)
		}

		if err := syscall.Mount("none", dev, "devtmpfs", syscall.MS_RDONLY, ""); err != nil {
			return fmt.Errorf("Couldn't mount /dev to %v: %v", dev, err)
		}
	}

	// Mount proc
	proc := filepath.Join(taskDir, "proc")
	if !e.pathExists(proc) {
		if err := os.Mkdir(proc, 0777); err != nil {
			return fmt.Errorf("Mkdir(%v) failed: %v", proc, err)
		}

		if err := syscall.Mount("none", proc, "proc", syscall.MS_RDONLY, ""); err != nil {
			return fmt.Errorf("Couldn't mount /proc to %v: %v", proc, err)
		}
	}

	// Set the tasks AllocDir environment variable.
	e.ctx.TaskEnv.SetAllocDir(filepath.Join("/", allocdir.SharedAllocName)).SetTaskLocalDir(filepath.Join("/", allocdir.TaskLocal)).Build()
	return nil
}

// runAs takes a user id as a string and looks up the user, and sets the command
// to execute as that user.
func (e *LinuxExecutor) runAs(userid string) error {
	u, err := user.Lookup(userid)
	if err != nil {
		return fmt.Errorf("Failed to identify user %v: %v", userid, err)
	}

	// Convert the uid and gid
	uid, err := strconv.ParseUint(u.Uid, 10, 32)
	if err != nil {
		return fmt.Errorf("Unable to convert userid to uint32: %s", err)
	}
	gid, err := strconv.ParseUint(u.Gid, 10, 32)
	if err != nil {
		return fmt.Errorf("Unable to convert groupid to uint32: %s", err)
	}

	// Set the command to run as that user and group.
	if e.cmd.SysProcAttr == nil {
		e.cmd.SysProcAttr = &syscall.SysProcAttr{}
	}
	if e.cmd.SysProcAttr.Credential == nil {
		e.cmd.SysProcAttr.Credential = &syscall.Credential{}
	}
	e.cmd.SysProcAttr.Credential.Uid = uint32(uid)
	e.cmd.SysProcAttr.Credential.Gid = uint32(gid)

	return nil
}

// pathExists is a helper function to check if the path exists.
func (e *LinuxExecutor) pathExists(path string) bool {
	if _, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}

// configureChroot enters the user command into a chroot if specified in the
// config and on an OS that supports Chroots.
func (e *LinuxExecutor) configureChroot() {
	if !e.ctx.Chroot {
		return
	}
	if e.cmd.SysProcAttr == nil {
		e.cmd.SysProcAttr = &syscall.SysProcAttr{}
	}

	e.cmd.SysProcAttr.Chroot = e.taskDir
	e.cmd.Dir = "/"
}

// cleanTaskDir is an idempotent operation to clean the task directory and
// should be called when tearing down the task.
func (e *LinuxExecutor) cleanTaskDir() error {
	// Prevent a race between Wait/ForceStop
	e.lock.Lock()
	defer e.lock.Unlock()

	// Unmount dev.
	errs := new(multierror.Error)
	dev := filepath.Join(e.taskDir, "dev")
	if e.pathExists(dev) {
		if err := syscall.Unmount(dev, 0); err != nil {
			errs = multierror.Append(errs, fmt.Errorf("Failed to unmount dev (%v): %v", dev, err))
		}

		if err := os.RemoveAll(dev); err != nil {
			errs = multierror.Append(errs, fmt.Errorf("Failed to delete dev directory (%v): %v", dev, err))
		}
	}

	// Unmount
	// proc.
	proc := filepath.Join(e.taskDir, "proc")
	if e.pathExists(proc) {
		if err := syscall.Unmount(proc, 0); err != nil {
			errs = multierror.Append(errs, fmt.Errorf("Failed to unmount proc (%v): %v", proc, err))
		}

		if err := os.RemoveAll(proc); err != nil {
			errs = multierror.Append(errs, fmt.Errorf("Failed to delete proc directory (%v): %v", dev, err))
		}
	}

	return errs.ErrorOrNil()
}

func (e *LinuxExecutor) Wait() (*ProcessState, error) {
	err := e.cmd.Wait()
	if err == nil {
		return &ProcessState{Pid: 0, ExitCode: 0, Time: time.Now()}, nil
	}
	exitCode := 1
	if exitErr, ok := err.(*exec.ExitError); ok {
		if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
			exitCode = status.ExitStatus()
		}
	}
	e.cleanTaskDir()
	e.destroyCgroup()
	return &ProcessState{Pid: 0, ExitCode: exitCode, Time: time.Now()}, nil
}

func (e *LinuxExecutor) Exit() error {
	e.logger.Printf("[INFO] Exiting plugin for task %q", e.ctx.Task.Name)
	proc, err := os.FindProcess(e.cmd.Process.Pid)
	if err != nil {
		return fmt.Errorf("failied to find user process %v: %v", e.cmd.Process.Pid, err)
	}
	e.cleanTaskDir()
	e.destroyCgroup()
	return proc.Kill()
}

func (e *LinuxExecutor) ShutDown() error {
	proc, err := os.FindProcess(e.cmd.Process.Pid)
	if err != nil {
		return err
	}
	if runtime.GOOS == "windows" {
		return proc.Kill()
	}
	return proc.Signal(os.Interrupt)
}

// configureCgroups converts a Nomad Resources specification into the equivalent
// cgroup configuration. It returns an error if the resources are invalid.
func (e *LinuxExecutor) configureCgroups(resources *structs.Resources) error {
	e.groups = &cgroupConfig.Cgroup{}
	e.groups.Resources = &cgroupConfig.Resources{}
	e.groups.Name = structs.GenerateUUID()

	// TODO: verify this is needed for things like network access
	e.groups.Resources.AllowAllDevices = true

	if resources.MemoryMB > 0 {
		// Total amount of memory allowed to consume
		e.groups.Resources.Memory = int64(resources.MemoryMB * 1024 * 1024)
		// Disable swap to avoid issues on the machine
		e.groups.Resources.MemorySwap = int64(-1)
	}

	if resources.CPU < 2 {
		return fmt.Errorf("resources.CPU must be equal to or greater than 2: %v", resources.CPU)
	}

	// Set the relative CPU shares for this cgroup.
	e.groups.Resources.CpuShares = int64(resources.CPU)

	if resources.IOPS != 0 {
		// Validate it is in an acceptable range.
		if resources.IOPS < 10 || resources.IOPS > 1000 {
			return fmt.Errorf("resources.IOPS must be between 10 and 1000: %d", resources.IOPS)
		}

		e.groups.Resources.BlkioWeight = uint16(resources.IOPS)
	}

	return nil
}

// destroyCgroup kills all processes in the cgroup and removes the cgroup
// configuration from the host.
func (e *LinuxExecutor) destroyCgroup() error {
	if e.groups == nil {
		return fmt.Errorf("Can't destroy: cgroup configuration empty")
	}

	// Prevent a race between Wait/ForceStop
	e.lock.Lock()
	defer e.lock.Unlock()

	manager := e.getCgroupManager(e.groups)
	pids, err := manager.GetPids()
	if err != nil {
		return fmt.Errorf("Failed to get pids in the cgroup %v: %v", e.groups.Name, err)
	}

	errs := new(multierror.Error)
	for _, pid := range pids {
		process, err := os.FindProcess(pid)
		if err != nil {
			multierror.Append(errs, fmt.Errorf("Failed to find Pid %v: %v", pid, err))
			continue
		}

		if err := process.Kill(); err != nil && err.Error() != "os: process already finished" {
			multierror.Append(errs, fmt.Errorf("Failed to kill Pid %v: %v", pid, err))
			continue
		}
	}

	// Remove the cgroup.
	if err := manager.Destroy(); err != nil {
		multierror.Append(errs, fmt.Errorf("Failed to delete the cgroup directories: %v", err))
	}

	if len(errs.Errors) != 0 {
		return fmt.Errorf("Failed to destroy cgroup: %v", errs)
	}

	return nil
}

// getCgroupManager returns the correct libcontainer cgroup manager.
func (e *LinuxExecutor) getCgroupManager(groups *cgroupConfig.Cgroup) cgroups.Manager {
	var manager cgroups.Manager
	manager = &cgroupFs.Manager{Cgroups: groups}
	if systemd.UseSystemd() {
		manager = &systemd.Manager{Cgroups: groups}
	}
	return manager
}
<|endoftext|>


<commit_before>package bourbon

import (
	"github.com/bmizerany/assert"
	"github.com/codegangsta/inject"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestMiddlewareContentTypeHandlerBlank(t *testing.T) {
	rw := httptest.NewRecorder()
	r, _ := http.NewRequest("GET", "http://example.com", nil)

	code, encodeable := ContentTypeHandler(rw, r)
	contentType := "application/json; charset=utf-8"

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
	assert.Equal(t, contentType, rw.Header().Get("Content-Type"))
}

func TestMiddlewareContentTypeHandlerWithJSON(t *testing.T) {
	contentType := "application/json; charset=utf-8"
	rw := httptest.NewRecorder()
	r, _ := http.NewRequest("GET", "http://example.com", nil)
	r.Header.Set("Content-Type", contentType)

	code, encodeable := ContentTypeHandler(rw, r)

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
	assert.Equal(t, contentType, rw.Header().Get("Content-Type"))
}

func TestMiddlewareContentTypeHandlerWithVndJSON(t *testing.T) {
	rw := httptest.NewRecorder()
	r, _ := http.NewRequest("GET", "http://example.com", nil)
	r.Header.Set("Content-Type", "application/vnd.example+json; param=2")

	code, encodeable := ContentTypeHandler(rw, r)
	contentType := "application/json; charset=utf-8"

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
	assert.Equal(t, contentType, rw.Header().Get("Content-Type"))
}

func TestMiddlewareContentTypeHandlerWithHTML(t *testing.T) {
	rw := httptest.NewRecorder()
	r, _ := http.NewRequest("GET", "http://example.com", nil)
	r.Header.Set("Content-Type", "text/html")

	code, encodeable := ContentTypeHandler(rw, r)
	contentType := "application/json; charset=utf-8"
	message := Message{
		Code:    415,
		Message: http.StatusText(415),
		Errors:  []string{`"text/html" is not a supported Content-Type`},
	}

	assert.Equal(t, 415, code)
	assert.Equal(t, message, encodeable)
	assert.Equal(t, contentType, rw.Header().Get("Content-Type"))
}

func TestMiddlewareDecodeHandlerWithJSON(t *testing.T) {
	type example struct {
		Name string `json:"name"`
	}
	reader := strings.NewReader(`{"name":"example"}`)
	fn := func(e example) {}

	r, _ := http.NewRequest("GET", "http://example.com", reader)
	c := createTestContext(fn, r)

	code, encodeable := DecodeHandler(c, r)

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
}

func TestMiddlewareDecodeHandlerWithMalformedJSON(t *testing.T) {
	type example struct {
		Name string `json:"name"`
	}
	reader := strings.NewReader(`{"name":""boom}`) // incorrect quoting
	fn := func(e example) {}

	r, _ := http.NewRequest("GET", "http://example.com", reader)
	c := createTestContext(fn, r)

	message := Message{
		Code:    400,
		Message: http.StatusText(400),
		Errors:  []string{"invalid character 'b' after object key:value pair"},
	}

	code, encodeable := DecodeHandler(c, r)

	assert.Equal(t, 400, code)
	assert.Equal(t, message, encodeable)
}

func TestMiddlewareDecodeHandlerWithBlankBody(t *testing.T) {
	type example struct {
		Name string `json:"name"`
	}
	reader := strings.NewReader("") // empty string
	fn := func(e example) {}

	r, _ := http.NewRequest("GET", "http://example.com", reader)
	c := createTestContext(fn, r)

	code, encodeable := DecodeHandler(c, r)

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
}

func TestMiddlewareDecodeHandlerWithContentLengthZero(t *testing.T) {
	fn := func() {}
	r, _ := http.NewRequest("GET", "http://example.com", nil)
	r.Header.Add("Content-Length", "0")
	c := createTestContext(fn, r)

	code, encodeable := DecodeHandler(c, r)

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
}

func createTestContext(fn Handler, r *http.Request) context {
	rt := &route{new(bourbon), "GET", "/", fn, nil}
	return context{inject.New(), fn, nil, rt, nil, r}
}
<commit_msg>fix middleware test<commit_after>package bourbon

import (
	"github.com/bmizerany/assert"
	"github.com/codegangsta/inject"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestMiddlewareContentTypeHandlerBlank(t *testing.T) {
	rw := httptest.NewRecorder()
	r, _ := http.NewRequest("GET", "http://example.com", nil)

	code, encodeable := ContentTypeHandler(rw, r)
	contentType := "application/json; charset=utf-8"

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
	assert.Equal(t, contentType, rw.Header().Get("Content-Type"))
}

func TestMiddlewareContentTypeHandlerWithJSON(t *testing.T) {
	contentType := "application/json; charset=utf-8"
	rw := httptest.NewRecorder()
	r, _ := http.NewRequest("GET", "http://example.com", nil)
	r.Header.Set("Content-Type", contentType)

	code, encodeable := ContentTypeHandler(rw, r)

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
	assert.Equal(t, contentType, rw.Header().Get("Content-Type"))
}

func TestMiddlewareContentTypeHandlerWithVndJSON(t *testing.T) {
	rw := httptest.NewRecorder()
	r, _ := http.NewRequest("GET", "http://example.com", nil)
	r.Header.Set("Content-Type", "application/vnd.example+json; param=2")

	code, encodeable := ContentTypeHandler(rw, r)
	contentType := "application/json; charset=utf-8"

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
	assert.Equal(t, contentType, rw.Header().Get("Content-Type"))
}

func TestMiddlewareContentTypeHandlerWithHTML(t *testing.T) {
	rw := httptest.NewRecorder()
	r, _ := http.NewRequest("GET", "http://example.com", nil)
	r.Header.Set("Content-Type", "text/html")

	code, encodeable := ContentTypeHandler(rw, r)
	contentType := "application/json; charset=utf-8"
	message := Message{
		Code:    415,
		Message: http.StatusText(415),
		Errors:  []string{`"text/html" is not a supported Content-Type`},
	}

	assert.Equal(t, 415, code)
	assert.Equal(t, message, encodeable)
	assert.Equal(t, contentType, rw.Header().Get("Content-Type"))
}

func TestMiddlewareDecodeHandlerWithJSON(t *testing.T) {
	type example struct {
		Name string `json:"name"`
	}
	reader := strings.NewReader(`{"name":"example"}`)
	fn := func(e example) {}

	r, _ := http.NewRequest("GET", "http://example.com", reader)
	c := createTestContext(fn, r)

	code, encodeable := DecodeHandler(c, r)

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
}

func TestMiddlewareDecodeHandlerWithMalformedJSON(t *testing.T) {
	type example struct {
		Name string `json:"name"`
	}
	reader := strings.NewReader(`{"name":""boom}`) // incorrect quoting
	fn := func(e example) {}

	r, _ := http.NewRequest("GET", "http://example.com", reader)
	c := createTestContext(fn, r)

	message := Message{
		Code:    400,
		Message: http.StatusText(400),
		Errors:  []string{"invalid character 'b' after object key:value pair"},
	}

	code, encodeable := DecodeHandler(c, r)

	assert.Equal(t, 400, code)
	assert.Equal(t, message, encodeable)
}

func TestMiddlewareDecodeHandlerWithBlankBody(t *testing.T) {
	type example struct {
		Name string `json:"name"`
	}
	reader := strings.NewReader("") // empty string
	fn := func(e example) {}

	r, _ := http.NewRequest("GET", "http://example.com", reader)
	c := createTestContext(fn, r)

	code, encodeable := DecodeHandler(c, r)

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
}

func TestMiddlewareDecodeHandlerWithContentLengthZero(t *testing.T) {
	fn := func() {}
	r, _ := http.NewRequest("GET", "http://example.com", nil)
	r.Header.Add("Content-Length", "0")
	c := createTestContext(fn, r)

	code, encodeable := DecodeHandler(c, r)

	assert.Equal(t, 0, code)
	assert.Equal(t, nil, encodeable)
}

func createTestContext(fn Handler, r *http.Request) *context {
	rt := &route{new(bourbon), "GET", "/", fn, nil}
	return &context{inject.New(), fn, nil, rt, nil, r}
}
<|endoftext|>


<commit_before>/*
	Script to handle an email response to the reminder.
	Steps:
		1. Get Sender - is it one of the registered senders in queue and are they the hosting group?
		2. Check if they should be responding or if someone is being snarky.
		3. Look for Yes/No/Skip
			3a. Yes - update the order in the group
			3b. No - send an email to the next in line, update group
			3c. Skip - respond with the current turn order for next week
	Notes: Super procedural right now.  I need to clean up the code

*/
package hostqueue

import (
	"appengine"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/mail"
	"regexp"
	"strings"
)

func IncomingMail(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)
	defer r.Body.Close()

	m, err := mail.ReadMessage(r.Body)
	if err != nil {
		log.Fatal(err)
	}

	from := m.Header.Get("From")
	//clean up responses formatted with a name first.  Ex: 'Ann Addicks <test@test.com>'
	from = strings.Split(from, "<")[1]
	from = strings.Split(from, ">")[0]
	ctx.Infof("from : " + from)
	g, err := findGroup(ctx, from)
	if err != nil {
		log.Fatal(err)
	} else if isValidResponder(from, g) {
		responseRegex := regexp.MustCompile(`(yes\b|no\b|skip\b)(.*?)`)
		body, err := ioutil.ReadAll(m.Body)

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

		s := string(body)
		bodyString := strings.ToLower(s)
		bodyString = strings.Split(bodyString, "it is your turn to host!")[0]

		switch responseRegex.FindString(bodyString) {
		case "yes":
			// Update the order in the group by deleting the current host and appending them to the end.
			hosts := g.Hosts
			currentHost := hosts[g.Next]
			currentHost.TimesHosted++

			hosts = append(hosts[:g.Next], hosts[g.Next+1:]...)
			hosts = append(hosts, currentHost)
			g.Hosts = hosts
			g.Next = 0

			g.save(ctx)
			ctx.Infof("Match Yes")
		case "no":
			//Send an email to the next in line && update the group
			hosts := g.Hosts

			if g.Next < (len(hosts) - 1) {
				g.Next++
			} else {
				g.Next = 0
			}

			g.save(ctx)
			sendReminder(g, r)
			ctx.Infof("Match No")
		case "skip":
			//Respond with the current turn order for next week
			sendSkipMessage(g, r)
			ctx.Infof("Match Skip")
		default:
			ctx.Infof("Could not find yes/no/skip")
		}
	} else {
		ctx.Infof("Sender is not valid")
	}
}

/*Pull all the groups and loop through to find the email address.
  I need to learn more about querying app engines datastore to make this nicer.*/
func findGroup(ctx appengine.Context, from string) (Group, error) {
	groups, err := GetGroups(ctx)
	var g Group
	if err != nil {
		return g, err
	}

	for _, group := range groups {
		for _, host := range group.Hosts {
			if strings.Contains(host.Emails, from) {
				return group, nil
			}
		}
	}
	return g, fmt.Errorf("Cannot find a group for: %s", from)
}

//Check if they should be responding or if someone is being snarky.
func isValidResponder(from string, g Group) bool {
	validEmails := g.Hosts[g.Next].Emails
	return strings.Contains(validEmails, from)
}

//http://stackoverflow.com/questions/10485743/contains-method-for-a-slice
func contains(slice []string, item string) bool {
	set := make(map[string]struct{}, len(slice))
	for _, s := range slice {
		set[s] = struct{}{}
	}

	_, ok := set[item]
	return ok
}
<commit_msg>call new host confirmed message to ack back.<commit_after>/*
	Script to handle an email response to the reminder.
	Steps:
		1. Get Sender - is it one of the registered senders in queue and are they the hosting group?
		2. Check if they should be responding or if someone is being snarky.
		3. Look for Yes/No/Skip
			3a. Yes - update the order in the group
			3b. No - send an email to the next in line, update group
			3c. Skip - respond with the current turn order for next week
	Notes: Super procedural right now.  I need to clean up the code

*/
package hostqueue

import (
	"appengine"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/mail"
	"regexp"
	"strings"
)

func IncomingMail(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)
	defer r.Body.Close()

	m, err := mail.ReadMessage(r.Body)
	if err != nil {
		log.Fatal(err)
	}

	from := m.Header.Get("From")
	//clean up responses formatted with a name first.  Ex: 'Ann Addicks <test@test.com>'
	from = strings.Split(from, "<")[1]
	from = strings.Split(from, ">")[0]
	ctx.Infof("from : " + from)
	g, err := findGroup(ctx, from)
	if err != nil {
		log.Fatal(err)
	} else if isValidResponder(from, g) {
		responseRegex := regexp.MustCompile(`(yes\b|no\b|skip\b)(.*?)`)
		body, err := ioutil.ReadAll(m.Body)

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

		s := string(body)
		bodyString := strings.ToLower(s)
		bodyString = strings.Split(bodyString, "it is your turn to host!")[0]

		switch responseRegex.FindString(bodyString) {
		case "yes":
			// Update the order in the group by deleting the current host and appending them to the end.
			hosts := g.Hosts
			currentHost := hosts[g.Next]
			currentHost.TimesHosted++

			hosts = append(hosts[:g.Next], hosts[g.Next+1:]...)
			hosts = append(hosts, currentHost)
			g.Hosts = hosts
			g.Next = 0

			g.save(ctx)
			sendHostConfirmedMessage(g, r)
			ctx.Infof("Match Yes")
		case "no":
			//Send an email to the next in line && update the group
			hosts := g.Hosts

			if g.Next < (len(hosts) - 1) {
				g.Next++
			} else {
				g.Next = 0
			}

			g.save(ctx)
			sendReminder(g, r)
			ctx.Infof("Match No")
		case "skip":
			//Respond with the current turn order for next week
			sendSkipMessage(g, r)
			ctx.Infof("Match Skip")
		default:
			ctx.Infof("Could not find yes/no/skip")
		}
	} else {
		ctx.Infof("Sender is not valid")
	}
}

/*Pull all the groups and loop through to find the email address.
  I need to learn more about querying app engines datastore to make this nicer.*/
func findGroup(ctx appengine.Context, from string) (Group, error) {
	groups, err := GetGroups(ctx)
	var g Group
	if err != nil {
		return g, err
	}

	for _, group := range groups {
		for _, host := range group.Hosts {
			if strings.Contains(host.Emails, from) {
				return group, nil
			}
		}
	}
	return g, fmt.Errorf("Cannot find a group for: %s", from)
}

//Check if they should be responding or if someone is being snarky.
func isValidResponder(from string, g Group) bool {
	validEmails := g.Hosts[g.Next].Emails
	return strings.Contains(validEmails, from)
}

//http://stackoverflow.com/questions/10485743/contains-method-for-a-slice
func contains(slice []string, item string) bool {
	set := make(map[string]struct{}, len(slice))
	for _, s := range slice {
		set[s] = struct{}{}
	}

	_, ok := set[item]
	return ok
}
<|endoftext|>