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

// +build arm64

package arch

import (
	"fmt"
	"io"
	"syscall"

	"gvisor.dev/gvisor/pkg/binary"
	"gvisor.dev/gvisor/pkg/cpuid"
	"gvisor.dev/gvisor/pkg/log"
	rpb "gvisor.dev/gvisor/pkg/sentry/arch/registers_go_proto"
	"gvisor.dev/gvisor/pkg/syserror"
	"gvisor.dev/gvisor/pkg/usermem"
)

const (
	// SyscallWidth is the width of insturctions.
	SyscallWidth = 4
)

// aarch64FPState is aarch64 floating point state.
type aarch64FPState []byte

// initAarch64FPState (defined in asm files) sets up initial state.
func initAarch64FPState(data *FloatingPointData) {
	// TODO(gvisor.dev/issue/1238): floating-point is not supported.
}

func newAarch64FPStateSlice() []byte {
	return alignedBytes(4096, 32)[:4096]
}

// newAarch64FPState returns an initialized floating point state.
//
// The returned state is large enough to store all floating point state
// supported by host, even if the app won't use much of it due to a restricted
// FeatureSet. Since they may still be able to see state not advertised by
// CPUID we must ensure it does not contain any sentry state.
func newAarch64FPState() aarch64FPState {
	f := aarch64FPState(newAarch64FPStateSlice())
	initAarch64FPState(f.FloatingPointData())
	return f
}

// fork creates and returns an identical copy of the aarch64 floating point state.
func (f aarch64FPState) fork() aarch64FPState {
	n := aarch64FPState(newAarch64FPStateSlice())
	copy(n, f)
	return n
}

// FloatingPointData returns the raw data pointer.
func (f aarch64FPState) FloatingPointData() *FloatingPointData {
	return (*FloatingPointData)(&f[0])
}

// NewFloatingPointData returns a new floating point data blob.
//
// This is primarily for use in tests.
func NewFloatingPointData() *FloatingPointData {
	return (*FloatingPointData)(&(newAarch64FPState()[0]))
}

// State contains the common architecture bits for aarch64 (the build tag of this
// file ensures it's only built on aarch64).
type State struct {
	// The system registers.
	Regs syscall.PtraceRegs `state:".(syscallPtraceRegs)"`

	// Our floating point state.
	aarch64FPState `state:"wait"`

	// FeatureSet is a pointer to the currently active feature set.
	FeatureSet *cpuid.FeatureSet
}

// Proto returns a protobuf representation of the system registers in State.
func (s State) Proto() *rpb.Registers {
	regs := &rpb.ARM64Registers{
		R0:     s.Regs.Regs[0],
		R1:     s.Regs.Regs[1],
		R2:     s.Regs.Regs[2],
		R3:     s.Regs.Regs[3],
		R4:     s.Regs.Regs[4],
		R5:     s.Regs.Regs[5],
		R6:     s.Regs.Regs[6],
		R7:     s.Regs.Regs[7],
		R8:     s.Regs.Regs[8],
		R9:     s.Regs.Regs[9],
		R10:    s.Regs.Regs[10],
		R11:    s.Regs.Regs[11],
		R12:    s.Regs.Regs[12],
		R13:    s.Regs.Regs[13],
		R14:    s.Regs.Regs[14],
		R15:    s.Regs.Regs[15],
		R16:    s.Regs.Regs[16],
		R17:    s.Regs.Regs[17],
		R18:    s.Regs.Regs[18],
		R19:    s.Regs.Regs[19],
		R20:    s.Regs.Regs[20],
		R21:    s.Regs.Regs[21],
		R22:    s.Regs.Regs[22],
		R23:    s.Regs.Regs[23],
		R24:    s.Regs.Regs[24],
		R25:    s.Regs.Regs[25],
		R26:    s.Regs.Regs[26],
		R27:    s.Regs.Regs[27],
		R28:    s.Regs.Regs[28],
		R29:    s.Regs.Regs[29],
		R30:    s.Regs.Regs[30],
		Sp:     s.Regs.Sp,
		Pc:     s.Regs.Pc,
		Pstate: s.Regs.Pstate,
	}
	return &rpb.Registers{Arch: &rpb.Registers_Arm64{Arm64: regs}}
}

// Fork creates and returns an identical copy of the state.
func (s *State) Fork() State {
	// TODO(gvisor.dev/issue/1238): floating-point is not supported.
	return State{
		Regs:       s.Regs,
		FeatureSet: s.FeatureSet,
	}
}

// StateData implements Context.StateData.
func (s *State) StateData() *State {
	return s
}

// CPUIDEmulate emulates a cpuid instruction.
func (s *State) CPUIDEmulate(l log.Logger) {
	// TODO(gvisor.dev/issue/1255): cpuid is not supported.
}

// SingleStep implements Context.SingleStep.
func (s *State) SingleStep() bool {
	return false
}

// SetSingleStep enables single stepping.
func (s *State) SetSingleStep() {
	// Set the trap flag.
	// TODO(gvisor.dev/issue/1239): ptrace single-step is not supported.
}

// ClearSingleStep enables single stepping.
func (s *State) ClearSingleStep() {
	// Clear the trap flag.
	// TODO(gvisor.dev/issue/1239): ptrace single-step is not supported.
}

// RegisterMap returns a map of all registers.
func (s *State) RegisterMap() (map[string]uintptr, error) {
	return map[string]uintptr{
		"R0":     uintptr(s.Regs.Regs[0]),
		"R1":     uintptr(s.Regs.Regs[1]),
		"R2":     uintptr(s.Regs.Regs[2]),
		"R3":     uintptr(s.Regs.Regs[3]),
		"R4":     uintptr(s.Regs.Regs[4]),
		"R5":     uintptr(s.Regs.Regs[5]),
		"R6":     uintptr(s.Regs.Regs[6]),
		"R7":     uintptr(s.Regs.Regs[7]),
		"R8":     uintptr(s.Regs.Regs[8]),
		"R9":     uintptr(s.Regs.Regs[9]),
		"R10":    uintptr(s.Regs.Regs[10]),
		"R11":    uintptr(s.Regs.Regs[11]),
		"R12":    uintptr(s.Regs.Regs[12]),
		"R13":    uintptr(s.Regs.Regs[13]),
		"R14":    uintptr(s.Regs.Regs[14]),
		"R15":    uintptr(s.Regs.Regs[15]),
		"R16":    uintptr(s.Regs.Regs[16]),
		"R17":    uintptr(s.Regs.Regs[17]),
		"R18":    uintptr(s.Regs.Regs[18]),
		"R19":    uintptr(s.Regs.Regs[19]),
		"R20":    uintptr(s.Regs.Regs[20]),
		"R21":    uintptr(s.Regs.Regs[21]),
		"R22":    uintptr(s.Regs.Regs[22]),
		"R23":    uintptr(s.Regs.Regs[23]),
		"R24":    uintptr(s.Regs.Regs[24]),
		"R25":    uintptr(s.Regs.Regs[25]),
		"R26":    uintptr(s.Regs.Regs[26]),
		"R27":    uintptr(s.Regs.Regs[27]),
		"R28":    uintptr(s.Regs.Regs[28]),
		"R29":    uintptr(s.Regs.Regs[29]),
		"R30":    uintptr(s.Regs.Regs[30]),
		"Sp":     uintptr(s.Regs.Sp),
		"Pc":     uintptr(s.Regs.Pc),
		"Pstate": uintptr(s.Regs.Pstate),
	}, nil
}

// PtraceGetRegs implements Context.PtraceGetRegs.
func (s *State) PtraceGetRegs(dst io.Writer) (int, error) {
	return dst.Write(binary.Marshal(nil, usermem.ByteOrder, s.ptraceGetRegs()))
}

func (s *State) ptraceGetRegs() syscall.PtraceRegs {
	return s.Regs
}

var ptraceRegsSize = int(binary.Size(syscall.PtraceRegs{}))

// PtraceSetRegs implements Context.PtraceSetRegs.
func (s *State) PtraceSetRegs(src io.Reader) (int, error) {
	var regs syscall.PtraceRegs
	buf := make([]byte, ptraceRegsSize)
	if _, err := io.ReadFull(src, buf); err != nil {
		return 0, err
	}
	binary.Unmarshal(buf, usermem.ByteOrder, &regs)
	s.Regs = regs
	return ptraceRegsSize, nil
}

// PtraceGetFPRegs implements Context.PtraceGetFPRegs.
func (s *State) PtraceGetFPRegs(dst io.Writer) (int, error) {
	// TODO(gvisor.dev/issue/1238): floating-point is not supported.
	return 0, nil
}

// PtraceSetFPRegs implements Context.PtraceSetFPRegs.
func (s *State) PtraceSetFPRegs(src io.Reader) (int, error) {
	// TODO(gvisor.dev/issue/1238): floating-point is not supported.
	return 0, nil
}

// Register sets defined in include/uapi/linux/elf.h.
const (
	_NT_PRSTATUS = 1
	_NT_PRFPREG  = 2
)

// PtraceGetRegSet implements Context.PtraceGetRegSet.
func (s *State) PtraceGetRegSet(regset uintptr, dst io.Writer, maxlen int) (int, error) {
	switch regset {
	case _NT_PRSTATUS:
		if maxlen < ptraceRegsSize {
			return 0, syserror.EFAULT
		}
		return s.PtraceGetRegs(dst)
	default:
		return 0, syserror.EINVAL
	}
}

// PtraceSetRegSet implements Context.PtraceSetRegSet.
func (s *State) PtraceSetRegSet(regset uintptr, src io.Reader, maxlen int) (int, error) {
	switch regset {
	case _NT_PRSTATUS:
		if maxlen < ptraceRegsSize {
			return 0, syserror.EFAULT
		}
		return s.PtraceSetRegs(src)
	default:
		return 0, syserror.EINVAL
	}
}

// FullRestore indicates whether a full restore is required.
func (s *State) FullRestore() bool {
	return false
}

// New returns a new architecture context.
func New(arch Arch, fs *cpuid.FeatureSet) Context {
	switch arch {
	case ARM64:
		return &context64{
			State{
				FeatureSet: fs,
			},
		}
	}
	panic(fmt.Sprintf("unknown architecture %v", arch))
}
<commit_msg>Add definition of arch.ARMTrapFlag.<commit_after>// Copyright 2020 The gVisor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build arm64

package arch

import (
	"fmt"
	"io"
	"syscall"

	"gvisor.dev/gvisor/pkg/binary"
	"gvisor.dev/gvisor/pkg/cpuid"
	"gvisor.dev/gvisor/pkg/log"
	rpb "gvisor.dev/gvisor/pkg/sentry/arch/registers_go_proto"
	"gvisor.dev/gvisor/pkg/syserror"
	"gvisor.dev/gvisor/pkg/usermem"
)

const (
	// SyscallWidth is the width of insturctions.
	SyscallWidth = 4
)

// ARMTrapFlag is the mask for the trap flag.
const ARMTrapFlag = uint64(1) << 21

// aarch64FPState is aarch64 floating point state.
type aarch64FPState []byte

// initAarch64FPState (defined in asm files) sets up initial state.
func initAarch64FPState(data *FloatingPointData) {
	// TODO(gvisor.dev/issue/1238): floating-point is not supported.
}

func newAarch64FPStateSlice() []byte {
	return alignedBytes(4096, 32)[:4096]
}

// newAarch64FPState returns an initialized floating point state.
//
// The returned state is large enough to store all floating point state
// supported by host, even if the app won't use much of it due to a restricted
// FeatureSet. Since they may still be able to see state not advertised by
// CPUID we must ensure it does not contain any sentry state.
func newAarch64FPState() aarch64FPState {
	f := aarch64FPState(newAarch64FPStateSlice())
	initAarch64FPState(f.FloatingPointData())
	return f
}

// fork creates and returns an identical copy of the aarch64 floating point state.
func (f aarch64FPState) fork() aarch64FPState {
	n := aarch64FPState(newAarch64FPStateSlice())
	copy(n, f)
	return n
}

// FloatingPointData returns the raw data pointer.
func (f aarch64FPState) FloatingPointData() *FloatingPointData {
	return (*FloatingPointData)(&f[0])
}

// NewFloatingPointData returns a new floating point data blob.
//
// This is primarily for use in tests.
func NewFloatingPointData() *FloatingPointData {
	return (*FloatingPointData)(&(newAarch64FPState()[0]))
}

// State contains the common architecture bits for aarch64 (the build tag of this
// file ensures it's only built on aarch64).
type State struct {
	// The system registers.
	Regs syscall.PtraceRegs `state:".(syscallPtraceRegs)"`

	// Our floating point state.
	aarch64FPState `state:"wait"`

	// FeatureSet is a pointer to the currently active feature set.
	FeatureSet *cpuid.FeatureSet
}

// Proto returns a protobuf representation of the system registers in State.
func (s State) Proto() *rpb.Registers {
	regs := &rpb.ARM64Registers{
		R0:     s.Regs.Regs[0],
		R1:     s.Regs.Regs[1],
		R2:     s.Regs.Regs[2],
		R3:     s.Regs.Regs[3],
		R4:     s.Regs.Regs[4],
		R5:     s.Regs.Regs[5],
		R6:     s.Regs.Regs[6],
		R7:     s.Regs.Regs[7],
		R8:     s.Regs.Regs[8],
		R9:     s.Regs.Regs[9],
		R10:    s.Regs.Regs[10],
		R11:    s.Regs.Regs[11],
		R12:    s.Regs.Regs[12],
		R13:    s.Regs.Regs[13],
		R14:    s.Regs.Regs[14],
		R15:    s.Regs.Regs[15],
		R16:    s.Regs.Regs[16],
		R17:    s.Regs.Regs[17],
		R18:    s.Regs.Regs[18],
		R19:    s.Regs.Regs[19],
		R20:    s.Regs.Regs[20],
		R21:    s.Regs.Regs[21],
		R22:    s.Regs.Regs[22],
		R23:    s.Regs.Regs[23],
		R24:    s.Regs.Regs[24],
		R25:    s.Regs.Regs[25],
		R26:    s.Regs.Regs[26],
		R27:    s.Regs.Regs[27],
		R28:    s.Regs.Regs[28],
		R29:    s.Regs.Regs[29],
		R30:    s.Regs.Regs[30],
		Sp:     s.Regs.Sp,
		Pc:     s.Regs.Pc,
		Pstate: s.Regs.Pstate,
	}
	return &rpb.Registers{Arch: &rpb.Registers_Arm64{Arm64: regs}}
}

// Fork creates and returns an identical copy of the state.
func (s *State) Fork() State {
	// TODO(gvisor.dev/issue/1238): floating-point is not supported.
	return State{
		Regs:       s.Regs,
		FeatureSet: s.FeatureSet,
	}
}

// StateData implements Context.StateData.
func (s *State) StateData() *State {
	return s
}

// CPUIDEmulate emulates a cpuid instruction.
func (s *State) CPUIDEmulate(l log.Logger) {
	// TODO(gvisor.dev/issue/1255): cpuid is not supported.
}

// SingleStep implements Context.SingleStep.
func (s *State) SingleStep() bool {
	return false
}

// SetSingleStep enables single stepping.
func (s *State) SetSingleStep() {
	// Set the trap flag.
	// TODO(gvisor.dev/issue/1239): ptrace single-step is not supported.
}

// ClearSingleStep enables single stepping.
func (s *State) ClearSingleStep() {
	// Clear the trap flag.
	// TODO(gvisor.dev/issue/1239): ptrace single-step is not supported.
}

// RegisterMap returns a map of all registers.
func (s *State) RegisterMap() (map[string]uintptr, error) {
	return map[string]uintptr{
		"R0":     uintptr(s.Regs.Regs[0]),
		"R1":     uintptr(s.Regs.Regs[1]),
		"R2":     uintptr(s.Regs.Regs[2]),
		"R3":     uintptr(s.Regs.Regs[3]),
		"R4":     uintptr(s.Regs.Regs[4]),
		"R5":     uintptr(s.Regs.Regs[5]),
		"R6":     uintptr(s.Regs.Regs[6]),
		"R7":     uintptr(s.Regs.Regs[7]),
		"R8":     uintptr(s.Regs.Regs[8]),
		"R9":     uintptr(s.Regs.Regs[9]),
		"R10":    uintptr(s.Regs.Regs[10]),
		"R11":    uintptr(s.Regs.Regs[11]),
		"R12":    uintptr(s.Regs.Regs[12]),
		"R13":    uintptr(s.Regs.Regs[13]),
		"R14":    uintptr(s.Regs.Regs[14]),
		"R15":    uintptr(s.Regs.Regs[15]),
		"R16":    uintptr(s.Regs.Regs[16]),
		"R17":    uintptr(s.Regs.Regs[17]),
		"R18":    uintptr(s.Regs.Regs[18]),
		"R19":    uintptr(s.Regs.Regs[19]),
		"R20":    uintptr(s.Regs.Regs[20]),
		"R21":    uintptr(s.Regs.Regs[21]),
		"R22":    uintptr(s.Regs.Regs[22]),
		"R23":    uintptr(s.Regs.Regs[23]),
		"R24":    uintptr(s.Regs.Regs[24]),
		"R25":    uintptr(s.Regs.Regs[25]),
		"R26":    uintptr(s.Regs.Regs[26]),
		"R27":    uintptr(s.Regs.Regs[27]),
		"R28":    uintptr(s.Regs.Regs[28]),
		"R29":    uintptr(s.Regs.Regs[29]),
		"R30":    uintptr(s.Regs.Regs[30]),
		"Sp":     uintptr(s.Regs.Sp),
		"Pc":     uintptr(s.Regs.Pc),
		"Pstate": uintptr(s.Regs.Pstate),
	}, nil
}

// PtraceGetRegs implements Context.PtraceGetRegs.
func (s *State) PtraceGetRegs(dst io.Writer) (int, error) {
	return dst.Write(binary.Marshal(nil, usermem.ByteOrder, s.ptraceGetRegs()))
}

func (s *State) ptraceGetRegs() syscall.PtraceRegs {
	return s.Regs
}

var ptraceRegsSize = int(binary.Size(syscall.PtraceRegs{}))

// PtraceSetRegs implements Context.PtraceSetRegs.
func (s *State) PtraceSetRegs(src io.Reader) (int, error) {
	var regs syscall.PtraceRegs
	buf := make([]byte, ptraceRegsSize)
	if _, err := io.ReadFull(src, buf); err != nil {
		return 0, err
	}
	binary.Unmarshal(buf, usermem.ByteOrder, &regs)
	s.Regs = regs
	return ptraceRegsSize, nil
}

// PtraceGetFPRegs implements Context.PtraceGetFPRegs.
func (s *State) PtraceGetFPRegs(dst io.Writer) (int, error) {
	// TODO(gvisor.dev/issue/1238): floating-point is not supported.
	return 0, nil
}

// PtraceSetFPRegs implements Context.PtraceSetFPRegs.
func (s *State) PtraceSetFPRegs(src io.Reader) (int, error) {
	// TODO(gvisor.dev/issue/1238): floating-point is not supported.
	return 0, nil
}

// Register sets defined in include/uapi/linux/elf.h.
const (
	_NT_PRSTATUS = 1
	_NT_PRFPREG  = 2
)

// PtraceGetRegSet implements Context.PtraceGetRegSet.
func (s *State) PtraceGetRegSet(regset uintptr, dst io.Writer, maxlen int) (int, error) {
	switch regset {
	case _NT_PRSTATUS:
		if maxlen < ptraceRegsSize {
			return 0, syserror.EFAULT
		}
		return s.PtraceGetRegs(dst)
	default:
		return 0, syserror.EINVAL
	}
}

// PtraceSetRegSet implements Context.PtraceSetRegSet.
func (s *State) PtraceSetRegSet(regset uintptr, src io.Reader, maxlen int) (int, error) {
	switch regset {
	case _NT_PRSTATUS:
		if maxlen < ptraceRegsSize {
			return 0, syserror.EFAULT
		}
		return s.PtraceSetRegs(src)
	default:
		return 0, syserror.EINVAL
	}
}

// FullRestore indicates whether a full restore is required.
func (s *State) FullRestore() bool {
	return false
}

// New returns a new architecture context.
func New(arch Arch, fs *cpuid.FeatureSet) Context {
	switch arch {
	case ARM64:
		return &context64{
			State{
				FeatureSet: fs,
			},
		}
	}
	panic(fmt.Sprintf("unknown architecture %v", arch))
}
<|endoftext|>


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

package storage

import (
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/pkg/errors"

	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
)

// createTestBookie creates a new bookie, stopper and manual clock for testing.
func createTestBookie(
	reservationTimeout time.Duration, maxReservations int, maxReservedBytes int64,
) *bookie {
	b := newBookie(newStoreMetrics(time.Hour))
	// Lock the bookie to prevent the main loop from running as we change some
	// of the bookie's state.
	b.mu.Lock()
	defer b.mu.Unlock()
	b.maxReservations = maxReservations
	b.maxReservedBytes = maxReservedBytes
	// Set a high number for a mocked total available space.
	b.metrics.Available.Update(defaultMaxReservedBytes * 10)
	return b
}

// verifyBookie ensures that there are the correct number of reservations and
// reserved bytes.
func verifyBookie(t *testing.T, b *bookie, reservations int, reservedBytes int64) {
	if e, a := reservedBytes, b.metrics.Reserved.Count(); e != a {
		t.Error(errors.Errorf("expected total bytes reserved to be %d, got %d", e, a))
	}
	if e, a := reservations, int(b.metrics.ReservedReplicaCount.Count()); e != a {
		t.Error(errors.Errorf("expected total reservations to be %d, got %d", e, a))
	}
}

// TestBookieReserve ensures that you can never have more than one reservation
// for a specific rangeID at a time, and that both `Reserve` and `Fill` function
// correctly.
func TestBookieReserve(t *testing.T) {
	defer leaktest.AfterTest(t)()
	b := createTestBookie(time.Hour, 5, defaultMaxReservedBytes)

	testCases := []struct {
		rangeID      int
		reserve      bool                   // true for reserve, false for fill
		expSuc       bool                   // is the operation expected to succeed
		expOut       int                    // expected number of reserved replicas
		expBytes     int64                  // expected number of bytes being reserved
		deadReplicas []roachpb.ReplicaIdent // dead replicas that we should not reserve over
	}{
		{rangeID: 1, reserve: true, expSuc: true, expOut: 1, expBytes: 1},
		{rangeID: 1, reserve: true, expSuc: false, expOut: 1, expBytes: 1},
		{rangeID: 1, reserve: false, expSuc: true, expOut: 0, expBytes: 0},
		{rangeID: 1, reserve: false, expSuc: false, expOut: 0, expBytes: 0},
		{rangeID: 2, reserve: true, expSuc: true, expOut: 1, expBytes: 2},
		{rangeID: 3, reserve: true, expSuc: true, expOut: 2, expBytes: 5},
		{rangeID: 1, reserve: true, expSuc: true, expOut: 3, expBytes: 6},
		{rangeID: 2, reserve: true, expSuc: false, expOut: 3, expBytes: 6},
		{rangeID: 2, reserve: false, expSuc: true, expOut: 2, expBytes: 4},
		{rangeID: 2, reserve: false, expSuc: false, expOut: 2, expBytes: 4},
		{rangeID: 3, reserve: false, expSuc: true, expOut: 1, expBytes: 1},
		{rangeID: 1, reserve: false, expSuc: true, expOut: 0, expBytes: 0},
		{rangeID: 2, reserve: false, expSuc: false, expOut: 0, expBytes: 0},
		{rangeID: 0, reserve: true, expSuc: false, expOut: 0, expBytes: 0, deadReplicas: []roachpb.ReplicaIdent{{RangeID: 0}}},
		{rangeID: 0, reserve: true, expSuc: true, expOut: 1, expBytes: 0, deadReplicas: []roachpb.ReplicaIdent{{RangeID: 1}}},
		{rangeID: 0, reserve: false, expSuc: true, expOut: 0, expBytes: 0},
	}

	ctx := context.Background()
	for i, testCase := range testCases {
		if testCase.reserve {
			// Try to reserve the range.
			req := reservationRequest{
				StoreRequestHeader: StoreRequestHeader{
					StoreID: roachpb.StoreID(i),
					NodeID:  roachpb.NodeID(i),
				},
				RangeID:   roachpb.RangeID(testCase.rangeID),
				RangeSize: int64(testCase.rangeID),
			}
			if resp := b.Reserve(ctx, req, testCase.deadReplicas); resp.Reserved != testCase.expSuc {
				if testCase.expSuc {
					t.Errorf("%d: expected a successful reservation, was rejected", i)
				} else {
					t.Errorf("%d: expected no reservation, but it was accepted", i)
				}
			}
		} else {
			// Fill the reservation.
			if filled := b.Fill(ctx, roachpb.RangeID(testCase.rangeID)); filled != testCase.expSuc {
				if testCase.expSuc {
					t.Errorf("%d: expected a successful filled reservation, was rejected", i)
				} else {
					t.Errorf("%d: expected no reservation to be filled, but it was accepted", i)
				}
			}
		}

		verifyBookie(t, b, testCase.expOut, testCase.expBytes)
	}

	// Test that repeated requests with the same store and node number extend
	// the timeout of the pre-existing reservation.
	repeatReq := reservationRequest{
		StoreRequestHeader: StoreRequestHeader{
			StoreID: 100,
			NodeID:  100,
		},
		RangeID:   100,
		RangeSize: 100,
	}
	for i := 1; i < 10; i++ {
		if !b.Reserve(context.Background(), repeatReq, nil).Reserved {
			t.Errorf("%d: could not add repeated reservation", i)
		}
		verifyBookie(t, b, 1, 100)
	}

	// Test rejecting a reservation due to disk space constraints.
	overfilledReq := reservationRequest{
		StoreRequestHeader: StoreRequestHeader{
			StoreID: 200,
			NodeID:  200,
		},
		RangeID:   200,
		RangeSize: 200,
	}

	b.mu.Lock()
	// Set the bytes have 1 less byte free than needed by the reservation.
	b.metrics.Available.Update(b.mu.size + (2 * overfilledReq.RangeSize) - 1)
	b.mu.Unlock()

	if b.Reserve(context.Background(), overfilledReq, nil).Reserved {
		t.Errorf("expected reservation to fail due to disk space constraints, but it succeeded")
	}
	verifyBookie(t, b, 1, 100) // The same numbers from the last call to verifyBookie.
}

// TestBookieReserveMaxRanges ensures that over-booking doesn't occur when there
// are already maxReservations.
func TestBookieReserveMaxRanges(t *testing.T) {
	defer leaktest.AfterTest(t)()

	previousReserved := 10

	b := createTestBookie(time.Hour, previousReserved, defaultMaxReservedBytes)

	// Load up reservations.
	for i := 1; i <= previousReserved; i++ {
		req := reservationRequest{
			StoreRequestHeader: StoreRequestHeader{
				StoreID: roachpb.StoreID(i),
				NodeID:  roachpb.NodeID(i),
			},
			RangeID:   roachpb.RangeID(i),
			RangeSize: 1,
		}
		if !b.Reserve(context.Background(), req, nil).Reserved {
			t.Errorf("%d: could not add reservation", i)
		}
		verifyBookie(t, b, i, int64(i))
	}

	overbookedReq := reservationRequest{
		StoreRequestHeader: StoreRequestHeader{
			StoreID: roachpb.StoreID(previousReserved + 1),
			NodeID:  roachpb.NodeID(previousReserved + 1),
		},
		RangeID:   roachpb.RangeID(previousReserved + 1),
		RangeSize: 1,
	}
	if b.Reserve(context.Background(), overbookedReq, nil).Reserved {
		t.Errorf("expected reservation to fail due to too many already existing reservations, but it succeeded")
	}
	// The same numbers from the last call to verifyBookie.
	verifyBookie(t, b, previousReserved, int64(previousReserved))
}

// TestBookieReserveMaxBytes ensures that over-booking doesn't occur when trying
// to reserve more bytes than maxReservedBytes.
func TestBookieReserveMaxBytes(t *testing.T) {
	defer leaktest.AfterTest(t)()

	previousReservedBytes := 10

	b := createTestBookie(time.Hour, previousReservedBytes*2, int64(previousReservedBytes))

	// Load up reservations with a size of 1 each.
	for i := 1; i <= previousReservedBytes; i++ {
		req := reservationRequest{
			StoreRequestHeader: StoreRequestHeader{
				StoreID: roachpb.StoreID(i),
				NodeID:  roachpb.NodeID(i),
			},
			RangeID:   roachpb.RangeID(i),
			RangeSize: 1,
		}
		if !b.Reserve(context.Background(), req, nil).Reserved {
			t.Errorf("%d: could not add reservation", i)
		}
		verifyBookie(t, b, i, int64(i))
	}

	overbookedReq := reservationRequest{
		StoreRequestHeader: StoreRequestHeader{
			StoreID: roachpb.StoreID(previousReservedBytes + 1),
			NodeID:  roachpb.NodeID(previousReservedBytes + 1),
		},
		RangeID:   roachpb.RangeID(previousReservedBytes + 1),
		RangeSize: 1,
	}
	if b.Reserve(context.Background(), overbookedReq, nil).Reserved {
		t.Errorf("expected reservation to fail due to too many already existing reservations, but it succeeded")
	}
	// The same numbers from the last call to verifyBookie.
	verifyBookie(t, b, previousReservedBytes, int64(previousReservedBytes))
}
<commit_msg>storage: remove reservationTimeout from test<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.

package storage

import (
	"math"
	"testing"

	"github.com/pkg/errors"
	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
)

func createTestBookie(maxReservations int, maxReservedBytes int64) *bookie {
	b := newBookie(newStoreMetrics(math.MaxInt32))
	b.maxReservations = maxReservations
	b.maxReservedBytes = maxReservedBytes
	// Set a high number for a mocked total available space.
	b.metrics.Available.Update(defaultMaxReservedBytes * 10)
	return b
}

// verifyBookie ensures that there are the correct number of reservations and
// reserved bytes.
func verifyBookie(t *testing.T, b *bookie, reservations int, reservedBytes int64) {
	if e, a := reservedBytes, b.metrics.Reserved.Count(); e != a {
		t.Error(errors.Errorf("expected total bytes reserved to be %d, got %d", e, a))
	}
	if e, a := reservations, int(b.metrics.ReservedReplicaCount.Count()); e != a {
		t.Error(errors.Errorf("expected total reservations to be %d, got %d", e, a))
	}
}

// TestBookieReserve ensures that you can never have more than one reservation
// for a specific rangeID at a time, and that both `Reserve` and `Fill` function
// correctly.
func TestBookieReserve(t *testing.T) {
	defer leaktest.AfterTest(t)()
	b := createTestBookie(5, defaultMaxReservedBytes)

	testCases := []struct {
		rangeID      int
		reserve      bool                   // true for reserve, false for fill
		expSuc       bool                   // is the operation expected to succeed
		expOut       int                    // expected number of reserved replicas
		expBytes     int64                  // expected number of bytes being reserved
		deadReplicas []roachpb.ReplicaIdent // dead replicas that we should not reserve over
	}{
		{rangeID: 1, reserve: true, expSuc: true, expOut: 1, expBytes: 1},
		{rangeID: 1, reserve: true, expSuc: false, expOut: 1, expBytes: 1},
		{rangeID: 1, reserve: false, expSuc: true, expOut: 0, expBytes: 0},
		{rangeID: 1, reserve: false, expSuc: false, expOut: 0, expBytes: 0},
		{rangeID: 2, reserve: true, expSuc: true, expOut: 1, expBytes: 2},
		{rangeID: 3, reserve: true, expSuc: true, expOut: 2, expBytes: 5},
		{rangeID: 1, reserve: true, expSuc: true, expOut: 3, expBytes: 6},
		{rangeID: 2, reserve: true, expSuc: false, expOut: 3, expBytes: 6},
		{rangeID: 2, reserve: false, expSuc: true, expOut: 2, expBytes: 4},
		{rangeID: 2, reserve: false, expSuc: false, expOut: 2, expBytes: 4},
		{rangeID: 3, reserve: false, expSuc: true, expOut: 1, expBytes: 1},
		{rangeID: 1, reserve: false, expSuc: true, expOut: 0, expBytes: 0},
		{rangeID: 2, reserve: false, expSuc: false, expOut: 0, expBytes: 0},
		{rangeID: 0, reserve: true, expSuc: false, expOut: 0, expBytes: 0, deadReplicas: []roachpb.ReplicaIdent{{RangeID: 0}}},
		{rangeID: 0, reserve: true, expSuc: true, expOut: 1, expBytes: 0, deadReplicas: []roachpb.ReplicaIdent{{RangeID: 1}}},
		{rangeID: 0, reserve: false, expSuc: true, expOut: 0, expBytes: 0},
	}

	ctx := context.Background()
	for i, testCase := range testCases {
		if testCase.reserve {
			// Try to reserve the range.
			req := reservationRequest{
				StoreRequestHeader: StoreRequestHeader{
					StoreID: roachpb.StoreID(i),
					NodeID:  roachpb.NodeID(i),
				},
				RangeID:   roachpb.RangeID(testCase.rangeID),
				RangeSize: int64(testCase.rangeID),
			}
			if resp := b.Reserve(ctx, req, testCase.deadReplicas); resp.Reserved != testCase.expSuc {
				if testCase.expSuc {
					t.Errorf("%d: expected a successful reservation, was rejected", i)
				} else {
					t.Errorf("%d: expected no reservation, but it was accepted", i)
				}
			}
		} else {
			// Fill the reservation.
			if filled := b.Fill(ctx, roachpb.RangeID(testCase.rangeID)); filled != testCase.expSuc {
				if testCase.expSuc {
					t.Errorf("%d: expected a successful filled reservation, was rejected", i)
				} else {
					t.Errorf("%d: expected no reservation to be filled, but it was accepted", i)
				}
			}
		}

		verifyBookie(t, b, testCase.expOut, testCase.expBytes)
	}

	// Test that repeated requests with the same store and node number extend
	// the timeout of the pre-existing reservation.
	repeatReq := reservationRequest{
		StoreRequestHeader: StoreRequestHeader{
			StoreID: 100,
			NodeID:  100,
		},
		RangeID:   100,
		RangeSize: 100,
	}
	for i := 1; i < 10; i++ {
		if !b.Reserve(context.Background(), repeatReq, nil).Reserved {
			t.Errorf("%d: could not add repeated reservation", i)
		}
		verifyBookie(t, b, 1, 100)
	}

	// Test rejecting a reservation due to disk space constraints.
	overfilledReq := reservationRequest{
		StoreRequestHeader: StoreRequestHeader{
			StoreID: 200,
			NodeID:  200,
		},
		RangeID:   200,
		RangeSize: 200,
	}

	b.mu.Lock()
	// Set the bytes have 1 less byte free than needed by the reservation.
	b.metrics.Available.Update(b.mu.size + (2 * overfilledReq.RangeSize) - 1)
	b.mu.Unlock()

	if b.Reserve(context.Background(), overfilledReq, nil).Reserved {
		t.Errorf("expected reservation to fail due to disk space constraints, but it succeeded")
	}
	verifyBookie(t, b, 1, 100) // The same numbers from the last call to verifyBookie.
}

// TestBookieReserveMaxRanges ensures that over-booking doesn't occur when there
// are already maxReservations.
func TestBookieReserveMaxRanges(t *testing.T) {
	defer leaktest.AfterTest(t)()

	previousReserved := 10

	b := createTestBookie(previousReserved, defaultMaxReservedBytes)

	// Load up reservations.
	for i := 1; i <= previousReserved; i++ {
		req := reservationRequest{
			StoreRequestHeader: StoreRequestHeader{
				StoreID: roachpb.StoreID(i),
				NodeID:  roachpb.NodeID(i),
			},
			RangeID:   roachpb.RangeID(i),
			RangeSize: 1,
		}
		if !b.Reserve(context.Background(), req, nil).Reserved {
			t.Errorf("%d: could not add reservation", i)
		}
		verifyBookie(t, b, i, int64(i))
	}

	overbookedReq := reservationRequest{
		StoreRequestHeader: StoreRequestHeader{
			StoreID: roachpb.StoreID(previousReserved + 1),
			NodeID:  roachpb.NodeID(previousReserved + 1),
		},
		RangeID:   roachpb.RangeID(previousReserved + 1),
		RangeSize: 1,
	}
	if b.Reserve(context.Background(), overbookedReq, nil).Reserved {
		t.Errorf("expected reservation to fail due to too many already existing reservations, but it succeeded")
	}
	// The same numbers from the last call to verifyBookie.
	verifyBookie(t, b, previousReserved, int64(previousReserved))
}

// TestBookieReserveMaxBytes ensures that over-booking doesn't occur when trying
// to reserve more bytes than maxReservedBytes.
func TestBookieReserveMaxBytes(t *testing.T) {
	defer leaktest.AfterTest(t)()

	previousReservedBytes := 10

	b := createTestBookie(previousReservedBytes*2, int64(previousReservedBytes))

	// Load up reservations with a size of 1 each.
	for i := 1; i <= previousReservedBytes; i++ {
		req := reservationRequest{
			StoreRequestHeader: StoreRequestHeader{
				StoreID: roachpb.StoreID(i),
				NodeID:  roachpb.NodeID(i),
			},
			RangeID:   roachpb.RangeID(i),
			RangeSize: 1,
		}
		if !b.Reserve(context.Background(), req, nil).Reserved {
			t.Errorf("%d: could not add reservation", i)
		}
		verifyBookie(t, b, i, int64(i))
	}

	overbookedReq := reservationRequest{
		StoreRequestHeader: StoreRequestHeader{
			StoreID: roachpb.StoreID(previousReservedBytes + 1),
			NodeID:  roachpb.NodeID(previousReservedBytes + 1),
		},
		RangeID:   roachpb.RangeID(previousReservedBytes + 1),
		RangeSize: 1,
	}
	if b.Reserve(context.Background(), overbookedReq, nil).Reserved {
		t.Errorf("expected reservation to fail due to too many already existing reservations, but it succeeded")
	}
	// The same numbers from the last call to verifyBookie.
	verifyBookie(t, b, previousReservedBytes, int64(previousReservedBytes))
}
<|endoftext|>


<commit_before>package schema

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"sync"

	"github.com/mitchellh/mapstructure"
)

// MapFieldWriter writes data into a single map[string]string structure.
type MapFieldWriter struct {
	Schema map[string]*Schema

	lock   sync.Mutex
	result map[string]string
}

// Map returns the underlying map that is being written to.
func (w *MapFieldWriter) Map() map[string]string {
	w.lock.Lock()
	defer w.lock.Unlock()
	if w.result == nil {
		w.result = make(map[string]string)
	}

	return w.result
}

func (w *MapFieldWriter) WriteField(addr []string, value interface{}) error {
	w.lock.Lock()
	defer w.lock.Unlock()
	if w.result == nil {
		w.result = make(map[string]string)
	}

	schemaList := addrToSchema(addr, w.Schema)
	if len(schemaList) == 0 {
		return fmt.Errorf("Invalid address to set: %#v", addr)
	}

	// If we're setting anything other than a list root or set root,
	// then disallow it.
	for _, schema := range schemaList[:len(schemaList)-1] {
		if schema.Type == TypeList {
			return fmt.Errorf(
				"%s: can only set full list",
				strings.Join(addr, "."))
		}

		if schema.Type == TypeMap {
			return fmt.Errorf(
				"%s: can only set full map",
				strings.Join(addr, "."))
		}

		if schema.Type == TypeSet {
			return fmt.Errorf(
				"%s: can only set full set",
				strings.Join(addr, "."))
		}
	}

	return w.set(addr, value)
}

func (w *MapFieldWriter) set(addr []string, value interface{}) error {
	schemaList := addrToSchema(addr, w.Schema)
	if len(schemaList) == 0 {
		return fmt.Errorf("Invalid address to set: %#v", addr)
	}

	schema := schemaList[len(schemaList)-1]
	switch schema.Type {
	case TypeBool, TypeInt, TypeFloat, TypeString:
		return w.setPrimitive(addr, value, schema)
	case TypeList:
		return w.setList(addr, value, schema)
	case TypeMap:
		return w.setMap(addr, value, schema)
	case TypeSet:
		return w.setSet(addr, value, schema)
	case typeObject:
		return w.setObject(addr, value, schema)
	default:
		panic(fmt.Sprintf("Unknown type: %#v", schema.Type))
	}
}

func (w *MapFieldWriter) setList(
	addr []string,
	v interface{},
	schema *Schema) error {
	k := strings.Join(addr, ".")
	setElement := func(idx string, value interface{}) error {
		addrCopy := make([]string, len(addr), len(addr)+1)
		copy(addrCopy, addr)
		return w.set(append(addrCopy, idx), value)
	}

	var vs []interface{}
	if err := mapstructure.Decode(v, &vs); err != nil {
		return fmt.Errorf("%s: %s", k, err)
	}

	// Set the entire list.
	var err error
	for i, elem := range vs {
		is := strconv.FormatInt(int64(i), 10)
		err = setElement(is, elem)
		if err != nil {
			break
		}
	}
	if err != nil {
		for i, _ := range vs {
			is := strconv.FormatInt(int64(i), 10)
			setElement(is, nil)
		}

		return err
	}

	w.result[k+".#"] = strconv.FormatInt(int64(len(vs)), 10)
	return nil
}

func (w *MapFieldWriter) setMap(
	addr []string,
	value interface{},
	schema *Schema) error {
	k := strings.Join(addr, ".")
	v := reflect.ValueOf(value)
	vs := make(map[string]interface{})

	if value == nil {
		// The empty string here means the map is removed.
		w.result[k] = ""
		return nil
	}

	if v.Kind() != reflect.Map {
		return fmt.Errorf("%s: must be a map", k)
	}
	if v.Type().Key().Kind() != reflect.String {
		return fmt.Errorf("%s: keys must strings", k)
	}
	for _, mk := range v.MapKeys() {
		mv := v.MapIndex(mk)
		vs[mk.String()] = mv.Interface()
	}

	// Remove the pure key since we're setting the full map value
	delete(w.result, k)

	// Set each subkey
	addrCopy := make([]string, len(addr), len(addr)+1)
	copy(addrCopy, addr)
	for subKey, v := range vs {
		if err := w.set(append(addrCopy, subKey), v); err != nil {
			return err
		}
	}

	// Set the count
	w.result[k+".#"] = strconv.Itoa(len(vs))

	return nil
}

func (w *MapFieldWriter) setObject(
	addr []string,
	value interface{},
	schema *Schema) error {
	// Set the entire object. First decode into a proper structure
	var v map[string]interface{}
	if err := mapstructure.Decode(value, &v); err != nil {
		return fmt.Errorf("%s: %s", strings.Join(addr, "."), err)
	}

	// Make space for additional elements in the address
	addrCopy := make([]string, len(addr), len(addr)+1)
	copy(addrCopy, addr)

	// Set each element in turn
	var err error
	for k1, v1 := range v {
		if err = w.set(append(addrCopy, k1), v1); err != nil {
			break
		}
	}
	if err != nil {
		for k1, _ := range v {
			w.set(append(addrCopy, k1), nil)
		}
	}

	return err
}

func (w *MapFieldWriter) setPrimitive(
	addr []string,
	v interface{},
	schema *Schema) error {
	k := strings.Join(addr, ".")

	if v == nil {
		delete(w.result, k)
		return nil
	}

	var set string
	switch schema.Type {
	case TypeBool:
		var b bool
		if err := mapstructure.Decode(v, &b); err != nil {
			return fmt.Errorf("%s: %s", k, err)
		}

		set = strconv.FormatBool(b)
	case TypeString:
		if err := mapstructure.Decode(v, &set); err != nil {
			return fmt.Errorf("%s: %s", k, err)
		}
	case TypeInt:
		var n int
		if err := mapstructure.Decode(v, &n); err != nil {
			return fmt.Errorf("%s: %s", k, err)
		}
		set = strconv.FormatInt(int64(n), 10)
	case TypeFloat:
		var n float64
		if err := mapstructure.Decode(v, &n); err != nil {
			return fmt.Errorf("%s: %s", k, err)
		}
		set = strconv.FormatFloat(float64(n), 'G', -1, 64)
	default:
		return fmt.Errorf("Unknown type: %#v", schema.Type)
	}

	w.result[k] = set
	return nil
}

func (w *MapFieldWriter) setSet(
	addr []string,
	value interface{},
	schema *Schema) error {
	addrCopy := make([]string, len(addr), len(addr)+1)
	copy(addrCopy, addr)
	k := strings.Join(addr, ".")

	if value == nil {
		w.result[k+".#"] = "0"
		return nil
	}

	// If it is a slice, then we have to turn it into a *Set so that
	// we get the proper order back based on the hash code.
	if v := reflect.ValueOf(value); v.Kind() == reflect.Slice {
		// Build a temp *ResourceData to use for the conversion
		tempSchema := *schema
		tempSchema.Type = TypeList
		tempSchemaMap := map[string]*Schema{addr[0]: &tempSchema}
		tempW := &MapFieldWriter{Schema: tempSchemaMap}

		// Set the entire list, this lets us get sane values out of it
		if err := tempW.WriteField(addr, value); err != nil {
			return err
		}

		// Build the set by going over the list items in order and
		// hashing them into the set. The reason we go over the list and
		// not the `value` directly is because this forces all types
		// to become []interface{} (generic) instead of []string, which
		// most hash functions are expecting.
		s := &Set{F: schema.Set}
		tempR := &MapFieldReader{
			Map:    BasicMapReader(tempW.Map()),
			Schema: tempSchemaMap,
		}
		for i := 0; i < v.Len(); i++ {
			is := strconv.FormatInt(int64(i), 10)
			result, err := tempR.ReadField(append(addrCopy, is))
			if err != nil {
				return err
			}
			if !result.Exists {
				panic("set item just set doesn't exist")
			}

			s.Add(result.Value)
		}

		value = s
	}

	for code, elem := range value.(*Set).m {
		codeStr := strconv.FormatInt(int64(code), 10)
		if err := w.set(append(addrCopy, codeStr), elem); err != nil {
			return err
		}
	}

	w.result[k+".#"] = strconv.Itoa(value.(*Set).Len())
	return nil
}
<commit_msg>schema: delete non existing values<commit_after>package schema

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"sync"

	"github.com/mitchellh/mapstructure"
)

// MapFieldWriter writes data into a single map[string]string structure.
type MapFieldWriter struct {
	Schema map[string]*Schema

	lock   sync.Mutex
	result map[string]string
}

// Map returns the underlying map that is being written to.
func (w *MapFieldWriter) Map() map[string]string {
	w.lock.Lock()
	defer w.lock.Unlock()
	if w.result == nil {
		w.result = make(map[string]string)
	}

	return w.result
}

func (w *MapFieldWriter) WriteField(addr []string, value interface{}) error {
	w.lock.Lock()
	defer w.lock.Unlock()
	if w.result == nil {
		w.result = make(map[string]string)
	}

	schemaList := addrToSchema(addr, w.Schema)
	if len(schemaList) == 0 {
		return fmt.Errorf("Invalid address to set: %#v", addr)
	}

	// If we're setting anything other than a list root or set root,
	// then disallow it.
	for _, schema := range schemaList[:len(schemaList)-1] {
		if schema.Type == TypeList {
			return fmt.Errorf(
				"%s: can only set full list",
				strings.Join(addr, "."))
		}

		if schema.Type == TypeMap {
			return fmt.Errorf(
				"%s: can only set full map",
				strings.Join(addr, "."))
		}

		if schema.Type == TypeSet {
			return fmt.Errorf(
				"%s: can only set full set",
				strings.Join(addr, "."))
		}
	}

	return w.set(addr, value)
}

func (w *MapFieldWriter) set(addr []string, value interface{}) error {
	schemaList := addrToSchema(addr, w.Schema)
	if len(schemaList) == 0 {
		return fmt.Errorf("Invalid address to set: %#v", addr)
	}

	schema := schemaList[len(schemaList)-1]
	switch schema.Type {
	case TypeBool, TypeInt, TypeFloat, TypeString:
		return w.setPrimitive(addr, value, schema)
	case TypeList:
		return w.setList(addr, value, schema)
	case TypeMap:
		return w.setMap(addr, value, schema)
	case TypeSet:
		return w.setSet(addr, value, schema)
	case typeObject:
		return w.setObject(addr, value, schema)
	default:
		panic(fmt.Sprintf("Unknown type: %#v", schema.Type))
	}
}

func (w *MapFieldWriter) setList(
	addr []string,
	v interface{},
	schema *Schema) error {
	k := strings.Join(addr, ".")
	setElement := func(idx string, value interface{}) error {
		addrCopy := make([]string, len(addr), len(addr)+1)
		copy(addrCopy, addr)
		return w.set(append(addrCopy, idx), value)
	}

	var vs []interface{}
	if err := mapstructure.Decode(v, &vs); err != nil {
		return fmt.Errorf("%s: %s", k, err)
	}

	// Set the entire list.
	var err error
	for i, elem := range vs {
		is := strconv.FormatInt(int64(i), 10)
		err = setElement(is, elem)
		if err != nil {
			break
		}
	}
	if err != nil {
		for i, _ := range vs {
			is := strconv.FormatInt(int64(i), 10)
			setElement(is, nil)
		}

		return err
	}

	w.result[k+".#"] = strconv.FormatInt(int64(len(vs)), 10)
	return nil
}

func (w *MapFieldWriter) setMap(
	addr []string,
	value interface{},
	schema *Schema) error {
	k := strings.Join(addr, ".")
	v := reflect.ValueOf(value)
	vs := make(map[string]interface{})

	if value == nil {
		// The empty string here means the map is removed.
		w.result[k] = ""
		return nil
	}

	if v.Kind() != reflect.Map {
		return fmt.Errorf("%s: must be a map", k)
	}
	if v.Type().Key().Kind() != reflect.String {
		return fmt.Errorf("%s: keys must strings", k)
	}
	for _, mk := range v.MapKeys() {
		mv := v.MapIndex(mk)
		vs[mk.String()] = mv.Interface()
	}

	// Remove the pure key since we're setting the full map value
	delete(w.result, k)

	// Set each subkey
	addrCopy := make([]string, len(addr), len(addr)+1)
	copy(addrCopy, addr)
	for subKey, v := range vs {
		if err := w.set(append(addrCopy, subKey), v); err != nil {
			return err
		}
	}

	// Set the count
	w.result[k+".#"] = strconv.Itoa(len(vs))

	return nil
}

func (w *MapFieldWriter) setObject(
	addr []string,
	value interface{},
	schema *Schema) error {
	// Set the entire object. First decode into a proper structure
	var v map[string]interface{}
	if err := mapstructure.Decode(value, &v); err != nil {
		return fmt.Errorf("%s: %s", strings.Join(addr, "."), err)
	}

	// Make space for additional elements in the address
	addrCopy := make([]string, len(addr), len(addr)+1)
	copy(addrCopy, addr)

	// Set each element in turn
	var err error
	for k1, v1 := range v {
		if err = w.set(append(addrCopy, k1), v1); err != nil {
			break
		}
	}
	if err != nil {
		for k1, _ := range v {
			w.set(append(addrCopy, k1), nil)
		}
	}

	return err
}

func (w *MapFieldWriter) setPrimitive(
	addr []string,
	v interface{},
	schema *Schema) error {
	k := strings.Join(addr, ".")

	if v == nil {
		// The empty string here means the value is removed.
		w.result[k] = ""
		return nil
	}

	var set string
	switch schema.Type {
	case TypeBool:
		var b bool
		if err := mapstructure.Decode(v, &b); err != nil {
			return fmt.Errorf("%s: %s", k, err)
		}

		set = strconv.FormatBool(b)
	case TypeString:
		if err := mapstructure.Decode(v, &set); err != nil {
			return fmt.Errorf("%s: %s", k, err)
		}
	case TypeInt:
		var n int
		if err := mapstructure.Decode(v, &n); err != nil {
			return fmt.Errorf("%s: %s", k, err)
		}
		set = strconv.FormatInt(int64(n), 10)
	case TypeFloat:
		var n float64
		if err := mapstructure.Decode(v, &n); err != nil {
			return fmt.Errorf("%s: %s", k, err)
		}
		set = strconv.FormatFloat(float64(n), 'G', -1, 64)
	default:
		return fmt.Errorf("Unknown type: %#v", schema.Type)
	}

	w.result[k] = set
	return nil
}

func (w *MapFieldWriter) setSet(
	addr []string,
	value interface{},
	schema *Schema) error {
	addrCopy := make([]string, len(addr), len(addr)+1)
	copy(addrCopy, addr)
	k := strings.Join(addr, ".")

	if value == nil {
		w.result[k+".#"] = "0"
		return nil
	}

	// If it is a slice, then we have to turn it into a *Set so that
	// we get the proper order back based on the hash code.
	if v := reflect.ValueOf(value); v.Kind() == reflect.Slice {
		// Build a temp *ResourceData to use for the conversion
		tempSchema := *schema
		tempSchema.Type = TypeList
		tempSchemaMap := map[string]*Schema{addr[0]: &tempSchema}
		tempW := &MapFieldWriter{Schema: tempSchemaMap}

		// Set the entire list, this lets us get sane values out of it
		if err := tempW.WriteField(addr, value); err != nil {
			return err
		}

		// Build the set by going over the list items in order and
		// hashing them into the set. The reason we go over the list and
		// not the `value` directly is because this forces all types
		// to become []interface{} (generic) instead of []string, which
		// most hash functions are expecting.
		s := &Set{F: schema.Set}
		tempR := &MapFieldReader{
			Map:    BasicMapReader(tempW.Map()),
			Schema: tempSchemaMap,
		}
		for i := 0; i < v.Len(); i++ {
			is := strconv.FormatInt(int64(i), 10)
			result, err := tempR.ReadField(append(addrCopy, is))
			if err != nil {
				return err
			}
			if !result.Exists {
				panic("set item just set doesn't exist")
			}

			s.Add(result.Value)
		}

		value = s
	}

	for code, elem := range value.(*Set).m {
		codeStr := strconv.FormatInt(int64(code), 10)
		if err := w.set(append(addrCopy, codeStr), elem); err != nil {
			return err
		}
	}

	w.result[k+".#"] = strconv.Itoa(value.(*Set).Len())
	return nil
}
<|endoftext|>


<commit_before>package sensu

import (
	"bytes"
	"encoding/json"
	"github.com/upfluence/sensu-client-go/sensu/check"
	"log"
	"strconv"
	"strings"
	"time"
)

const (
	MAX_FAILS = 100
	MAX_TIME  = 60 * time.Second
)

type Subscriber struct {
	Subscription string
	Client       *Client
}

func (s *Subscriber) SetClient(c *Client) error {
	s.Client = c

	return nil
}

func (s *Subscriber) Start() error {
	var output check.CheckOutput
	var b []byte

	funnel := strings.Join(
		[]string{
			s.Client.Config.Name(),
			CurrentVersion,
			strconv.Itoa(int(time.Now().Unix())),
		},
		"-",
	)

	msgChan := make(chan []byte)
	stopChan := make(chan bool)

	go s.Client.Transport.Subscribe("#", s.Subscription, funnel, msgChan, stopChan)

	log.Printf("Subscribed to %s", s.Subscription)

	failures := 0

	for {
		if failures >= MAX_FAILS {
			stopChan <- true
			msgChan = make(chan []byte)
			failures = 0
			time.Sleep(MAX_TIME)
			go s.Client.Transport.Subscribe(
				"#",
				s.Subscription,
				funnel,
				msgChan,
				stopChan,
			)

		}

		b = <-msgChan

		payload := make(map[string]interface{})

		log.Printf("Check received : %s", bytes.NewBuffer(b).String())
		json.Unmarshal(b, &payload)

		if _, ok := payload["name"]; !ok {
			log.Printf("The name field is not filled")
			failures++
			continue
		}

		if ch, ok := check.Store[payload["name"].(string)]; ok {
			output = ch.Execute()
		} else if _, ok := payload["command"]; !ok {
			log.Printf("The command field is not filled")
			continue
		} else {
			output = (&check.ExternalCheck{payload["command"].(string)}).Execute()
		}

		p, err := json.Marshal(s.forgeCheckResponse(payload, &output))

		if err != nil {
			log.Printf("something goes wrong : %s", err.Error())
		} else {
			log.Printf("Payload sent: %s", bytes.NewBuffer(p).String())

			err = s.Client.Transport.Publish("direct", "results", "", p)
		}
	}

	return nil
}

func (s *Subscriber) forgeCheckResponse(payload map[string]interface{}, output *check.CheckOutput) map[string]interface{} {
	result := make(map[string]interface{})

	result["client"] = s.Client.Config.Name()

	formattedOuput := make(map[string]interface{})

	formattedOuput["name"] = payload["name"]
	formattedOuput["issued"] = int(payload["issued"].(float64))
	formattedOuput["output"] = output.Output
	formattedOuput["duration"] = output.Duration
	formattedOuput["status"] = output.Status
	formattedOuput["executed"] = output.Executed

	result["check"] = formattedOuput

	return result
}
<commit_msg>kill connection and reconnect if failures<commit_after>package sensu

import (
	"bytes"
	"encoding/json"
	"github.com/upfluence/sensu-client-go/sensu/check"
	"log"
	"strconv"
	"strings"
	"time"
)

const (
	MAX_FAILS = 100
	MAX_TIME  = 60 * time.Second
)

type Subscriber struct {
	Subscription string
	Client       *Client
}

func (s *Subscriber) SetClient(c *Client) error {
	s.Client = c

	return nil
}

func (s *Subscriber) Start() error {
	var output check.CheckOutput
	var b []byte

	funnel := strings.Join(
		[]string{
			s.Client.Config.Name(),
			CurrentVersion,
			strconv.Itoa(int(time.Now().Unix())),
		},
		"-",
	)

	msgChan := make(chan []byte)
	stopChan := make(chan bool)

	go s.Client.Transport.Subscribe("#", s.Subscription, funnel, msgChan, stopChan)

	log.Printf("Subscribed to %s", s.Subscription)

	failures := 0

	for {
		if failures >= MAX_FAILS {
			stopChan <- true
			msgChan = make(chan []byte)
			failures = 0
			s.Client.Transport.Close()
			time.Sleep(MAX_TIME)
			s.Client.Transport.Connect()
			go s.Client.Transport.Subscribe(
				"#",
				s.Subscription,
				funnel,
				msgChan,
				stopChan,
			)

		}

		b = <-msgChan

		payload := make(map[string]interface{})

		log.Printf("Check received : %s", bytes.NewBuffer(b).String())
		json.Unmarshal(b, &payload)

		if _, ok := payload["name"]; !ok {
			log.Printf("The name field is not filled")
			failures++
			continue
		}

		if ch, ok := check.Store[payload["name"].(string)]; ok {
			output = ch.Execute()
		} else if _, ok := payload["command"]; !ok {
			log.Printf("The command field is not filled")
			continue
		} else {
			output = (&check.ExternalCheck{payload["command"].(string)}).Execute()
		}

		p, err := json.Marshal(s.forgeCheckResponse(payload, &output))

		if err != nil {
			log.Printf("something goes wrong : %s", err.Error())
		} else {
			log.Printf("Payload sent: %s", bytes.NewBuffer(p).String())

			err = s.Client.Transport.Publish("direct", "results", "", p)
		}
	}

	return nil
}

func (s *Subscriber) forgeCheckResponse(payload map[string]interface{}, output *check.CheckOutput) map[string]interface{} {
	result := make(map[string]interface{})

	result["client"] = s.Client.Config.Name()

	formattedOuput := make(map[string]interface{})

	formattedOuput["name"] = payload["name"]
	formattedOuput["issued"] = int(payload["issued"].(float64))
	formattedOuput["output"] = output.Output
	formattedOuput["duration"] = output.Duration
	formattedOuput["status"] = output.Status
	formattedOuput["executed"] = output.Executed

	result["check"] = formattedOuput

	return result
}
<|endoftext|>


<commit_before>package sensu

import (
	"bytes"
	"encoding/json"
	"github.com/upfluence/sensu-client-go/sensu/check"
	"log"
	"strconv"
	"strings"
	"time"
)

type Subscriber struct {
	Subscription string
	Client       *Client
}

func (s *Subscriber) SetClient(c *Client) error {
	s.Client = c

	return nil
}

func (s *Subscriber) Start() error {
	var output check.CheckOutput
	var b []byte

	funnel := strings.Join(
		[]string{
			s.Client.Config.Name(),
			CurrentVersion,
			strconv.Itoa(int(time.Now().Unix())),
		},
		"-",
	)

	msgChan := make(chan []byte)
	stopChan := make(chan bool)

	go s.Client.Transport.Subscribe("#", s.Subscription, funnel, msgChan, stopChan)

	log.Printf("Subscribed to %s", s.Subscription)

	for {
		b = <-msgChan

		payload := make(map[string]interface{})

		log.Printf("Check received : %s", bytes.NewBuffer(b).String())
		json.Unmarshal(b, &payload)

		if ch, ok := check.Store[payload["name"].(string)]; ok {
			output = ch.Execute()
		} else {
			output = (&check.ExternalCheck{payload["command"].(string)}).Execute()
		}

		p, err := json.Marshal(s.forgeCheckResponse(payload, &output))

		if err != nil {
			log.Printf("something goes wrong : %s", err.Error())
		} else {
			log.Printf("Payload sent: %s", bytes.NewBuffer(p).String())

			err = s.Client.Transport.Publish("direct", "results", "", p)
		}
	}

	return nil
}

func (s *Subscriber) forgeCheckResponse(payload map[string]interface{}, output *check.CheckOutput) map[string]interface{} {
	result := make(map[string]interface{})

	result["client"] = s.Client.Config.Name()

	formattedOuput := make(map[string]interface{})

	formattedOuput["name"] = payload["name"]
	formattedOuput["issued"] = int(payload["issued"].(float64))
	formattedOuput["output"] = output.Output
	formattedOuput["duration"] = output.Duration
	formattedOuput["status"] = output.Status
	formattedOuput["executed"] = output.Executed

	result["check"] = formattedOuput

	return result
}
<commit_msg>Check if the command exists before execute it<commit_after>package sensu

import (
	"bytes"
	"encoding/json"
	"github.com/upfluence/sensu-client-go/sensu/check"
	"log"
	"strconv"
	"strings"
	"time"
)

type Subscriber struct {
	Subscription string
	Client       *Client
}

func (s *Subscriber) SetClient(c *Client) error {
	s.Client = c

	return nil
}

func (s *Subscriber) Start() error {
	var output check.CheckOutput
	var b []byte

	funnel := strings.Join(
		[]string{
			s.Client.Config.Name(),
			CurrentVersion,
			strconv.Itoa(int(time.Now().Unix())),
		},
		"-",
	)

	msgChan := make(chan []byte)
	stopChan := make(chan bool)

	go s.Client.Transport.Subscribe("#", s.Subscription, funnel, msgChan, stopChan)

	log.Printf("Subscribed to %s", s.Subscription)

	for {
		b = <-msgChan

		payload := make(map[string]interface{})

		log.Printf("Check received : %s", bytes.NewBuffer(b).String())
		json.Unmarshal(b, &payload)

		if ch, ok := check.Store[payload["name"].(string)]; ok {
			output = ch.Execute()
		} else if _, ok := payload["command"]; !ok {
			log.Printf("The command field is not filled")

			continue
		} else {
			output = (&check.ExternalCheck{payload["command"].(string)}).Execute()
		}

		p, err := json.Marshal(s.forgeCheckResponse(payload, &output))

		if err != nil {
			log.Printf("something goes wrong : %s", err.Error())
		} else {
			log.Printf("Payload sent: %s", bytes.NewBuffer(p).String())

			err = s.Client.Transport.Publish("direct", "results", "", p)
		}
	}

	return nil
}

func (s *Subscriber) forgeCheckResponse(payload map[string]interface{}, output *check.CheckOutput) map[string]interface{} {
	result := make(map[string]interface{})

	result["client"] = s.Client.Config.Name()

	formattedOuput := make(map[string]interface{})

	formattedOuput["name"] = payload["name"]
	formattedOuput["issued"] = int(payload["issued"].(float64))
	formattedOuput["output"] = output.Output
	formattedOuput["duration"] = output.Duration
	formattedOuput["status"] = output.Status
	formattedOuput["executed"] = output.Executed

	result["check"] = formattedOuput

	return result
}
<|endoftext|>


<commit_before>package shell

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"

	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// An inline script to execute. Multiple strings are all executed
	// in the context of a single shell.
	Inline []string `mapstructure:"inline"`

	// The shebang value used when running inline scripts.
	InlineShebang string `mapstructure:"inline_shebang"`

	// The local path of the shell script to upload and execute.
	Script string `mapstructure:"script"`

	// An array of environment variables that will be injected before
	// your command(s) are executed.
	Vars []string `mapstructure:"environment_vars"`

	// An array of multiple scripts to run.
	Scripts []string `mapstructure:"scripts"`

	TargetPath string `mapstructure:"target"`

	tpl *packer.ConfigTemplate
}

type ShellPostProcessor struct {
	cfg Config
}

type OutputPathTemplate struct {
	ArtifactId string
	BuildName  string
	Provider   string
}

func (p *ShellPostProcessor) Configure(raws ...interface{}) error {
	_, err := common.DecodeConfig(&p.cfg, raws...)
	if err != nil {
		return err
	}

	errs := new(packer.MultiError)

	if p.cfg.InlineShebang == "" {
		p.cfg.InlineShebang = "/bin/sh"
	}

	if p.cfg.Scripts == nil {
		p.cfg.Scripts = make([]string, 0)
	}

	if p.cfg.Vars == nil {
		p.cfg.Vars = make([]string, 0)
	}

	if p.cfg.Script != "" && len(p.cfg.Scripts) > 0 {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only one of script or scripts can be specified."))
	}

	if p.cfg.Script != "" {
		p.cfg.Scripts = []string{p.cfg.Script}
	}

	p.cfg.tpl, err = packer.NewConfigTemplate()
	if err != nil {
		return err
	}
	p.cfg.tpl.UserVars = p.cfg.PackerUserVars

	if p.cfg.TargetPath == "" {
		p.cfg.TargetPath = "packer_{{ .BuildName }}_{{.Provider}}"
	}

	if err = p.cfg.tpl.Validate(p.cfg.TargetPath); err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Error parsing target template: %s", err))
	}

	templates := map[string]*string{
		"inline_shebang": &p.cfg.InlineShebang,
		"script":         &p.cfg.Script,
	}

	for n, ptr := range templates {
		var err error
		*ptr, err = p.cfg.tpl.Process(*ptr, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", n, err))
		}
	}

	sliceTemplates := map[string][]string{
		"inline":           p.cfg.Inline,
		"scripts":          p.cfg.Scripts,
		"environment_vars": p.cfg.Vars,
	}

	for n, slice := range sliceTemplates {
		for i, elem := range slice {
			var err error
			slice[i], err = p.cfg.tpl.Process(elem, nil)
			if err != nil {
				errs = packer.MultiErrorAppend(
					errs, fmt.Errorf("Error processing %s[%d]: %s", n, i, err))
			}
		}
	}

	if len(p.cfg.Scripts) == 0 && p.cfg.Inline == nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Either a script file or inline script must be specified."))
	} else if len(p.cfg.Scripts) > 0 && p.cfg.Inline != nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only a script file or an inline script can be specified, not both."))
	}

	for _, path := range p.cfg.Scripts {
		if _, err := os.Stat(path); err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Bad script '%s': %s", path, err))
		}
	}

	// Do a check for bad environment variables, such as '=foo', 'foobar'
	for _, kv := range p.cfg.Vars {
		vs := strings.SplitN(kv, "=", 2)
		if len(vs) != 2 || vs[0] == "" {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Environment variable not in format 'key=value': %s", kv))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *ShellPostProcessor) PostProcess(ui packer.Ui, artifact packer.Artifact) (packer.Artifact, bool, error) {
	scripts := make([]string, len(p.cfg.Scripts))
	copy(scripts, p.cfg.Scripts)

	if p.cfg.Inline != nil {
		tf, err := ioutil.TempFile("", "packer-shell")
		if err != nil {
			return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
		}
		defer os.Remove(tf.Name())

		// Set the path to the temporary file
		scripts = append(scripts, tf.Name())

		// Write our contents to it
		writer := bufio.NewWriter(tf)
		writer.WriteString(fmt.Sprintf("#!%s\n", p.cfg.InlineShebang))
		for _, command := range p.cfg.Inline {
			if _, err := writer.WriteString(command + "\n"); err != nil {
				return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
			}
		}

		if err := writer.Flush(); err != nil {
			return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
		}

		tf.Close()
	}

	envVars := make([]string, len(p.cfg.Vars)+2)
	envVars[0] = "PACKER_BUILD_NAME=" + p.cfg.PackerBuildName
	envVars[1] = "PACKER_BUILDER_TYPE=" + p.cfg.PackerBuilderType
	copy(envVars[2:], p.cfg.Vars)

	for _, artifact := range artifact.Files() {

		for _, path := range scripts {
			ui.Say(fmt.Sprintf("Process with shell script: %s", path))

			log.Printf("Opening %s for reading", path)
			f, err := os.Open(path)
			if err != nil {
				return nil, false, fmt.Errorf("Error opening shell script: %s", err)
			}
			defer f.Close()

			ui.Message(fmt.Sprintf("Executing script with artifact: %s", artifact))
			cmd := exec.Command("/bin/sh", []string{path}...)
			var buffer bytes.Buffer
			cmd.Stdout = &buffer
			cmd.Stderr = &buffer
			cmd.Env = envVars
			err = cmd.Run()
			if err != nil {
				return nil, false, fmt.Errorf("Unable to execute script: %s", buffer.String())
			}
			ui.Message(fmt.Sprintf("%s", buffer.String()))
		}
	}
	return artifact, false, nil
}
<commit_msg>fix run<commit_after>package shell

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"

	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// An inline script to execute. Multiple strings are all executed
	// in the context of a single shell.
	Inline []string `mapstructure:"inline"`

	// The shebang value used when running inline scripts.
	InlineShebang string `mapstructure:"inline_shebang"`

	// The local path of the shell script to upload and execute.
	Script string `mapstructure:"script"`

	// An array of environment variables that will be injected before
	// your command(s) are executed.
	Vars []string `mapstructure:"environment_vars"`

	// An array of multiple scripts to run.
	Scripts []string `mapstructure:"scripts"`

	TargetPath string `mapstructure:"target"`

	tpl *packer.ConfigTemplate
}

type ShellPostProcessor struct {
	cfg Config
}

type OutputPathTemplate struct {
	ArtifactId string
	BuildName  string
	Provider   string
}

func (p *ShellPostProcessor) Configure(raws ...interface{}) error {
	_, err := common.DecodeConfig(&p.cfg, raws...)
	if err != nil {
		return err
	}

	errs := new(packer.MultiError)

	if p.cfg.InlineShebang == "" {
		p.cfg.InlineShebang = "/bin/sh"
	}

	if p.cfg.Scripts == nil {
		p.cfg.Scripts = make([]string, 0)
	}

	if p.cfg.Vars == nil {
		p.cfg.Vars = make([]string, 0)
	}

	if p.cfg.Script != "" && len(p.cfg.Scripts) > 0 {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only one of script or scripts can be specified."))
	}

	if p.cfg.Script != "" {
		p.cfg.Scripts = []string{p.cfg.Script}
	}

	p.cfg.tpl, err = packer.NewConfigTemplate()
	if err != nil {
		return err
	}
	p.cfg.tpl.UserVars = p.cfg.PackerUserVars

	if p.cfg.TargetPath == "" {
		p.cfg.TargetPath = "packer_{{ .BuildName }}_{{.Provider}}"
	}

	if err = p.cfg.tpl.Validate(p.cfg.TargetPath); err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Error parsing target template: %s", err))
	}

	templates := map[string]*string{
		"inline_shebang": &p.cfg.InlineShebang,
		"script":         &p.cfg.Script,
	}

	for n, ptr := range templates {
		var err error
		*ptr, err = p.cfg.tpl.Process(*ptr, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", n, err))
		}
	}

	sliceTemplates := map[string][]string{
		"inline":           p.cfg.Inline,
		"scripts":          p.cfg.Scripts,
		"environment_vars": p.cfg.Vars,
	}

	for n, slice := range sliceTemplates {
		for i, elem := range slice {
			var err error
			slice[i], err = p.cfg.tpl.Process(elem, nil)
			if err != nil {
				errs = packer.MultiErrorAppend(
					errs, fmt.Errorf("Error processing %s[%d]: %s", n, i, err))
			}
		}
	}

	if len(p.cfg.Scripts) == 0 && p.cfg.Inline == nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Either a script file or inline script must be specified."))
	} else if len(p.cfg.Scripts) > 0 && p.cfg.Inline != nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only a script file or an inline script can be specified, not both."))
	}

	for _, path := range p.cfg.Scripts {
		if _, err := os.Stat(path); err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Bad script '%s': %s", path, err))
		}
	}

	// Do a check for bad environment variables, such as '=foo', 'foobar'
	for _, kv := range p.cfg.Vars {
		vs := strings.SplitN(kv, "=", 2)
		if len(vs) != 2 || vs[0] == "" {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Environment variable not in format 'key=value': %s", kv))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *ShellPostProcessor) PostProcess(ui packer.Ui, artifact packer.Artifact) (packer.Artifact, bool, error) {
	scripts := make([]string, len(p.cfg.Scripts))
	copy(scripts, p.cfg.Scripts)

	if p.cfg.Inline != nil {
		tf, err := ioutil.TempFile("", "packer-shell")
		if err != nil {
			return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
		}
		defer os.Remove(tf.Name())

		// Set the path to the temporary file
		scripts = append(scripts, tf.Name())

		// Write our contents to it
		writer := bufio.NewWriter(tf)
		writer.WriteString(fmt.Sprintf("#!%s\n", p.cfg.InlineShebang))
		for _, command := range p.cfg.Inline {
			if _, err := writer.WriteString(command + "\n"); err != nil {
				return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
			}
		}

		if err := writer.Flush(); err != nil {
			return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
		}

		tf.Close()
	}

	envVars := make([]string, len(p.cfg.Vars)+2)
	envVars[0] = "PACKER_BUILD_NAME=" + p.cfg.PackerBuildName
	envVars[1] = "PACKER_BUILDER_TYPE=" + p.cfg.PackerBuilderType
	copy(envVars[2:], p.cfg.Vars)

	for _, artifact := range artifact.Files() {

		for _, path := range scripts {
			ui.Say(fmt.Sprintf("Process with shell script: %s", path))

			log.Printf("Opening %s for reading", path)
			f, err := os.Open(path)
			if err != nil {
				return nil, false, fmt.Errorf("Error opening shell script: %s", err)
			}
			defer f.Close()

			ui.Message(fmt.Sprintf("Executing script with artifact: %s", artifact))
			args := []string{path, artifact}
			cmd := exec.Command("/bin/sh", args...)
			var buffer bytes.Buffer
			cmd.Stdout = &buffer
			cmd.Stderr = &buffer
			cmd.Env = envVars
			err = cmd.Run()
			if err != nil {
				return nil, false, fmt.Errorf("Unable to execute script: %s", buffer.String())
			}
			ui.Message(fmt.Sprintf("%s", buffer.String()))
		}
	}
	return artifact, false, nil
}
<|endoftext|>


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

// Package shlib provides the functionality of the mcm-shellify tool.
package shlib

import (
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"strconv"

	"github.com/zombiezen/mcm/catalog"
	"github.com/zombiezen/mcm/internal/depgraph"
	"github.com/zombiezen/mcm/third_party/golang/capnproto"
)

// WriteScript converts a catalog into a bash script and writes it to w.
func WriteScript(w io.Writer, c catalog.Catalog) error {
	g := newGen(w)
	g.p(script("#!/bin/bash"))
	g.p(script("# Autogenerated by mcm-shellify"))
	g.p()
	res, _ := c.Resources()
	if res.Len() == 0 {
		g.p(script("# Empty catalog"))
		return g.ew.err
	}
	graph, err := depgraph.New(res)
	if err != nil {
		return err
	}
	g.p(script("_() {"))
	g.in()
	for i := 0; i < res.Len(); i++ {
		v := resourceStatusVar(res.At(i).ID())
		g.p(script("local"), assignment{v, -2})
	}
	for g.ew.err == nil && !graph.Done() {
		ready := append([]uint64(nil), graph.Ready()...)
		if len(ready) == 0 {
			return errors.New("graph not done, but has nothing to do")
		}
		for _, id := range ready {
			if err := g.resource(graph.Resource(id)); err != nil {
				return fmt.Errorf("resource ID=%d: %v", id, err)
			}
			graph.Mark(id)
		}
	}
	g.exitStatusCheck(res)
	g.out()
	g.p(script("}"))
	g.p(script(`_ "$0" "$@"`))
	return g.ew.err
}

func resourceStatusVar(id uint64) string {
	return fmt.Sprintf("status%d", id)
}

func (g *gen) resource(r catalog.Resource) error {
	g.p()
	id := r.ID()
	if c, _ := r.Comment(); c != "" {
		// TODO(someday): trim newlines?
		g.p(script("#"), script(c))
	} else {
		g.p(script(fmt.Sprintf("# Resource ID=%d", id)))
	}
	statVar := resourceStatusVar(id)
	if deps, _ := r.Dependencies(); deps.Len() > 0 {
		g.p(script("if [["), depsPrecondition(deps), script("]]; then"))
		g.in()
		defer func() {
			g.out()
			g.p(script("else"))
			g.in()
			g.p(assignment{statVar, -1})
			g.out()
			g.p(script("fi"))
		}()
	}
	switch r.Which() {
	case catalog.Resource_Which_noop:
		if deps, _ := r.Dependencies(); deps.Len() > 0 {
			g.p(depsChangedCondition(deps), script("&&"), assignment{statVar, 1}, script("||"), assignment{statVar, 0})
		} else {
			g.p(assignment{statVar, 0})
		}
		return nil
	case catalog.Resource_Which_file:
		f, err := r.File()
		if err != nil {
			return fmt.Errorf("read from catalog: %v", err)
		}
		return g.file(id, f)
	default:
		return fmt.Errorf("unsupported resource %v", r.Which())
	}
}

func depsPrecondition(deps capnp.UInt64List) script {
	var buf []byte
	for i, n := 0, deps.Len(); i < n; i++ {
		if i > 0 {
			buf = append(buf, " && "...)
		}
		buf = append(buf, "status"...)
		buf = strconv.AppendUint(buf, deps.At(i), 10)
		buf = append(buf, " -ge 0"...)
	}
	return script(buf)
}

func depsChangedCondition(deps capnp.UInt64List) script {
	var buf []byte
	for i, n := 0, deps.Len(); i < n; i++ {
		if i > 0 {
			buf = append(buf, " || "...)
		}
		buf = append(buf, "status"...)
		buf = strconv.AppendUint(buf, deps.At(i), 10)
		buf = append(buf, " -gt 0"...)
	}
	return script(buf)
}

func (g *gen) exitStatusCheck(res catalog.Resource_List) {
	var buf []byte
	for i, n := 0, res.Len(); i < n; i++ {
		if i > 0 {
			buf = append(buf, " || "...)
		}
		buf = append(buf, "status"...)
		buf = strconv.AppendUint(buf, res.At(i).ID(), 10)
		buf = append(buf, " -lt 0"...)
	}
	g.p(script("if [["), script(buf), script("]]; then"))
	g.in()
	g.p(script("return 1"))
	g.out()
	g.p(script("fi"))
	g.p(script("return 0"))
}

func updateStatus(id uint64) script {
	var buf []byte
	buf = append(buf, "&& "...)
	v := resourceStatusVar(id)
	buf = appendPArg(buf, assignment{v, 1})
	buf = append(buf, " || "...)
	buf = appendPArg(buf, assignment{v, -1})
	return script(buf)
}

func (g *gen) file(id uint64, f catalog.File) error {
	path, err := f.Path()
	if err != nil {
		return fmt.Errorf("reading file path: %v", err)
	} else if path == "" {
		return errors.New("file path is empty")
	}
	switch f.Which() {
	case catalog.File_Which_plain:
		// TODO(soon): handle no content case
		// TODO(soon): respect file mode
		if f.Plain().HasContent() {
			content, err := f.Plain().Content()
			if err != nil {
				return fmt.Errorf("read content from catalog: %v", err)
			}
			enc := make([]byte, base64.StdEncoding.EncodedLen(len(content)))
			base64.StdEncoding.Encode(enc, content)
			g.p(script("("))
			g.p(script("base64 -d >"), path, heredoc{marker: "!EOF!", data: enc})
			// TODO(now): check existing content and lstat
			g.p(script(")"), updateStatus(id))
		}
	case catalog.File_Which_directory:
		// TODO(soon): respect file mode
		g.p(script("if [[ ! -e"), path, script("]]; then"))
		g.in()
		g.p(script("mkdir"), path, updateStatus(id))
		g.out()
		g.p(script("elif [[ -d"), path, script("]]; then"))
		g.in()
		g.p(assignment{resourceStatusVar(id), 0})
		g.out()
		g.p(script("else"))
		g.in()
		g.p(script("echo"), path, script("'is not a directory' 1>&2"))
		g.p(assignment{resourceStatusVar(id), -1})
		g.out()
		g.p(script("fi"))
	case catalog.File_Which_symlink:
		target, _ := f.Symlink().Target()
		if target == "" {
			return errors.New("symlink target is empty")
		}
		g.p(script("if [[ -h"), path, script("]]; then"))
		g.in()
		g.p(script("if [[ \"$(readlink"), path, script(")\" !="), target, script("]]; then"))
		g.in()
		g.p(script("ln -f -s"), target, path, updateStatus(id))
		g.out()
		g.p(script("else"))
		g.in()
		g.p(assignment{resourceStatusVar(id), 0})
		g.out()
		g.p(script("fi"))
		g.out()
		g.p(script("elif [[ ! -e"), path, script("]]; then"))
		g.in()
		g.p(script("ln -s"), target, path, updateStatus(id))
		g.out()
		g.p(script("else"))
		g.in()
		g.p(script("echo"), path, script("'is not a symlink' 1>&2"))
		g.p(assignment{resourceStatusVar(id), -1})
		g.out()
		g.p(script("fi"))
	default:
		return fmt.Errorf("unsupported file directive %v", f.Which())
	}
	return nil
}
<commit_msg>shellify: one function per resource<commit_after>// Copyright 2016 The Minimal Configuration Manager Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package shlib provides the functionality of the mcm-shellify tool.
package shlib

import (
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"strconv"

	"github.com/zombiezen/mcm/catalog"
	"github.com/zombiezen/mcm/internal/depgraph"
	"github.com/zombiezen/mcm/third_party/golang/capnproto"
)

// WriteScript converts a catalog into a bash script and writes it to w.
func WriteScript(w io.Writer, c catalog.Catalog) error {
	g := newGen(w)
	g.p(script("#!/bin/bash"))
	g.p(script("# Autogenerated by mcm-shellify"))
	g.p()
	res, _ := c.Resources()
	if res.Len() == 0 {
		g.p(script("# Empty catalog"))
		return g.ew.err
	}
	graph, err := depgraph.New(res)
	if err != nil {
		return err
	}
	for i := 0; i < res.Len(); i++ {
		v := resourceStatusVar(res.At(i).ID())
		g.p(assignment{v, -2})
	}
	for i := 0; i < res.Len(); i++ {
		r := res.At(i)
		if err := g.resourceFunc(r); err != nil {
			return fmt.Errorf("resource ID=%d: %v", r.ID(), err)
		}
	}
	g.p(script("_() {"))
	g.in()
	for g.ew.err == nil && !graph.Done() {
		ready := append([]uint64(nil), graph.Ready()...)
		if len(ready) == 0 {
			return errors.New("graph not done, but has nothing to do")
		}
		for _, id := range ready {
			graph.Mark(id)
			deps, _ := graph.Resource(id).Dependencies()
			if deps.Len() == 0 {
				g.p(resourceFuncName(id))
				continue
			}
			g.p(script("if [["), depsPrecondition(deps), script("]]; then"))
			g.in()
			g.p(resourceFuncName(id))
			g.out()
			g.p(script("else"))
			g.in()
			g.p(assignment{resourceStatusVar(id), -1})
			g.out()
			g.p(script("fi"))
		}
	}
	g.exitStatusCheck(res)
	g.out()
	g.p(script("}"))
	g.p(script(`_ "$0" "$@"`))
	return g.ew.err
}

func resourceStatusVar(id uint64) string {
	return fmt.Sprintf("status%d", id)
}

func resourceFuncName(id uint64) script {
	return script(fmt.Sprintf("resource%d", id))
}

func (g *gen) resourceFunc(r catalog.Resource) error {
	id := r.ID()
	if c, _ := r.Comment(); c != "" {
		// TODO(someday): trim newlines?
		g.p(script("#"), script(c))
	}
	g.p(script(resourceFuncName(id) + "() {"))
	defer g.p(script("}"))
	g.in()
	defer g.out()

	statVar := resourceStatusVar(id)
	switch r.Which() {
	case catalog.Resource_Which_noop:
		if deps, _ := r.Dependencies(); deps.Len() > 0 {
			g.p(depsChangedCondition(deps), script("&&"), assignment{statVar, 1}, script("||"), assignment{statVar, 0})
		} else {
			g.p(assignment{statVar, 0})
		}
		return nil
	case catalog.Resource_Which_file:
		f, err := r.File()
		if err != nil {
			return fmt.Errorf("read from catalog: %v", err)
		}
		return g.file(id, f)
	default:
		return fmt.Errorf("unsupported resource %v", r.Which())
	}
}

func depsPrecondition(deps capnp.UInt64List) script {
	var buf []byte
	for i, n := 0, deps.Len(); i < n; i++ {
		if i > 0 {
			buf = append(buf, " && "...)
		}
		buf = append(buf, "status"...)
		buf = strconv.AppendUint(buf, deps.At(i), 10)
		buf = append(buf, " -ge 0"...)
	}
	return script(buf)
}

func depsChangedCondition(deps capnp.UInt64List) script {
	var buf []byte
	for i, n := 0, deps.Len(); i < n; i++ {
		if i > 0 {
			buf = append(buf, " || "...)
		}
		buf = append(buf, "status"...)
		buf = strconv.AppendUint(buf, deps.At(i), 10)
		buf = append(buf, " -gt 0"...)
	}
	return script(buf)
}

func (g *gen) exitStatusCheck(res catalog.Resource_List) {
	var buf []byte
	for i, n := 0, res.Len(); i < n; i++ {
		if i > 0 {
			buf = append(buf, " || "...)
		}
		buf = append(buf, "status"...)
		buf = strconv.AppendUint(buf, res.At(i).ID(), 10)
		buf = append(buf, " -lt 0"...)
	}
	g.p(script("if [["), script(buf), script("]]; then"))
	g.in()
	g.p(script("return 1"))
	g.out()
	g.p(script("fi"))
	g.p(script("return 0"))
}

func updateStatus(id uint64) script {
	var buf []byte
	buf = append(buf, "&& "...)
	v := resourceStatusVar(id)
	buf = appendPArg(buf, assignment{v, 1})
	buf = append(buf, " || "...)
	buf = appendPArg(buf, assignment{v, -1})
	return script(buf)
}

func (g *gen) returnStatus(id uint64, val int) {
	g.p(assignment{resourceStatusVar(id), val})
	if val >= 0 {
		g.p(script("return 0"))
	} else {
		g.p(script("return 1"))
	}
}

func resourceFuncReturn(id uint64) script {
	var buf []byte
	buf = append(buf, "[[ $"...)
	buf = append(buf, resourceStatusVar(id)...)
	buf = append(buf, " -eq 0 ]] && return 0 || return 1"...)
	return script(buf)
}

func (g *gen) file(id uint64, f catalog.File) error {
	path, err := f.Path()
	if err != nil {
		return fmt.Errorf("reading file path: %v", err)
	} else if path == "" {
		return errors.New("file path is empty")
	}
	switch f.Which() {
	case catalog.File_Which_plain:
		// TODO(soon): handle no content case
		// TODO(soon): respect file mode
		if f.Plain().HasContent() {
			content, err := f.Plain().Content()
			if err != nil {
				return fmt.Errorf("read content from catalog: %v", err)
			}
			enc := make([]byte, base64.StdEncoding.EncodedLen(len(content)))
			base64.StdEncoding.Encode(enc, content)
			g.p(script("("))
			g.p(script("base64 -d >"), path, heredoc{marker: "!EOF!", data: enc})
			// TODO(soon): check existing content and lstat
			g.p(script(")"), updateStatus(id))
			g.p(resourceFuncReturn(id))
		}
	case catalog.File_Which_directory:
		// TODO(soon): respect file mode
		g.p(script("if [[ -d"), path, script("]]; then"))
		g.in()
		g.returnStatus(id, 0)
		g.out()
		g.p(script("fi"))
		g.p(script("if [[ -e"), path, script("]]; then"))
		g.in()
		g.p(script("echo"), path, script("'is not a directory' 1>&2"))
		g.returnStatus(id, -1)
		g.out()
		g.p("fi")
		g.p(script("mkdir"), path, updateStatus(id))
		g.p(resourceFuncReturn(id))
	case catalog.File_Which_symlink:
		target, _ := f.Symlink().Target()
		if target == "" {
			return errors.New("symlink target is empty")
		}
		g.p(script("if [[ -h"), path, script("]]; then"))
		g.in()
		g.p(script("if [[ \"$(readlink"), path, script(")\" !="), target, script("]]; then"))
		g.in()
		g.p(script("ln -f -s"), target, path, updateStatus(id))
		g.p(resourceFuncReturn(id))
		g.out()
		g.p(script("else"))
		g.in()
		g.returnStatus(id, 0)
		g.out()
		g.p(script("fi"))
		g.out()
		g.p(script("fi"))

		g.p(script("if [[ -e"), path, script("]]; then"))
		g.in()
		g.p(script("echo"), path, script("'is not a symlink' 1>&2"))
		g.returnStatus(id, -1)
		g.out()
		g.p(script("fi"))

		g.p(script("ln -s"), target, path, updateStatus(id))
		g.p(resourceFuncReturn(id))
	default:
		return fmt.Errorf("unsupported file directive %v", f.Which())
	}
	return nil
}
<|endoftext|>