<commit_before>// Copyright 2021 Google LLC
//
// 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 bitlocker provides functionality for managing Bitlocker.
package bitlocker

import (
	"fmt"

	"github.com/scjalliance/comshim"
	"github.com/google/logger"
	"github.com/go-ole/go-ole"
	"github.com/go-ole/go-ole/oleutil"
	"github.com/iamacarpet/go-win64api"
)

var (
	// Test Helpers
	funcBackup       = winapi.BackupBitLockerRecoveryKeys
	funcRecoveryInfo = winapi.GetBitLockerRecoveryInfo
)

// BackupToAD backs up Bitlocker recovery keys to Active Directory.
func BackupToAD() error {
	infos, err := funcRecoveryInfo()
	if err != nil {
		return err
	}
	volIDs := []string{}
	for _, i := range infos {
		if i.ConversionStatus != 1 {
			logger.Warningf("Skipping volume %s due to conversion status (%d).", i.DriveLetter, i.ConversionStatus)
			continue
		}
		logger.Infof("Backing up Bitlocker recovery password for drive %q.", i.DriveLetter)
		volIDs = append(volIDs, i.PersistentVolumeID)
	}
	return funcBackup(volIDs)
}

// Encryption Methods
// https://docs.microsoft.com/en-us/windows/win32/secprov/getencryptionmethod-win32-encryptablevolume
type EncryptionMethod int32

const (
	None EncryptionMethod = iota
	AES128WithDiffuser
	AES256WithDiffuser
	AES128
	AES256
	HardwareEncryption
	XtsAES128
	XtsAES256
)

// Encryption Flags
// https://docs.microsoft.com/en-us/windows/win32/secprov/encrypt-win32-encryptablevolume
type EncryptionFlag int32

const (
	EncryptDataOnly    EncryptionFlag = 0x00000001
	EncryptDemandWipe  EncryptionFlag = 0x00000002
	EncryptSynchronous EncryptionFlag = 0x00010000

	// Error Codes
	ERROR_IO_DEVICE                     int32 = -2147023779
	FVE_E_EDRIVE_INCOMPATIBLE_VOLUME    int32 = -2144272206
	FVE_E_NO_TPM_WITH_PASSPHRASE        int32 = -2144272212
	FVE_E_PASSPHRASE_TOO_LONG           int32 = -2144272214
	FVE_E_POLICY_PASSPHRASE_NOT_ALLOWED int32 = -2144272278
	FVE_E_NOT_DECRYPTED                 int32 = -2144272327
	FVE_E_INVALID_PASSWORD_FORMAT       int32 = -2144272331
	FVE_E_BOOTABLE_CDDVD                int32 = -2144272336
	FVE_E_PROTECTOR_EXISTS              int32 = -2144272335
)

func encryptErrHandler(val int32) error {
	switch val {
	case ERROR_IO_DEVICE:
		return fmt.Errorf("an I/O error has occurred during encryption; the device may need to be reset")
	case FVE_E_EDRIVE_INCOMPATIBLE_VOLUME:
		return fmt.Errorf("the drive specified does not support hardware-based encryption")
	case FVE_E_NO_TPM_WITH_PASSPHRASE:
		return fmt.Errorf("a TPM key protector cannot be added because a password protector exists on the drive")
	case FVE_E_PASSPHRASE_TOO_LONG:
		return fmt.Errorf("the passphrase cannot exceed 256 characters")
	case FVE_E_POLICY_PASSPHRASE_NOT_ALLOWED:
		return fmt.Errorf("Group Policy settings do not permit the creation of a password")
	case FVE_E_NOT_DECRYPTED:
		return fmt.Errorf("the drive must be fully decrypted to complete this operation")
	case FVE_E_INVALID_PASSWORD_FORMAT:
		return fmt.Errorf("the format of the recovery password provided is invalid")
	case FVE_E_BOOTABLE_CDDVD:
		return fmt.Errorf("BitLocker Drive Encryption detected bootable media (CD or DVD) in the computer. " +
			"Remove the media and restart the computer before configuring BitLocker.")
	case FVE_E_PROTECTOR_EXISTS:
		return fmt.Errorf("key protector cannot be added; only one key protector of this type is allowed for this drive")
	default:
		return fmt.Errorf("error code returned during encryption: %d", val)
	}
}

// A Volume tracks an open encryptable volume.
type Volume struct {
	letter  string
	handle  *ole.IDispatch
	wmiIntf *ole.IDispatch
	wmiSvc  *ole.IDispatch
}

// Close frees all resources associated with a volume.
func (v *Volume) Close() {
	v.handle.Release()
	v.wmiIntf.Release()
	v.wmiSvc.Release()
	comshim.Done()
}

// Connect connects to an encryptable volume in order to manage it.
// You must call Close() to release the volume when finished.
//
// Example: bitlocker.Connect("c:")
func Connect(driveLetter string) (Volume, error) {
	comshim.Add(1)
	v := Volume{letter: driveLetter}

	unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
	if err != nil {
		comshim.Done()
		return v, fmt.Errorf("CreateObject: %w", err)
	}
	defer unknown.Release()
	v.wmiIntf, err = unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		comshim.Done()
		return v, fmt.Errorf("QueryInterface: %w", err)
	}
	serviceRaw, err := oleutil.CallMethod(v.wmiIntf, "ConnectServer", nil, `\\.\ROOT\CIMV2\Security\MicrosoftVolumeEncryption`)
	if err != nil {
		v.Close()
		return v, fmt.Errorf("ConnectServer: %w", err)
	}
	v.wmiSvc = serviceRaw.ToIDispatch()

	raw, err := oleutil.CallMethod(v.wmiSvc, "ExecQuery", "SELECT * FROM Win32_EncryptableVolume WHERE DriveLetter = '"+driveLetter+"'")
	if err != nil {
		v.Close()
		return v, fmt.Errorf("ExecQuery: %w", err)
	}
	result := raw.ToIDispatch()
	defer result.Release()

	itemRaw, err := oleutil.CallMethod(result, "ItemIndex", 0)
	if err != nil {
		v.Close()
		return v, fmt.Errorf("failed to fetch result row while processing BitLocker info: %w", err)
	}
	v.handle = itemRaw.ToIDispatch()

	return v, nil
}

// Encrypt encrypts the volume.
//
// Example: vol.Encrypt(bitlocker.XtsAES256, bitlocker.EncryptDataOnly)
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/protectkeywithtpm-win32-encryptablevolume
func (v *Volume) Encrypt(method EncryptionMethod, flags EncryptionFlag) error {
	resultRaw, err := oleutil.CallMethod(v.handle, "Encrypt", int32(method), int32(flags))
	if err != nil {
		return fmt.Errorf("Encrypt(%s): %w", v.letter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("Encrypt(%s): %w", v.letter, encryptErrHandler(val))
	}

	return nil
}

// DiscoveryVolumeType specifies the type of discovery volume to be used by Prepare.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/preparevolume-win32-encryptablevolume
type DiscoveryVolumeType string

const (
	// VolumeTypeNone indicates no discovery volume. This value creates a native BitLocker volume.
	VolumeTypeNone DiscoveryVolumeType = "<none>"
	// VolumeTypeDefault indicates the default behavior.
	VolumeTypeDefault DiscoveryVolumeType = "<default>"
	// VolumeTypeFAT32 creates a FAT32 discovery volume.
	VolumeTypeFAT32 DiscoveryVolumeType = "FAT32"
)

// ForceEncryptionType specifies the encryption type to be used when calling Prepare on the volume.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/preparevolume-win32-encryptablevolume
type ForceEncryptionType int32

const (
	// EncryptionTypeUnspecified indicates that the encryption type is not specified.
	EncryptionTypeUnspecified ForceEncryptionType = 0
	// EncryptionTypeSoftware specifies software encryption.
	EncryptionTypeSoftware ForceEncryptionType = 1
	// EncryptionTypeHardware specifies hardware encryption.
	EncryptionTypeHardware ForceEncryptionType = 2
)

// Prepare prepares a new Bitlocker Volume. This should be called BEFORE any key protectors are added.
//
// Example: vol.Prepare(bitlocker.VolumeTypeDefault, bitlocker.EncryptionTypeHardware)
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/preparevolume-win32-encryptablevolume
func (v *Volume) Prepare(volType DiscoveryVolumeType, encType ForceEncryptionType) error {
	resultRaw, err := oleutil.CallMethod(v.handle, "PrepareVolume", string(volType), int32(encType))
	if err != nil {
		return fmt.Errorf("PrepareVolume(%s): %w", v.letter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("PrepareVolume(%s): %w", v.letter, encryptErrHandler(val))
	}
	return nil
}

// ProtectWithNumericalPassword adds a numerical password key protector.
//
// Leave password as a blank string to have one auto-generated by Windows. (Recommended)
//
// In Powershell this is referred to as a RecoveryPasswordProtector.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/protectkeywithnumericalpassword-win32-encryptablevolume
func (v *Volume) ProtectWithNumericalPassword(password string) error {
	var volumeKeyProtectorID ole.VARIANT
	ole.VariantInit(&volumeKeyProtectorID)
	var resultRaw *ole.VARIANT
	var err error
	if password != "" {
		resultRaw, err = oleutil.CallMethod(v.handle, "ProtectWithNumericalPassword", nil, password, &volumeKeyProtectorID)
	} else {
		resultRaw, err = oleutil.CallMethod(v.handle, "ProtectWithNumericalPassword", nil, nil, &volumeKeyProtectorID)
	}
	if err != nil {
		return fmt.Errorf("ProtectWithNumericalPassword(%s): %w", v.letter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("ProtectWithNumericalPassword(%s): %w", v.letter, encryptErrHandler(val))
	}

	return nil
}

// ProtectWithPassphrase adds a passphrase key protector.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/protectkeywithpassphrase-win32-encryptablevolume
func (v *Volume) ProtectWithPassphrase(passphrase string) error {
	var volumeKeyProtectorID ole.VARIANT
	ole.VariantInit(&volumeKeyProtectorID)
	resultRaw, err := oleutil.CallMethod(v.handle, "ProtectKeyWithPassphrase", nil, passphrase, &volumeKeyProtectorID)
	if err != nil {
		return fmt.Errorf("ProtectWithPassphrase(%s): %w", v.letter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("ProtectWithPassphrase(%s): %w", v.letter, encryptErrHandler(val))
	}

	return nil
}

// ProtectWithTPM adds the TPM key protector.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/protectkeywithtpm-win32-encryptablevolume
func (v *Volume) ProtectWithTPM() error {
	var volumeKeyProtectorID ole.VARIANT
	ole.VariantInit(&volumeKeyProtectorID)
	resultRaw, err := oleutil.CallMethod(v.handle, "ProtectKeyWithTPM", nil, nil, &volumeKeyProtectorID)
	if err != nil {
		return fmt.Errorf("ProtectKeyWithTPM(%s): %w", v.letter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("ProtectKeyWithTPM(%s): %w", v.letter, encryptErrHandler(val))
	}

	return nil
}
<commit_msg>Fix ProtectWithNumericalPassword<commit_after>// Copyright 2021 Google LLC
//
// 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 bitlocker provides functionality for managing Bitlocker.
package bitlocker

import (
	"fmt"

	"github.com/scjalliance/comshim"
	"github.com/google/logger"
	"github.com/go-ole/go-ole"
	"github.com/go-ole/go-ole/oleutil"
	"github.com/iamacarpet/go-win64api"
)

var (
	// Test Helpers
	funcBackup       = winapi.BackupBitLockerRecoveryKeys
	funcRecoveryInfo = winapi.GetBitLockerRecoveryInfo
)

// BackupToAD backs up Bitlocker recovery keys to Active Directory.
func BackupToAD() error {
	infos, err := funcRecoveryInfo()
	if err != nil {
		return err
	}
	volIDs := []string{}
	for _, i := range infos {
		if i.ConversionStatus != 1 {
			logger.Warningf("Skipping volume %s due to conversion status (%d).", i.DriveLetter, i.ConversionStatus)
			continue
		}
		logger.Infof("Backing up Bitlocker recovery password for drive %q.", i.DriveLetter)
		volIDs = append(volIDs, i.PersistentVolumeID)
	}
	return funcBackup(volIDs)
}

// Encryption Methods
// https://docs.microsoft.com/en-us/windows/win32/secprov/getencryptionmethod-win32-encryptablevolume
type EncryptionMethod int32

const (
	None EncryptionMethod = iota
	AES128WithDiffuser
	AES256WithDiffuser
	AES128
	AES256
	HardwareEncryption
	XtsAES128
	XtsAES256
)

// Encryption Flags
// https://docs.microsoft.com/en-us/windows/win32/secprov/encrypt-win32-encryptablevolume
type EncryptionFlag int32

const (
	EncryptDataOnly    EncryptionFlag = 0x00000001
	EncryptDemandWipe  EncryptionFlag = 0x00000002
	EncryptSynchronous EncryptionFlag = 0x00010000

	// Error Codes
	ERROR_IO_DEVICE                     int32 = -2147023779
	FVE_E_EDRIVE_INCOMPATIBLE_VOLUME    int32 = -2144272206
	FVE_E_NO_TPM_WITH_PASSPHRASE        int32 = -2144272212
	FVE_E_PASSPHRASE_TOO_LONG           int32 = -2144272214
	FVE_E_POLICY_PASSPHRASE_NOT_ALLOWED int32 = -2144272278
	FVE_E_NOT_DECRYPTED                 int32 = -2144272327
	FVE_E_INVALID_PASSWORD_FORMAT       int32 = -2144272331
	FVE_E_BOOTABLE_CDDVD                int32 = -2144272336
	FVE_E_PROTECTOR_EXISTS              int32 = -2144272335
)

func encryptErrHandler(val int32) error {
	switch val {
	case ERROR_IO_DEVICE:
		return fmt.Errorf("an I/O error has occurred during encryption; the device may need to be reset")
	case FVE_E_EDRIVE_INCOMPATIBLE_VOLUME:
		return fmt.Errorf("the drive specified does not support hardware-based encryption")
	case FVE_E_NO_TPM_WITH_PASSPHRASE:
		return fmt.Errorf("a TPM key protector cannot be added because a password protector exists on the drive")
	case FVE_E_PASSPHRASE_TOO_LONG:
		return fmt.Errorf("the passphrase cannot exceed 256 characters")
	case FVE_E_POLICY_PASSPHRASE_NOT_ALLOWED:
		return fmt.Errorf("Group Policy settings do not permit the creation of a password")
	case FVE_E_NOT_DECRYPTED:
		return fmt.Errorf("the drive must be fully decrypted to complete this operation")
	case FVE_E_INVALID_PASSWORD_FORMAT:
		return fmt.Errorf("the format of the recovery password provided is invalid")
	case FVE_E_BOOTABLE_CDDVD:
		return fmt.Errorf("BitLocker Drive Encryption detected bootable media (CD or DVD) in the computer. " +
			"Remove the media and restart the computer before configuring BitLocker.")
	case FVE_E_PROTECTOR_EXISTS:
		return fmt.Errorf("key protector cannot be added; only one key protector of this type is allowed for this drive")
	default:
		return fmt.Errorf("error code returned during encryption: %d", val)
	}
}

// A Volume tracks an open encryptable volume.
type Volume struct {
	letter  string
	handle  *ole.IDispatch
	wmiIntf *ole.IDispatch
	wmiSvc  *ole.IDispatch
}

// Close frees all resources associated with a volume.
func (v *Volume) Close() {
	v.handle.Release()
	v.wmiIntf.Release()
	v.wmiSvc.Release()
	comshim.Done()
}

// Connect connects to an encryptable volume in order to manage it.
// You must call Close() to release the volume when finished.
//
// Example: bitlocker.Connect("c:")
func Connect(driveLetter string) (Volume, error) {
	comshim.Add(1)
	v := Volume{letter: driveLetter}

	unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
	if err != nil {
		comshim.Done()
		return v, fmt.Errorf("CreateObject: %w", err)
	}
	defer unknown.Release()
	v.wmiIntf, err = unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		comshim.Done()
		return v, fmt.Errorf("QueryInterface: %w", err)
	}
	serviceRaw, err := oleutil.CallMethod(v.wmiIntf, "ConnectServer", nil, `\\.\ROOT\CIMV2\Security\MicrosoftVolumeEncryption`)
	if err != nil {
		v.Close()
		return v, fmt.Errorf("ConnectServer: %w", err)
	}
	v.wmiSvc = serviceRaw.ToIDispatch()

	raw, err := oleutil.CallMethod(v.wmiSvc, "ExecQuery", "SELECT * FROM Win32_EncryptableVolume WHERE DriveLetter = '"+driveLetter+"'")
	if err != nil {
		v.Close()
		return v, fmt.Errorf("ExecQuery: %w", err)
	}
	result := raw.ToIDispatch()
	defer result.Release()

	itemRaw, err := oleutil.CallMethod(result, "ItemIndex", 0)
	if err != nil {
		v.Close()
		return v, fmt.Errorf("failed to fetch result row while processing BitLocker info: %w", err)
	}
	v.handle = itemRaw.ToIDispatch()

	return v, nil
}

// Encrypt encrypts the volume.
//
// Example: vol.Encrypt(bitlocker.XtsAES256, bitlocker.EncryptDataOnly)
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/protectkeywithtpm-win32-encryptablevolume
func (v *Volume) Encrypt(method EncryptionMethod, flags EncryptionFlag) error {
	resultRaw, err := oleutil.CallMethod(v.handle, "Encrypt", int32(method), int32(flags))
	if err != nil {
		return fmt.Errorf("Encrypt(%s): %w", v.letter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("Encrypt(%s): %w", v.letter, encryptErrHandler(val))
	}

	return nil
}

// DiscoveryVolumeType specifies the type of discovery volume to be used by Prepare.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/preparevolume-win32-encryptablevolume
type DiscoveryVolumeType string

const (
	// VolumeTypeNone indicates no discovery volume. This value creates a native BitLocker volume.
	VolumeTypeNone DiscoveryVolumeType = "<none>"
	// VolumeTypeDefault indicates the default behavior.
	VolumeTypeDefault DiscoveryVolumeType = "<default>"
	// VolumeTypeFAT32 creates a FAT32 discovery volume.
	VolumeTypeFAT32 DiscoveryVolumeType = "FAT32"
)

// ForceEncryptionType specifies the encryption type to be used when calling Prepare on the volume.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/preparevolume-win32-encryptablevolume
type ForceEncryptionType int32

const (
	// EncryptionTypeUnspecified indicates that the encryption type is not specified.
	EncryptionTypeUnspecified ForceEncryptionType = 0
	// EncryptionTypeSoftware specifies software encryption.
	EncryptionTypeSoftware ForceEncryptionType = 1
	// EncryptionTypeHardware specifies hardware encryption.
	EncryptionTypeHardware ForceEncryptionType = 2
)

// Prepare prepares a new Bitlocker Volume. This should be called BEFORE any key protectors are added.
//
// Example: vol.Prepare(bitlocker.VolumeTypeDefault, bitlocker.EncryptionTypeHardware)
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/preparevolume-win32-encryptablevolume
func (v *Volume) Prepare(volType DiscoveryVolumeType, encType ForceEncryptionType) error {
	resultRaw, err := oleutil.CallMethod(v.handle, "PrepareVolume", string(volType), int32(encType))
	if err != nil {
		return fmt.Errorf("PrepareVolume(%s): %w", v.letter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("PrepareVolume(%s): %w", v.letter, encryptErrHandler(val))
	}
	return nil
}

// ProtectWithNumericalPassword adds a numerical password key protector.
//
// Leave password as a blank string to have one auto-generated by Windows. (Recommended)
//
// In Powershell this is referred to as a RecoveryPasswordProtector.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/protectkeywithnumericalpassword-win32-encryptablevolume
func (v *Volume) ProtectWithNumericalPassword(password string) error {
	var volumeKeyProtectorID ole.VARIANT
	ole.VariantInit(&volumeKeyProtectorID)
	var resultRaw *ole.VARIANT
	var err error
	if password != "" {
		resultRaw, err = oleutil.CallMethod(v.handle, "ProtectKeyWithNumericalPassword", nil, password, &volumeKeyProtectorID)
	} else {
		resultRaw, err = oleutil.CallMethod(v.handle, "ProtectKeyWithNumericalPassword", nil, nil, &volumeKeyProtectorID)
	}
	if err != nil {
		return fmt.Errorf("ProtectKeyWithNumericalPassword(%s): %w", v.letter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("ProtectKeyWithNumericalPassword(%s): %w", v.letter, encryptErrHandler(val))
	}

	return nil
}

// ProtectWithPassphrase adds a passphrase key protector.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/protectkeywithpassphrase-win32-encryptablevolume
func (v *Volume) ProtectWithPassphrase(passphrase string) error {
	var volumeKeyProtectorID ole.VARIANT
	ole.VariantInit(&volumeKeyProtectorID)
	resultRaw, err := oleutil.CallMethod(v.handle, "ProtectKeyWithPassphrase", nil, passphrase, &volumeKeyProtectorID)
	if err != nil {
		return fmt.Errorf("ProtectWithPassphrase(%s): %w", v.letter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("ProtectWithPassphrase(%s): %w", v.letter, encryptErrHandler(val))
	}

	return nil
}

// ProtectWithTPM adds the TPM key protector.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/secprov/protectkeywithtpm-win32-encryptablevolume
func (v *Volume) ProtectWithTPM() error {
	var volumeKeyProtectorID ole.VARIANT
	ole.VariantInit(&volumeKeyProtectorID)
	resultRaw, err := oleutil.CallMethod(v.handle, "ProtectKeyWithTPM", nil, nil, &volumeKeyProtectorID)
	if err != nil {
		return fmt.Errorf("ProtectKeyWithTPM(%s): %w", v.letter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("ProtectKeyWithTPM(%s): %w", v.letter, encryptErrHandler(val))
	}

	return nil
}
<|endoftext|>


<commit_before>package log

import (
	"io"
	"log"
	"os"
)

var (
	logger = NewLogger()
)

type Logger struct {
	debug bool
}

func NewLogger() *Logger {
	log.SetFlags(0)
	log.SetOutput(os.Stderr)
	return &Logger{
		debug: false,
	}
}

func (l *Logger) SetDebug(debug bool) {
	l.debug = debug
}

func (l *Logger) Debug(v ...interface{}) {
	if l.debug {
		log.Println(v...)
	}
}

func (l *Logger) Debugf(format string, v ...interface{}) {
	if l.debug {
		log.Printf(format+"\n", v...)
	}
}

func (l *Logger) Println(v ...interface{}) {
	log.Println(v...)
}

func (l *Logger) Printf(format string, v ...interface{}) {
	log.Printf(format+"\n", v...)
}

func SetFlags(flag int) {
	log.SetFlags(flag)
}

func SetOutput(w io.Writer) {
	log.SetOutput(w)
}

func SetDebug(debug bool) {
	logger.SetDebug(debug)
}

func Debug(args ...interface{}) {
	logger.Debug(args...)
}

func Debugf(format string, args ...interface{}) {
	logger.Debugf(format, args...)
}

func Println(args ...interface{}) {
	logger.Println(args...)
}

func Printf(format string, args ...interface{}) {
	logger.Printf(format, args...)
}
<commit_msg>Add log package comment<commit_after>package log

import (
	"io"
	"log"
	"os"
)

var (
	logger = NewLogger()
)

// Logger is logger
type Logger struct {
	debug bool
}

// NewLogger returns a logger object
func NewLogger() *Logger {
	log.SetFlags(0)
	log.SetOutput(os.Stderr)
	return &Logger{
		debug: false,
	}
}

// SetDebug enables debug mode
func (l *Logger) SetDebug(debug bool) {
	l.debug = debug
}

// Debug prints a debug log
func (l *Logger) Debug(v ...interface{}) {
	if l.debug {
		log.Println(v...)
	}
}

// Debugf prints a formatted debug log
func (l *Logger) Debugf(format string, v ...interface{}) {
	if l.debug {
		log.Printf(format+"\n", v...)
	}
}

// Println prints a stdout log
func (l *Logger) Println(v ...interface{}) {
	log.Println(v...)
}

// Printf prints a formatted stdout log
func (l *Logger) Printf(format string, v ...interface{}) {
	log.Printf(format+"\n", v...)
}

// SetFlags set log package's flags
func SetFlags(flag int) {
	log.SetFlags(flag)
}

// SetOutput set output destination
func SetOutput(w io.Writer) {
	log.SetOutput(w)
}

// SetDebug enables debug mode
func SetDebug(debug bool) {
	logger.SetDebug(debug)
}

// Debug prints a debug log
func Debug(args ...interface{}) {
	logger.Debug(args...)
}

// Debugf prints a formatted debug log
func Debugf(format string, args ...interface{}) {
	logger.Debugf(format, args...)
}

// Println prints a stdout log
func Println(args ...interface{}) {
	logger.Println(args...)
}

// Printf prints a formatted stdout log
func Printf(format string, args ...interface{}) {
	logger.Printf(format, args...)
}
<|endoftext|>


<commit_before>package vtadmin

import (
	"context"
	"database/sql"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/vt/vitessdriver"
	"vitess.io/vitess/go/vt/vtadmin/cluster"
	"vitess.io/vitess/go/vt/vtadmin/cluster/discovery/fakediscovery"
	"vitess.io/vitess/go/vt/vtadmin/grpcserver"
	"vitess.io/vitess/go/vt/vtadmin/http"
	"vitess.io/vitess/go/vt/vtadmin/vtsql"
	"vitess.io/vitess/go/vt/vtadmin/vtsql/fakevtsql"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	vtadminpb "vitess.io/vitess/go/vt/proto/vtadmin"
)

func TestGetGates(t *testing.T) {
	fakedisco1 := fakediscovery.New()
	cluster1 := &cluster.Cluster{
		ID:        "c1",
		Name:      "cluster1",
		Discovery: fakedisco1,
	}
	cluster1Gates := []*vtadminpb.VTGate{
		{
			Hostname: "cluster1-gate1",
		},
		{
			Hostname: "cluster1-gate2",
		},
		{
			Hostname: "cluster1-gate3",
		},
	}

	fakedisco1.AddTaggedGates(nil, cluster1Gates...)

	fakedisco2 := fakediscovery.New()
	cluster2 := &cluster.Cluster{
		ID:        "c2",
		Name:      "cluster2",
		Discovery: fakedisco2,
	}
	cluster2Gates := []*vtadminpb.VTGate{
		{
			Hostname: "cluster2-gate1",
		},
	}

	fakedisco2.AddTaggedGates(nil, cluster2Gates...)

	api := NewAPI([]*cluster.Cluster{cluster1, cluster2}, grpcserver.Options{}, http.Options{})
	ctx := context.Background()

	resp, err := api.GetGates(ctx, &vtadminpb.GetGatesRequest{})
	assert.NoError(t, err)
	assert.ElementsMatch(t, append(cluster1Gates, cluster2Gates...), resp.Gates)

	resp, err = api.GetGates(ctx, &vtadminpb.GetGatesRequest{ClusterIds: []string{cluster1.ID}})
	assert.NoError(t, err)
	assert.ElementsMatch(t, cluster1Gates, resp.Gates)

	fakedisco1.SetGatesError(true)

	resp, err = api.GetGates(ctx, &vtadminpb.GetGatesRequest{})
	assert.Error(t, err)
	assert.Nil(t, resp)
}

func TestGetTablets(t *testing.T) {
	type dbcfg struct {
		shouldErr bool
	}

	tests := []struct {
		name           string
		clusterTablets [][]*vtadminpb.Tablet
		dbconfigs      map[string]*dbcfg
		req            *vtadminpb.GetTabletsRequest
		expected       []*vtadminpb.Tablet
		shouldErr      bool
	}{
		{
			name: "single cluster",
			clusterTablets: [][]*vtadminpb.Tablet{
				{
					/* cluster 0 */
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  100,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
			},
			dbconfigs: map[string]*dbcfg{},
			req:       &vtadminpb.GetTabletsRequest{},
			expected: []*vtadminpb.Tablet{
				{
					Cluster: &vtadminpb.Cluster{
						Id:   "c0",
						Name: "cluster0",
					},
					State: vtadminpb.Tablet_SERVING,
					Tablet: &topodatapb.Tablet{
						Alias: &topodatapb.TabletAlias{
							Uid:  100,
							Cell: "zone1",
						},
						Hostname: "ks1-00-00-zone1-a",
						Keyspace: "ks1",
						Shard:    "-",
						Type:     topodatapb.TabletType_MASTER,
					},
				},
			},
			shouldErr: false,
		},
		{
			name: "one cluster errors",
			clusterTablets: [][]*vtadminpb.Tablet{
				/* cluster 0 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  100,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
				/* cluster 1 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  200,
								Cell: "zone1",
							},
							Hostname: "ks2-00-00-zone1-a",
							Keyspace: "ks2",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
			},
			dbconfigs: map[string]*dbcfg{
				"c1": {shouldErr: true},
			},
			req:       &vtadminpb.GetTabletsRequest{},
			expected:  nil,
			shouldErr: true,
		},
		{
			name: "multi cluster, selecting one",
			clusterTablets: [][]*vtadminpb.Tablet{
				/* cluster 0 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  100,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
				/* cluster 1 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  200,
								Cell: "zone1",
							},
							Hostname: "ks2-00-00-zone1-a",
							Keyspace: "ks2",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
			},
			dbconfigs: map[string]*dbcfg{},
			req:       &vtadminpb.GetTabletsRequest{ClusterIds: []string{"c0"}},
			expected: []*vtadminpb.Tablet{
				{
					Cluster: &vtadminpb.Cluster{
						Id:   "c0",
						Name: "cluster0",
					},
					State: vtadminpb.Tablet_SERVING,
					Tablet: &topodatapb.Tablet{
						Alias: &topodatapb.TabletAlias{
							Uid:  100,
							Cell: "zone1",
						},
						Hostname: "ks1-00-00-zone1-a",
						Keyspace: "ks1",
						Shard:    "-",
						Type:     topodatapb.TabletType_MASTER,
					},
				},
			},
			shouldErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			clusters := make([]*cluster.Cluster, len(tt.clusterTablets))
			for i, tablets := range tt.clusterTablets {
				tablets := tablets // avoid loop shadowing in the dialer closure below

				disco := fakediscovery.New()
				disco.AddTaggedGates(nil, &vtadminpb.VTGate{Hostname: fmt.Sprintf("cluster%d-gate", i)})

				cluster := &cluster.Cluster{
					ID:        fmt.Sprintf("c%d", i),
					Name:      fmt.Sprintf("cluster%d", i),
					Discovery: disco,
				}

				vtsqlCfg, err := vtsql.Parse(cluster.ID, cluster.Name, disco, []string{})
				require.NoError(t, err)

				dbconfig, ok := tt.dbconfigs[cluster.ID]
				if !ok {
					dbconfig = &dbcfg{shouldErr: false}
				}

				db := vtsql.New(cluster.ID, vtsqlCfg)
				db.DialFunc = func(cfg vitessdriver.Configuration) (*sql.DB, error) {
					return sql.OpenDB(&fakevtsql.Connector{Tablets: tablets, ShouldErr: dbconfig.shouldErr}), nil
				}

				cluster.DB = db

				clusters[i] = cluster
			}

			api := NewAPI(clusters, grpcserver.Options{}, http.Options{})
			resp, err := api.GetTablets(context.Background(), tt.req)
			if tt.shouldErr {
				assert.Error(t, err)
				return
			}

			assert.NoError(t, err)
			assert.ElementsMatch(t, tt.expected, resp.Tablets)
		})
	}
}

// This test only validates the error handling on dialing database connections.
// Other cases are covered by one or both of TestGetTablets and TestGetTablet.
func Test_getTablets(t *testing.T) {
	api := &API{}
	disco := fakediscovery.New()
	disco.AddTaggedGates(nil, &vtadminpb.VTGate{Hostname: "gate"})

	dbcfg, err := vtsql.Parse("1", "one", disco, []string{})
	require.NoError(t, err)

	db := vtsql.New("one", dbcfg)
	db.DialFunc = func(cfg vitessdriver.Configuration) (*sql.DB, error) {
		return nil, assert.AnError
	}

	_, err = api.getTablets(context.Background(), &cluster.Cluster{
		DB: db,
	})
	assert.Error(t, err)
}
<commit_msg>Add GetTablet tests<commit_after>package vtadmin

import (
	"context"
	"database/sql"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/vt/vitessdriver"
	"vitess.io/vitess/go/vt/vtadmin/cluster"
	"vitess.io/vitess/go/vt/vtadmin/cluster/discovery/fakediscovery"
	"vitess.io/vitess/go/vt/vtadmin/grpcserver"
	"vitess.io/vitess/go/vt/vtadmin/http"
	"vitess.io/vitess/go/vt/vtadmin/vtsql"
	"vitess.io/vitess/go/vt/vtadmin/vtsql/fakevtsql"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	vtadminpb "vitess.io/vitess/go/vt/proto/vtadmin"
)

func TestGetGates(t *testing.T) {
	fakedisco1 := fakediscovery.New()
	cluster1 := &cluster.Cluster{
		ID:        "c1",
		Name:      "cluster1",
		Discovery: fakedisco1,
	}
	cluster1Gates := []*vtadminpb.VTGate{
		{
			Hostname: "cluster1-gate1",
		},
		{
			Hostname: "cluster1-gate2",
		},
		{
			Hostname: "cluster1-gate3",
		},
	}

	fakedisco1.AddTaggedGates(nil, cluster1Gates...)

	fakedisco2 := fakediscovery.New()
	cluster2 := &cluster.Cluster{
		ID:        "c2",
		Name:      "cluster2",
		Discovery: fakedisco2,
	}
	cluster2Gates := []*vtadminpb.VTGate{
		{
			Hostname: "cluster2-gate1",
		},
	}

	fakedisco2.AddTaggedGates(nil, cluster2Gates...)

	api := NewAPI([]*cluster.Cluster{cluster1, cluster2}, grpcserver.Options{}, http.Options{})
	ctx := context.Background()

	resp, err := api.GetGates(ctx, &vtadminpb.GetGatesRequest{})
	assert.NoError(t, err)
	assert.ElementsMatch(t, append(cluster1Gates, cluster2Gates...), resp.Gates)

	resp, err = api.GetGates(ctx, &vtadminpb.GetGatesRequest{ClusterIds: []string{cluster1.ID}})
	assert.NoError(t, err)
	assert.ElementsMatch(t, cluster1Gates, resp.Gates)

	fakedisco1.SetGatesError(true)

	resp, err = api.GetGates(ctx, &vtadminpb.GetGatesRequest{})
	assert.Error(t, err)
	assert.Nil(t, resp)
}

func TestGetTablets(t *testing.T) {
	type dbcfg struct {
		shouldErr bool
	}

	tests := []struct {
		name           string
		clusterTablets [][]*vtadminpb.Tablet
		dbconfigs      map[string]*dbcfg
		req            *vtadminpb.GetTabletsRequest
		expected       []*vtadminpb.Tablet
		shouldErr      bool
	}{
		{
			name: "single cluster",
			clusterTablets: [][]*vtadminpb.Tablet{
				{
					/* cluster 0 */
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  100,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
			},
			dbconfigs: map[string]*dbcfg{},
			req:       &vtadminpb.GetTabletsRequest{},
			expected: []*vtadminpb.Tablet{
				{
					Cluster: &vtadminpb.Cluster{
						Id:   "c0",
						Name: "cluster0",
					},
					State: vtadminpb.Tablet_SERVING,
					Tablet: &topodatapb.Tablet{
						Alias: &topodatapb.TabletAlias{
							Uid:  100,
							Cell: "zone1",
						},
						Hostname: "ks1-00-00-zone1-a",
						Keyspace: "ks1",
						Shard:    "-",
						Type:     topodatapb.TabletType_MASTER,
					},
				},
			},
			shouldErr: false,
		},
		{
			name: "one cluster errors",
			clusterTablets: [][]*vtadminpb.Tablet{
				/* cluster 0 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  100,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
				/* cluster 1 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  200,
								Cell: "zone1",
							},
							Hostname: "ks2-00-00-zone1-a",
							Keyspace: "ks2",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
			},
			dbconfigs: map[string]*dbcfg{
				"c1": {shouldErr: true},
			},
			req:       &vtadminpb.GetTabletsRequest{},
			expected:  nil,
			shouldErr: true,
		},
		{
			name: "multi cluster, selecting one",
			clusterTablets: [][]*vtadminpb.Tablet{
				/* cluster 0 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  100,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
				/* cluster 1 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  200,
								Cell: "zone1",
							},
							Hostname: "ks2-00-00-zone1-a",
							Keyspace: "ks2",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
			},
			dbconfigs: map[string]*dbcfg{},
			req:       &vtadminpb.GetTabletsRequest{ClusterIds: []string{"c0"}},
			expected: []*vtadminpb.Tablet{
				{
					Cluster: &vtadminpb.Cluster{
						Id:   "c0",
						Name: "cluster0",
					},
					State: vtadminpb.Tablet_SERVING,
					Tablet: &topodatapb.Tablet{
						Alias: &topodatapb.TabletAlias{
							Uid:  100,
							Cell: "zone1",
						},
						Hostname: "ks1-00-00-zone1-a",
						Keyspace: "ks1",
						Shard:    "-",
						Type:     topodatapb.TabletType_MASTER,
					},
				},
			},
			shouldErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			clusters := make([]*cluster.Cluster, len(tt.clusterTablets))

			for i, tablets := range tt.clusterTablets { // nolint:dupl
				tablets := tablets // avoid loop shadowing in the dialer closure below

				disco := fakediscovery.New()
				disco.AddTaggedGates(nil, &vtadminpb.VTGate{Hostname: fmt.Sprintf("cluster%d-gate", i)})

				cluster := &cluster.Cluster{
					ID:        fmt.Sprintf("c%d", i),
					Name:      fmt.Sprintf("cluster%d", i),
					Discovery: disco,
				}

				vtsqlCfg, err := vtsql.Parse(cluster.ID, cluster.Name, disco, []string{})
				require.NoError(t, err)

				dbconfig, ok := tt.dbconfigs[cluster.ID]
				if !ok {
					dbconfig = &dbcfg{shouldErr: false}
				}

				db := vtsql.New(cluster.ID, vtsqlCfg)
				db.DialFunc = func(cfg vitessdriver.Configuration) (*sql.DB, error) {
					return sql.OpenDB(&fakevtsql.Connector{Tablets: tablets, ShouldErr: dbconfig.shouldErr}), nil
				}

				cluster.DB = db

				clusters[i] = cluster
			}

			api := NewAPI(clusters, grpcserver.Options{}, http.Options{})
			resp, err := api.GetTablets(context.Background(), tt.req)
			if tt.shouldErr {
				assert.Error(t, err)
				return
			}

			assert.NoError(t, err)
			assert.ElementsMatch(t, tt.expected, resp.Tablets)
		})
	}
}

// This test only validates the error handling on dialing database connections.
// Other cases are covered by one or both of TestGetTablets and TestGetTablet.
func Test_getTablets(t *testing.T) {
	api := &API{}
	disco := fakediscovery.New()
	disco.AddTaggedGates(nil, &vtadminpb.VTGate{Hostname: "gate"})

	dbcfg, err := vtsql.Parse("1", "one", disco, []string{})
	require.NoError(t, err)

	db := vtsql.New("one", dbcfg)
	db.DialFunc = func(cfg vitessdriver.Configuration) (*sql.DB, error) {
		return nil, assert.AnError
	}

	_, err = api.getTablets(context.Background(), &cluster.Cluster{
		DB: db,
	})
	assert.Error(t, err)
}

func TestGetTabet(t *testing.T) {
	type dbcfg struct {
		shouldErr bool
	}

	tests := []struct {
		name           string
		clusterTablets [][]*vtadminpb.Tablet
		dbconfigs      map[string]*dbcfg
		req            *vtadminpb.GetTabletRequest
		expected       *vtadminpb.Tablet
		shouldErr      bool
	}{
		{
			name: "single cluster",
			clusterTablets: [][]*vtadminpb.Tablet{
				{
					/* cluster 0 */
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  100,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
			},
			dbconfigs: map[string]*dbcfg{},
			req: &vtadminpb.GetTabletRequest{
				Hostname: "ks1-00-00-zone1-a",
			},
			expected: &vtadminpb.Tablet{
				Cluster: &vtadminpb.Cluster{
					Id:   "c0",
					Name: "cluster0",
				},
				State: vtadminpb.Tablet_SERVING,
				Tablet: &topodatapb.Tablet{
					Alias: &topodatapb.TabletAlias{
						Uid:  100,
						Cell: "zone1",
					},
					Hostname: "ks1-00-00-zone1-a",
					Keyspace: "ks1",
					Shard:    "-",
					Type:     topodatapb.TabletType_MASTER,
				},
			},
			shouldErr: false,
		},
		{
			name: "one cluster errors",
			clusterTablets: [][]*vtadminpb.Tablet{
				/* cluster 0 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  100,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
				/* cluster 1 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  200,
								Cell: "zone1",
							},
							Hostname: "ks2-00-00-zone1-a",
							Keyspace: "ks2",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
			},
			dbconfigs: map[string]*dbcfg{
				"c1": {shouldErr: true},
			},
			req: &vtadminpb.GetTabletRequest{
				Hostname: "doesn't matter",
			},
			expected:  nil,
			shouldErr: true,
		},
		{
			name: "multi cluster, selecting one with tablet",
			clusterTablets: [][]*vtadminpb.Tablet{
				/* cluster 0 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  100,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
				/* cluster 1 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  200,
								Cell: "zone1",
							},
							Hostname: "ks2-00-00-zone1-a",
							Keyspace: "ks2",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
			},
			dbconfigs: map[string]*dbcfg{},
			req: &vtadminpb.GetTabletRequest{
				Hostname:   "ks1-00-00-zone1-a",
				ClusterIds: []string{"c0"},
			},
			expected: &vtadminpb.Tablet{
				Cluster: &vtadminpb.Cluster{
					Id:   "c0",
					Name: "cluster0",
				},
				State: vtadminpb.Tablet_SERVING,
				Tablet: &topodatapb.Tablet{
					Alias: &topodatapb.TabletAlias{
						Uid:  100,
						Cell: "zone1",
					},
					Hostname: "ks1-00-00-zone1-a",
					Keyspace: "ks1",
					Shard:    "-",
					Type:     topodatapb.TabletType_MASTER,
				},
			},
			shouldErr: false,
		},
		{
			name: "multi cluster, multiple results",
			clusterTablets: [][]*vtadminpb.Tablet{
				/* cluster 0 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  100,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
				/* cluster 1 */
				{
					{
						State: vtadminpb.Tablet_SERVING,
						Tablet: &topodatapb.Tablet{
							Alias: &topodatapb.TabletAlias{
								Uid:  200,
								Cell: "zone1",
							},
							Hostname: "ks1-00-00-zone1-a",
							Keyspace: "ks1",
							Shard:    "-",
							Type:     topodatapb.TabletType_MASTER,
						},
					},
				},
			},
			dbconfigs: map[string]*dbcfg{},
			req: &vtadminpb.GetTabletRequest{
				Hostname: "ks1-00-00-zone1-a",
			},
			expected:  nil,
			shouldErr: true,
		},
		{
			name: "no results",
			clusterTablets: [][]*vtadminpb.Tablet{
				/* cluster 0 */
				{},
			},
			dbconfigs: map[string]*dbcfg{},
			req: &vtadminpb.GetTabletRequest{
				Hostname: "ks1-00-00-zone1-a",
			},
			expected:  nil,
			shouldErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			clusters := make([]*cluster.Cluster, len(tt.clusterTablets))

			for i, tablets := range tt.clusterTablets { // nolint:dupl
				tablets := tablets // avoid loop shadowing in the dialer closure below

				disco := fakediscovery.New()
				disco.AddTaggedGates(nil, &vtadminpb.VTGate{Hostname: fmt.Sprintf("cluster%d-gate", i)})

				cluster := &cluster.Cluster{
					ID:        fmt.Sprintf("c%d", i),
					Name:      fmt.Sprintf("cluster%d", i),
					Discovery: disco,
				}

				vtsqlCfg, err := vtsql.Parse(cluster.ID, cluster.Name, disco, []string{})
				require.NoError(t, err)

				dbconfig, ok := tt.dbconfigs[cluster.ID]
				if !ok {
					dbconfig = &dbcfg{shouldErr: false}
				}

				db := vtsql.New(cluster.ID, vtsqlCfg)
				db.DialFunc = func(cfg vitessdriver.Configuration) (*sql.DB, error) {
					return sql.OpenDB(&fakevtsql.Connector{Tablets: tablets, ShouldErr: dbconfig.shouldErr}), nil
				}

				cluster.DB = db

				clusters[i] = cluster
			}

			api := NewAPI(clusters, grpcserver.Options{}, http.Options{})
			resp, err := api.GetTablet(context.Background(), tt.req)
			if tt.shouldErr {
				assert.Error(t, err)
				return
			}

			assert.NoError(t, err)
			assert.Equal(t, tt.expected, resp)
		})
	}
}
<|endoftext|>


<commit_before>package vegeta

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
)

// Target is an HTTP request blueprint.
type Target struct {
	Method string
	URL    string
	Body   []byte
	Header http.Header
}

// Request creates an *http.Request out of Target and returns it along with an
// error in case of failure.
func (t *Target) Request() (*http.Request, error) {
	req, err := http.NewRequest(t.Method, t.URL, bytes.NewReader(t.Body))
	if err != nil {
		return nil, err
	}
	for k, vs := range t.Header {
		req.Header[k] = make([]string, len(vs))
		copy(req.Header[k], vs)
	}
	if host := req.Header.Get("Host"); host != "" {
		req.Host = host
	}
	return req, nil
}

var (
	// ErrNoTargets is returned when not enough Targets are available.
	ErrNoTargets = errors.New("no targets to attack")
	// ErrNilTarget is returned when the passed Target pointer is nil.
	ErrNilTarget = errors.New("nil target")
)

// A Targeter decodes a Target or returns an error in case of failure.
// Implementations must be safe for concurrent use.
type Targeter func(*Target) error

// NewStaticTargeter returns a Targeter which round-robins over the passed
// Targets.
func NewStaticTargeter(tgts ...Target) Targeter {
	i := int64(-1)
	return func(tgt *Target) error {
		if tgt == nil {
			return ErrNilTarget
		}
		*tgt = tgts[atomic.AddInt64(&i, 1)%int64(len(tgts))]
		return nil
	}
}

// NewEagerTargeter eagerly reads all Targets out of the provided io.Reader and
// returns a NewStaticTargeter with them.
//
// body will be set as the Target's body if no body is provided.
// hdr will be merged with the each Target's headers.
func NewEagerTargeter(src io.Reader, body []byte, header http.Header) (Targeter, error) {
	var (
		sc   = NewLazyTargeter(src, body, header)
		tgts []Target
		tgt  Target
		err  error
	)
	for {
		if err = sc(&tgt); err == ErrNoTargets {
			break
		} else if err != nil {
			return nil, err
		}
		tgts = append(tgts, tgt)
	}
	if len(tgts) == 0 {
		return nil, ErrNoTargets
	}
	return NewStaticTargeter(tgts...), nil
}

// NewLazyTargeter returns a new Targeter that lazily scans Targets from the
// provided io.Reader on every invocation.
//
// body will be set as the Target's body if no body is provided.
// hdr will be merged with the each Target's headers.
func NewLazyTargeter(src io.Reader, body []byte, hdr http.Header) Targeter {
	var mu sync.Mutex
	sc := peekingScanner{src: bufio.NewScanner(src)}
	return func(tgt *Target) (err error) {
		mu.Lock()
		defer mu.Unlock()

		if tgt == nil {
			return ErrNilTarget
		}

		var line string
		for {
			if !sc.Scan() {
				return ErrNoTargets
			}
			line = strings.TrimSpace(sc.Text())
			if len(line) != 0 {
				break
			}
		}

		tgt.Body = body
		tgt.Header = http.Header{}
		for k, vs := range hdr {
			tgt.Header[k] = vs
		}

		tokens := strings.SplitN(line, " ", 2)
		if len(tokens) < 2 {
			return fmt.Errorf("bad target: %s", line)
		}
		switch tokens[0] {
		case "HEAD", "GET", "PUT", "POST", "PATCH", "OPTIONS", "DELETE":
			tgt.Method = tokens[0]
		default:
			return fmt.Errorf("bad method: %s", tokens[0])
		}
		if _, err = url.ParseRequestURI(tokens[1]); err != nil {
			return fmt.Errorf("bad URL: %s", tokens[1])
		}
		tgt.URL = tokens[1]
		line = strings.TrimSpace(sc.Peek())
		if line == "" || startsWithHTTPMethod(line) {
			return nil
		}
		for sc.Scan() {
			if line = strings.TrimSpace(sc.Text()); line == "" {
				break
			} else if strings.HasPrefix(line, "@") {
				if tgt.Body, err = ioutil.ReadFile(line[1:]); err != nil {
					return fmt.Errorf("bad body: %s", err)
				}
				break
			}
			tokens = strings.SplitN(line, ":", 2)
			if len(tokens) < 2 {
				return fmt.Errorf("bad header: %s", line)
			}
			for i := range tokens {
				if tokens[i] = strings.TrimSpace(tokens[i]); tokens[i] == "" {
					return fmt.Errorf("bad header: %s", line)
				}
			}
			tgt.Header.Add(tokens[0], tokens[1])
		}
		if err = sc.Err(); err != nil {
			return ErrNoTargets
		}
		return nil
	}
}

var httpMethodChecker = regexp.MustCompile("^(HEAD|GET|PUT|POST|PATCH|OPTIONS|DELETE) ")

func startsWithHTTPMethod(t string) bool {
	return httpMethodChecker.MatchString(t)
}

// Wrap a Scanner so we can cheat and look at the next value and react accordingly,
// but still have it be around the next time we Scan() + Text()
type peekingScanner struct {
	src    *bufio.Scanner
	peeked string
}

func (s *peekingScanner) Err() error {
	return s.src.Err()
}

func (s *peekingScanner) Peek() string {
	if !s.src.Scan() {
		return ""
	}
	s.peeked = s.src.Text()
	return s.peeked
}

func (s *peekingScanner) Scan() bool {
	if s.peeked == "" {
		return s.src.Scan()
	}
	return true
}

func (s *peekingScanner) Text() string {
	if s.peeked == "" {
		return s.src.Text()
	}
	t := s.peeked
	s.peeked = ""
	return t
}
<commit_msg>Made headers in target definitions retain case sensitivity<commit_after>package vegeta

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
)

// Target is an HTTP request blueprint.
type Target struct {
	Method string
	URL    string
	Body   []byte
	Header http.Header
}

// Request creates an *http.Request out of Target and returns it along with an
// error in case of failure.
func (t *Target) Request() (*http.Request, error) {
	req, err := http.NewRequest(t.Method, t.URL, bytes.NewReader(t.Body))
	if err != nil {
		return nil, err
	}
	for k, vs := range t.Header {
		req.Header[k] = make([]string, len(vs))
		copy(req.Header[k], vs)
	}
	if host := req.Header.Get("Host"); host != "" {
		req.Host = host
	}
	return req, nil
}

var (
	// ErrNoTargets is returned when not enough Targets are available.
	ErrNoTargets = errors.New("no targets to attack")
	// ErrNilTarget is returned when the passed Target pointer is nil.
	ErrNilTarget = errors.New("nil target")
)

// A Targeter decodes a Target or returns an error in case of failure.
// Implementations must be safe for concurrent use.
type Targeter func(*Target) error

// NewStaticTargeter returns a Targeter which round-robins over the passed
// Targets.
func NewStaticTargeter(tgts ...Target) Targeter {
	i := int64(-1)
	return func(tgt *Target) error {
		if tgt == nil {
			return ErrNilTarget
		}
		*tgt = tgts[atomic.AddInt64(&i, 1)%int64(len(tgts))]
		return nil
	}
}

// NewEagerTargeter eagerly reads all Targets out of the provided io.Reader and
// returns a NewStaticTargeter with them.
//
// body will be set as the Target's body if no body is provided.
// hdr will be merged with the each Target's headers.
func NewEagerTargeter(src io.Reader, body []byte, header http.Header) (Targeter, error) {
	var (
		sc   = NewLazyTargeter(src, body, header)
		tgts []Target
		tgt  Target
		err  error
	)
	for {
		if err = sc(&tgt); err == ErrNoTargets {
			break
		} else if err != nil {
			return nil, err
		}
		tgts = append(tgts, tgt)
	}
	if len(tgts) == 0 {
		return nil, ErrNoTargets
	}
	return NewStaticTargeter(tgts...), nil
}

// NewLazyTargeter returns a new Targeter that lazily scans Targets from the
// provided io.Reader on every invocation.
//
// body will be set as the Target's body if no body is provided.
// hdr will be merged with the each Target's headers.
func NewLazyTargeter(src io.Reader, body []byte, hdr http.Header) Targeter {
	var mu sync.Mutex
	sc := peekingScanner{src: bufio.NewScanner(src)}
	return func(tgt *Target) (err error) {
		mu.Lock()
		defer mu.Unlock()

		if tgt == nil {
			return ErrNilTarget
		}

		var line string
		for {
			if !sc.Scan() {
				return ErrNoTargets
			}
			line = strings.TrimSpace(sc.Text())
			if len(line) != 0 {
				break
			}
		}

		tgt.Body = body
		tgt.Header = http.Header{}
		for k, vs := range hdr {
			tgt.Header[k] = vs
		}

		tokens := strings.SplitN(line, " ", 2)
		if len(tokens) < 2 {
			return fmt.Errorf("bad target: %s", line)
		}
		switch tokens[0] {
		case "HEAD", "GET", "PUT", "POST", "PATCH", "OPTIONS", "DELETE":
			tgt.Method = tokens[0]
		default:
			return fmt.Errorf("bad method: %s", tokens[0])
		}
		if _, err = url.ParseRequestURI(tokens[1]); err != nil {
			return fmt.Errorf("bad URL: %s", tokens[1])
		}
		tgt.URL = tokens[1]
		line = strings.TrimSpace(sc.Peek())
		if line == "" || startsWithHTTPMethod(line) {
			return nil
		}
		for sc.Scan() {
			if line = strings.TrimSpace(sc.Text()); line == "" {
				break
			} else if strings.HasPrefix(line, "@") {
				if tgt.Body, err = ioutil.ReadFile(line[1:]); err != nil {
					return fmt.Errorf("bad body: %s", err)
				}
				break
			}
			tokens = strings.SplitN(line, ":", 2)
			if len(tokens) < 2 {
				return fmt.Errorf("bad header: %s", line)
			}
			for i := range tokens {
				if tokens[i] = strings.TrimSpace(tokens[i]); tokens[i] == "" {
					return fmt.Errorf("bad header: %s", line)
				}
			}
			// Add key/value directly to the http.Header (map[string][]string).
			// http.Header.Add() canonicalizes keys but vegeta is used
			// to test systems that require case-sensitive headers.
			tgt.Header[tokens[0]] = append(tgt.Header[tokens[0]], tokens[1])
		}
		if err = sc.Err(); err != nil {
			return ErrNoTargets
		}
		return nil
	}
}

var httpMethodChecker = regexp.MustCompile("^(HEAD|GET|PUT|POST|PATCH|OPTIONS|DELETE) ")

func startsWithHTTPMethod(t string) bool {
	return httpMethodChecker.MatchString(t)
}

// Wrap a Scanner so we can cheat and look at the next value and react accordingly,
// but still have it be around the next time we Scan() + Text()
type peekingScanner struct {
	src    *bufio.Scanner
	peeked string
}

func (s *peekingScanner) Err() error {
	return s.src.Err()
}

func (s *peekingScanner) Peek() string {
	if !s.src.Scan() {
		return ""
	}
	s.peeked = s.src.Text()
	return s.peeked
}

func (s *peekingScanner) Scan() bool {
	if s.peeked == "" {
		return s.src.Scan()
	}
	return true
}

func (s *peekingScanner) Text() string {
	if s.peeked == "" {
		return s.src.Text()
	}
	t := s.peeked
	s.peeked = ""
	return t
}
<|endoftext|>


<commit_before>package golang

import (
	"github.com/mdempsky/gocode/suggest"
	"go/build"
	"go/types"
	"margo.sh/mg"
	"runtime/debug"
	"strings"
	"sync"
	"time"
)

type gsuOpts struct {
	ProposeBuiltins bool
	Debug           bool
}

type gsuImpRes struct {
	pkg *types.Package
	err error
}

type gcSuggest struct {
	gsuOpts
	sync.Mutex
	imp *gsuImporter
}

func newGcSuggest(mx *mg.Ctx, o gsuOpts) *gcSuggest {
	gsu := &gcSuggest{gsuOpts: o}
	gsu.imp = gsu.newGsuImporter(mx)
	return gsu
}

func (gsu *gcSuggest) newGsuImporter(mx *mg.Ctx) *gsuImporter {
	gi := &gsuImporter{
		mx:  mx,
		bld: BuildContext(mx),
		gsu: gsu,
		res: map[mgcCacheKey]gsuImpRes{},
	}
	return gi
}

func (gsu *gcSuggest) candidates(mx *mg.Ctx) []suggest.Candidate {
	defer mx.Profile.Push("candidates").Pop()
	gsu.Lock()
	defer gsu.Unlock()

	defer func() {
		if e := recover(); e != nil {
			mx.Log.Printf("gocode/suggest panic: %s\n%s\n", e, debug.Stack())
		}
	}()

	cfg := suggest.Config{
		// we no longer support contextual build env :(
		// GoSublime works around this for other packages by restarting the agent
		// if GOPATH changes, so we should be ok
		Importer:   gsu.imp,
		Builtin:    gsu.ProposeBuiltins,
		IgnoreCase: true,
	}
	if gsu.Debug {
		cfg.Logf = func(f string, a ...interface{}) {
			f = "Gocode: " + f
			if !strings.HasSuffix(f, "\n") {
				f += "\n"
			}
			mx.Log.Dbg.Printf(f, a...)
		}
	}

	v := mx.View
	src, _ := v.ReadAll()
	if len(src) == 0 {
		return nil
	}

	l, _ := cfg.Suggest(v.Filename(), src, v.Pos)
	return l
}

type gsuPkgInfo struct {
	// the import path
	Path string

	// the abs path to the package directory
	Dir string

	// whether or not this is a stdlib package
	Std bool
}

func (p gsuPkgInfo) cacheKey(source bool) mgcCacheKey {
	return mgcCacheKey{gsuPkgInfo: p, Source: source}
}

type gsuImporter struct {
	mx  *mg.Ctx
	bld *build.Context
	gsu *gcSuggest
	res map[mgcCacheKey]gsuImpRes
}

func (gi *gsuImporter) Import(path string) (*types.Package, error) {
	return gi.ImportFrom(path, ".", 0)
}

func (gi *gsuImporter) ImportFrom(impPath, srcDir string, mode types.ImportMode) (pkg *types.Package, err error) {
	// TODO: add mode to the key somehow?
	// mode is reserved, but currently not used so it's not a problem
	// but if it's used in the future, the importer result could depend on it
	//
	// adding it to the key might complicate the pkginfo api because it's called
	// by code that doesn't know anything about mode
	pkgInf, err := mctl.pkgInfo(gi.mx, impPath, srcDir)
	if err != nil {
		mctl.dbgf("pkgInfo(%q, %q): %s\n", impPath, srcDir, err)
		return nil, err
	}
	newDefImpr, newFbkImpr, srcMode := mctl.importerFactories()
	k := pkgInf.cacheKey(srcMode)

	// we cache the results of the underlying importer for this *session*
	// because if it fails, we could potentialy end up in a loop
	// trying to import the package again.
	if res, ok := gi.res[k]; ok {
		return res.pkg, res.err
	}
	defer func() { gi.res[k] = gsuImpRes{pkg: pkg, err: err} }()

	defImpr := newDefImpr(gi.mx, gi)
	pkg, err = gi.importFrom(defImpr, k, mode)
	complete := err == nil && pkg.Complete()
	if complete {
		return pkg, nil
	}

	mctl.dbgf("importFrom(%q, %q): default=%T: complete=%v, err=%v\n",
		k.Path, k.Dir, defImpr, complete, err,
	)

	// no fallback allowed
	if newFbkImpr == nil {
		return pkg, err
	}

	// problem1:
	// if the pkg import fails we will offer no completion
	//
	// problem 2:
	// if it succeeds, but is incomplete we offer completion with `invalid-type` failures
	// i.e. completion stops working at random points for no obvious reason
	//
	// assumption:
	//   it's better to risk using stale data (bin imports)
	//   as opposed to offering no completion at all
	//
	// risks:
	// we will end up caching the result, but that shouldn't be a big deal
	// because if the pkg is edited, thus (possibly) making it importable,
	// we will remove it from the cache anyway.
	// there is the issue about mixing binary (potentially incomplete) pkgs with src pkgs
	// but we were already not going to return anything, so it *shouldn't* apply here

	fbkImpr := newFbkImpr(gi.mx, gi)
	fbkPkg, fbkErr := gi.importFrom(fbkImpr, k.fallback(), mode)
	fbkComplete := fbkErr == nil && fbkPkg.Complete()
	switch {
	case fbkComplete:
		pkg, err = fbkPkg, nil
	case fbkPkg != nil && pkg == nil:
		pkg, err = fbkPkg, fbkErr
	}

	mctl.dbgf("importFrom(%q, %q): fallback=%T: complete=%v, err=%v\n",
		k.Path, k.Dir, fbkImpr, fbkComplete, fbkErr,
	)

	return pkg, err
}

func (gi *gsuImporter) importFrom(underlying types.ImporterFrom, k mgcCacheKey, mode types.ImportMode) (*types.Package, error) {
	defer gi.mx.Profile.Push("gsuImport: " + k.Path).Pop()

	if k.Std && k.Path == "unsafe" {
		return types.Unsafe, nil
	}

	if e, ok := mctl.pkgs.get(k); ok {
		return e.Pkg, nil
	}

	impStart := time.Now()
	pkg, err := underlying.ImportFrom(k.Path, k.Dir, mode)
	impDur := time.Since(impStart)

	if err == nil {
		mctl.pkgs.put(mgcCacheEnt{Key: k, Pkg: pkg, Dur: impDur})
	} else {
		mctl.dbgf("%T.ImportFrom(%q, %q): %s\n", underlying, k.Path, k.Dir, err)
	}

	return pkg, err
}
<commit_msg>guard again import cycles in gsuImporter<commit_after>package golang

import (
	"errors"
	"github.com/mdempsky/gocode/suggest"
	"go/build"
	"go/types"
	"margo.sh/mg"
	"runtime/debug"
	"strings"
	"sync"
	"time"
)

var (
	errImportCycleDetected = errors.New("import cycle detected")
)

type gsuOpts struct {
	ProposeBuiltins bool
	Debug           bool
}

type gsuImpRes struct {
	pkg *types.Package
	err error
}

type gcSuggest struct {
	gsuOpts
	sync.Mutex
	imp *gsuImporter
}

func newGcSuggest(mx *mg.Ctx, o gsuOpts) *gcSuggest {
	gsu := &gcSuggest{gsuOpts: o}
	gsu.imp = gsu.newGsuImporter(mx)
	return gsu
}

func (gsu *gcSuggest) newGsuImporter(mx *mg.Ctx) *gsuImporter {
	gi := &gsuImporter{
		mx:  mx,
		bld: BuildContext(mx),
		gsu: gsu,
	}
	gi.res.m = map[mgcCacheKey]gsuImpRes{}
	return gi
}

func (gsu *gcSuggest) candidates(mx *mg.Ctx) []suggest.Candidate {
	defer mx.Profile.Push("candidates").Pop()
	gsu.Lock()
	defer gsu.Unlock()

	defer func() {
		if e := recover(); e != nil {
			mx.Log.Printf("gocode/suggest panic: %s\n%s\n", e, debug.Stack())
		}
	}()

	cfg := suggest.Config{
		// we no longer support contextual build env :(
		// GoSublime works around this for other packages by restarting the agent
		// if GOPATH changes, so we should be ok
		Importer:   gsu.imp,
		Builtin:    gsu.ProposeBuiltins,
		IgnoreCase: true,
	}
	if gsu.Debug {
		cfg.Logf = func(f string, a ...interface{}) {
			f = "Gocode: " + f
			if !strings.HasSuffix(f, "\n") {
				f += "\n"
			}
			mx.Log.Dbg.Printf(f, a...)
		}
	}

	v := mx.View
	src, _ := v.ReadAll()
	if len(src) == 0 {
		return nil
	}

	l, _ := cfg.Suggest(v.Filename(), src, v.Pos)
	return l
}

type gsuPkgInfo struct {
	// the import path
	Path string

	// the abs path to the package directory
	Dir string

	// whether or not this is a stdlib package
	Std bool
}

func (p gsuPkgInfo) cacheKey(source bool) mgcCacheKey {
	return mgcCacheKey{gsuPkgInfo: p, Source: source}
}

type gsuImporter struct {
	mx  *mg.Ctx
	bld *build.Context
	gsu *gcSuggest

	res struct {
		sync.Mutex
		m map[mgcCacheKey]gsuImpRes
	}
}

func (gi *gsuImporter) Import(path string) (*types.Package, error) {
	return gi.ImportFrom(path, ".", 0)
}

func (gi *gsuImporter) ImportFrom(impPath, srcDir string, mode types.ImportMode) (pkg *types.Package, err error) {
	// TODO: add mode to the key somehow?
	// mode is reserved, but currently not used so it's not a problem
	// but if it's used in the future, the importer result could depend on it
	//
	// adding it to the key might complicate the pkginfo api because it's called
	// by code that doesn't know anything about mode
	pkgInf, err := mctl.pkgInfo(gi.mx, impPath, srcDir)
	if err != nil {
		mctl.dbgf("pkgInfo(%q, %q): %s\n", impPath, srcDir, err)
		return nil, err
	}
	newDefImpr, newFbkImpr, srcMode := mctl.importerFactories()
	k := pkgInf.cacheKey(srcMode)

	gi.res.Lock()
	res, seen := gi.res.m[k]
	if !seen {
		gi.res.m[k] = gsuImpRes{err: errImportCycleDetected}
	}
	gi.res.Unlock()

	// we cache the results of the underlying importer for this *session*
	// because if it fails, or there's an import cycle, we could potentialy end up in a loop
	// trying to import the package again.
	if seen {
		return res.pkg, res.err
	}
	defer func() {
		gi.res.Lock()
		defer gi.res.Unlock()

		gi.res.m[k] = gsuImpRes{pkg: pkg, err: err}
	}()

	defImpr := newDefImpr(gi.mx, gi)
	pkg, err = gi.importFrom(defImpr, k, mode)
	complete := err == nil && pkg.Complete()
	if complete {
		return pkg, nil
	}

	mctl.dbgf("importFrom(%q, %q): default=%T: complete=%v, err=%v\n",
		k.Path, k.Dir, defImpr, complete, err,
	)

	// no fallback allowed
	if newFbkImpr == nil {
		return pkg, err
	}

	// problem1:
	// if the pkg import fails we will offer no completion
	//
	// problem 2:
	// if it succeeds, but is incomplete we offer completion with `invalid-type` failures
	// i.e. completion stops working at random points for no obvious reason
	//
	// assumption:
	//   it's better to risk using stale data (bin imports)
	//   as opposed to offering no completion at all
	//
	// risks:
	// we will end up caching the result, but that shouldn't be a big deal
	// because if the pkg is edited, thus (possibly) making it importable,
	// we will remove it from the cache anyway.
	// there is the issue about mixing binary (potentially incomplete) pkgs with src pkgs
	// but we were already not going to return anything, so it *shouldn't* apply here

	fbkImpr := newFbkImpr(gi.mx, gi)
	fbkPkg, fbkErr := gi.importFrom(fbkImpr, k.fallback(), mode)
	fbkComplete := fbkErr == nil && fbkPkg.Complete()
	switch {
	case fbkComplete:
		pkg, err = fbkPkg, nil
	case fbkPkg != nil && pkg == nil:
		pkg, err = fbkPkg, fbkErr
	}

	mctl.dbgf("importFrom(%q, %q): fallback=%T: complete=%v, err=%v\n",
		k.Path, k.Dir, fbkImpr, fbkComplete, fbkErr,
	)

	return pkg, err
}

func (gi *gsuImporter) importFrom(underlying types.ImporterFrom, k mgcCacheKey, mode types.ImportMode) (*types.Package, error) {
	defer gi.mx.Profile.Push("gsuImport: " + k.Path).Pop()

	if k.Std && k.Path == "unsafe" {
		return types.Unsafe, nil
	}

	if e, ok := mctl.pkgs.get(k); ok {
		return e.Pkg, nil
	}

	impStart := time.Now()
	pkg, err := underlying.ImportFrom(k.Path, k.Dir, mode)
	impDur := time.Since(impStart)

	if err == nil {
		mctl.pkgs.put(mgcCacheEnt{Key: k, Pkg: pkg, Dur: impDur})
	} else {
		mctl.dbgf("%T.ImportFrom(%q, %q): %s\n", underlying, k.Path, k.Dir, err)
	}

	return pkg, err
}
<|endoftext|>


<commit_before>package main

// import "testing"
<commit_msg>add some config tests<commit_after>package main

import "testing"

func TestEmptyConfigFails(t *testing.T) {
	config := Config{}
	errs := config.Errors()

	if len(errs) != 4 {
		t.Error("Expected validation to fail.")
	}
}

func TestIcingaServerIsAddressAndPort(t *testing.T) {
	config := Config{}
	errs := config.Errors()

	var found bool

	for _, e := range errs {
		if e.Error() == "Icinga server should be in format `host:port` (e.g. 127.0.0.1:5665)" {
			found = true
		}
	}

	if !found {
		t.Error("Expected Icinga server address:port to be validated.")
	}
}

func TestRedisServerIsAddressAndPort(t *testing.T) {
	config := Config{}
	errs := config.Errors()

	var found bool

	for _, e := range errs {
		if e.Error() == "Redis server should be in format `host:port` (e.g. 127.0.0.1:6380)" {
			found = true
		}
	}

	if !found {
		t.Error("Expected Redis server address:port to be validated.")
	}
}

func TestIcingaAPIUserIsSet(t *testing.T) {
	config := Config{}
	errs := config.Errors()

	var found bool

	for _, e := range errs {
		if e.Error() == "No Icinga2 API user specified in ICINGA2_API_USER env variable or --user option" {
			found = true
		}
	}

	if !found {
		t.Error("Expected Icinga API user presence to be checked.")
	}
}

func TestIcingaAPIPasswordIsSet(t *testing.T) {
	config := Config{}
	errs := config.Errors()

	var found bool

	for _, e := range errs {
		if e.Error() == "No Icinga2 API password specified in ICINGA2_API_PASSWORD env variable or --password option" {
			found = true
		}
	}

	if !found {
		t.Error("Expected Icinga API password presence to be checked.")
	}
}
<|endoftext|>


<commit_before>package graylog

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
)

type Streams struct {
	Total      int      `json:"total"`
	StreamList []Stream `json:"streams"`
}

type Stream struct {
	Id                             string `json:"id"`
	Title                          string `json:"title"`
	Description                    string `json:"description"`
	Rules                          []Rule `json:"rules"`
	ContentPack                    string `json:"content_pack"`
	MatchingType                   string `json:"matching_type"`
	RemoveMatchesFromDefaultStream bool   `json:"remove_matches_from_default_stream"`
	IndexSetId                     string `json:"index_set_id"`
}

type StreamId struct {
	StreamId string `json:"stream_id"`
}

type StreamCreate struct {
	Title                          string `json:"title"`
	Description                    string `json:"description"`
	RemoveMatchesFromDefaultStream bool   `json:"remove_matches_from_default_stream"`
	IndexSetId                     string `json:"index_set_id"`
}

type Rules struct {
	Total 		int `json:"total"`
	StreamRules []Rule `json:"stream_rules"`
}

type Rule struct {
	Type        int    `json:"type"`
	Value       string `json:"value"`
	Field       string `json:"field"`
	Inverted    bool   `json:"inverted"`
	Description string `json:"description"`
}

type IndexSets struct {
	Total     int        `json:"total"`
	IndexSets []IndexSet `json:"index_sets"`
}

type IndexSet struct {
	Id    string `json:"id"`
	Title string `json:"title"`
}

type UserUpdate struct {
	Roles []string `json:"roles"`
}

type User struct {
	Id       string   `json:"id"`
	Username string   `json:"username"`
	Roles    []string `json:"roles"`
}

type Role struct {
	Name        string   `json:"name"`
	Description string   `json:"description"`
	Permissions []string `json:"permissions"`
	ReadOnly    bool     `json:"read_only"`
}

func CreateStream(namespaceName string) (bool, string) {
	if !isGrayLogActive() {
		return false, ""
	}

	if cond, id := isStreamAlreadyCreated(namespaceName); cond == true {
		return true, id
	}

	client := http.DefaultClient
	indexSetId := getIndexSetId()
	requestObject := StreamCreate{
		Title:                          namespaceName,
		Description:                    fmt.Sprintf("Logs for namespace %s", namespaceName),
		IndexSetId:                     indexSetId,
		RemoveMatchesFromDefaultStream: true,
	}

	body, err := json.Marshal(requestObject)

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

	req, err := http.NewRequest(http.MethodPost, getGraylogBaseUrl()+"/api/streams", bytes.NewBuffer(body))
	if err != nil {
		log.Fatal(err.Error())
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Accept", "application/json")
	req.SetBasicAuth(getGraylogSessionToken(), "session")

	resp, err := client.Do(req)
	if err != nil {
		log.Println(fmt.Sprintf("Error occured while calling Graylog for Stream creation. Error was: %s", err.Error()))
	}

	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err.Error())
	}

	switch resp.StatusCode {
	case 201:
		var stream StreamId

		err = json.Unmarshal(content, &stream)
		if err != nil {
			log.Fatal(err.Error())
		}

		createRuleForNamespace(namespaceName, stream.StreamId)

		createRoleForStreamReaders(namespaceName, stream.StreamId)
		// start stream
		startStream(stream.StreamId)
		return true, stream.StreamId
	case 403:
		log.Println("Graylog communication failed due to permission denied for user.")
		return false, ""
	default:
		log.Println(fmt.Sprintf("Graylog returned a not-OK status code when creating a stream. Code was: %d , message was: %s", resp.StatusCode, content))
		return false, ""
	}
}

func DeleteStream(namespaceName string) {
	if !isGrayLogActive() {
		return
	}

	stream, err := getStreamByNamespaceName(namespaceName)
	if err != nil {
		log.Println(fmt.Sprintf("An error occured while fetching information about the stream to be deleted. Error was: %s", err.Error()))
	}

	streamId := stream.Id

	client := http.DefaultClient

	req, err := http.NewRequest(http.MethodDelete, getGraylogBaseUrl()+fmt.Sprintf("/api/streams/%s", streamId), nil)
	if err != nil {
		log.Fatal(err.Error())
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Accept", "application/json")
	req.SetBasicAuth(getGraylogSessionToken(), "session")

	resp, err := client.Do(req)
	if err != nil {
		log.Println(fmt.Sprintf("Error occured while calling Graylog for Stream Start. Error was: %s", err.Error()))
	}

	switch resp.StatusCode {
	case 204:
		// stream deleted successfully, so delete role for it and reload local Stream cache
		deleteRoleForStreamReaders(namespaceName)
		reloadStreams()
	case 404:
		log.Println(fmt.Sprintf("Error while deleting stream: Stream %s could not be found", streamId))
	case 400:
		log.Println(fmt.Sprintf("Error while deleting stream: Stream %s was invalid", streamId))
	}
}

func GrantPermissionForStream(namespaceName, username string) bool {
	success := false
	if !isGrayLogActive() {
		return success
	}

	client := http.DefaultClient

	role, err := getRoleForNamespace(namespaceName)
	if err != nil {
		log.Println(fmt.Sprintf("Error occured while calling Graylog for retrieval of Role for Namespace. Error was: %s", err.Error()))
		return success
	}

	user, err := getGraylogUser(username)
	if err != nil {
		log.Println(fmt.Sprintf("Error occured while calling Graylog for retrieval of User. Error was: %s", err.Error()))
		return success
	}

	if contained, _ := contains(user.Roles, role.Name); contained == false {
		updatedRoles := append(user.Roles, role.Name)
		userup := UserUpdate{Roles: updatedRoles}

		body, err := json.Marshal(userup)

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

		req, err := http.NewRequest(http.MethodPut, getGraylogBaseUrl()+"/api/users/"+username, bytes.NewBuffer(body))
		if err != nil {
			log.Fatal(err.Error())
		}

		req.Header.Add("Content-Type", "application/json")
		req.Header.Add("Accept", "application/json")
		req.SetBasicAuth(getGraylogSessionToken(), "session")

		resp, err := client.Do(req)
		if err != nil {
			log.Println(fmt.Sprintf("Error occured while calling Graylog for PermissionGrant on Stream. Error was: %s", err.Error()))
		}

		switch resp.StatusCode {
		case 204:
			success = true
		case 400:
			log.Println("Graylog communication for PermissionGrant on Stream failed due to permission denied for user.")
			success = false

		case 404:
			success = false
		}
	} else {
		success = true
	}

	return success
}

func TakePermissionForStream(namespaceName, username string) bool {
	success := false
	if !isGrayLogActive() {
		return success
	}

	client := http.DefaultClient

	user, err := getGraylogUser(username)
	if err != nil {
		log.Println(fmt.Sprintf("Error occured while calling Graylog for retrieval of User. Error was: %s", err.Error()))
		return success
	}

	if contained, index := contains(user.Roles, getRoleNameForNamespace(namespaceName)); contained == true {
		// remove role from Roles slice
		updatedRoles := append(user.Roles[:index], user.Roles[index+1:]...)
		userup := UserUpdate{Roles: updatedRoles}

		body, err := json.Marshal(userup)

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

		req, err := http.NewRequest(http.MethodPut, getGraylogBaseUrl()+"/api/users/"+username, bytes.NewBuffer(body))
		if err != nil {
			log.Fatal(err.Error())
		}

		req.Header.Add("Content-Type", "application/json")
		req.Header.Add("Accept", "application/json")
		req.SetBasicAuth(getGraylogSessionToken(), "session")

		resp, err := client.Do(req)
		if err != nil {
			log.Println(fmt.Sprintf("Error occured while calling Graylog for PermissionGrant on Stream. Error was: %s", err.Error()))
		}

		switch resp.StatusCode {
		case 204:
			success = true
		case 400:
			log.Println("Graylog communication for PermissionGrant on Stream failed due to permission denied for user.")
			success = false
		}
	} else {
		success = true
	}
	return success
}

<commit_msg>added rule creation logic<commit_after>package graylog

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
)

type Streams struct {
	Total      int      `json:"total"`
	StreamList []Stream `json:"streams"`
}

type Stream struct {
	Id                             string `json:"id"`
	Title                          string `json:"title"`
	Description                    string `json:"description"`
	Rules                          []Rule `json:"rules"`
	ContentPack                    string `json:"content_pack"`
	MatchingType                   string `json:"matching_type"`
	RemoveMatchesFromDefaultStream bool   `json:"remove_matches_from_default_stream"`
	IndexSetId                     string `json:"index_set_id"`
}

type StreamId struct {
	StreamId string `json:"stream_id"`
}

type StreamCreate struct {
	Title                          string `json:"title"`
	Description                    string `json:"description"`
	RemoveMatchesFromDefaultStream bool   `json:"remove_matches_from_default_stream"`
	IndexSetId                     string `json:"index_set_id"`
}

type Rules struct {
	Total 		int `json:"total"`
	StreamRules []Rule `json:"stream_rules"`
}

type Rule struct {
	Type        int    `json:"type"`
	Value       string `json:"value"`
	Field       string `json:"field"`
	Inverted    bool   `json:"inverted"`
	Description string `json:"description"`
}

type IndexSets struct {
	Total     int        `json:"total"`
	IndexSets []IndexSet `json:"index_sets"`
}

type IndexSet struct {
	Id    string `json:"id"`
	Title string `json:"title"`
}

type UserUpdate struct {
	Roles []string `json:"roles"`
}

type User struct {
	Id       string   `json:"id"`
	Username string   `json:"username"`
	Roles    []string `json:"roles"`
}

type Role struct {
	Name        string   `json:"name"`
	Description string   `json:"description"`
	Permissions []string `json:"permissions"`
	ReadOnly    bool     `json:"read_only"`
}

func CreateStream(namespaceName string) (bool, string) {
	if !isGrayLogActive() {
		return false, ""
	}

	if cond, id := isStreamAlreadyCreated(namespaceName); cond == true {
		createRuleForNamespace(namespaceName, id)
		createRoleForStreamReaders(namespaceName, id)
		return true, id
	}

	client := http.DefaultClient
	indexSetId := getIndexSetId()
	requestObject := StreamCreate{
		Title:                          namespaceName,
		Description:                    fmt.Sprintf("Logs for namespace %s", namespaceName),
		IndexSetId:                     indexSetId,
		RemoveMatchesFromDefaultStream: true,
	}

	body, err := json.Marshal(requestObject)

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

	req, err := http.NewRequest(http.MethodPost, getGraylogBaseUrl()+"/api/streams", bytes.NewBuffer(body))
	if err != nil {
		log.Fatal(err.Error())
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Accept", "application/json")
	req.SetBasicAuth(getGraylogSessionToken(), "session")

	resp, err := client.Do(req)
	if err != nil {
		log.Println(fmt.Sprintf("Error occured while calling Graylog for Stream creation. Error was: %s", err.Error()))
	}

	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err.Error())
	}

	switch resp.StatusCode {
	case 201:
		var stream StreamId

		err = json.Unmarshal(content, &stream)
		if err != nil {
			log.Fatal(err.Error())
		}

		createRuleForNamespace(namespaceName, stream.StreamId)

		createRoleForStreamReaders(namespaceName, stream.StreamId)
		// start stream
		startStream(stream.StreamId)
		return true, stream.StreamId
	case 403:
		log.Println("Graylog communication failed due to permission denied for user.")
		return false, ""
	default:
		log.Println(fmt.Sprintf("Graylog returned a not-OK status code when creating a stream. Code was: %d , message was: %s", resp.StatusCode, content))
		return false, ""
	}
}

func DeleteStream(namespaceName string) {
	if !isGrayLogActive() {
		return
	}

	stream, err := getStreamByNamespaceName(namespaceName)
	if err != nil {
		log.Println(fmt.Sprintf("An error occured while fetching information about the stream to be deleted. Error was: %s", err.Error()))
	}

	streamId := stream.Id

	client := http.DefaultClient

	req, err := http.NewRequest(http.MethodDelete, getGraylogBaseUrl()+fmt.Sprintf("/api/streams/%s", streamId), nil)
	if err != nil {
		log.Fatal(err.Error())
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Accept", "application/json")
	req.SetBasicAuth(getGraylogSessionToken(), "session")

	resp, err := client.Do(req)
	if err != nil {
		log.Println(fmt.Sprintf("Error occured while calling Graylog for Stream Start. Error was: %s", err.Error()))
	}

	switch resp.StatusCode {
	case 204:
		// stream deleted successfully, so delete role for it and reload local Stream cache
		deleteRoleForStreamReaders(namespaceName)
		reloadStreams()
	case 404:
		log.Println(fmt.Sprintf("Error while deleting stream: Stream %s could not be found", streamId))
	case 400:
		log.Println(fmt.Sprintf("Error while deleting stream: Stream %s was invalid", streamId))
	}
}

func GrantPermissionForStream(namespaceName, username string) bool {
	success := false
	if !isGrayLogActive() {
		return success
	}

	client := http.DefaultClient

	role, err := getRoleForNamespace(namespaceName)
	if err != nil {
		log.Println(fmt.Sprintf("Error occured while calling Graylog for retrieval of Role for Namespace. Error was: %s", err.Error()))
		return success
	}

	user, err := getGraylogUser(username)
	if err != nil {
		log.Println(fmt.Sprintf("Error occured while calling Graylog for retrieval of User. Error was: %s", err.Error()))
		return success
	}

	if contained, _ := contains(user.Roles, role.Name); contained == false {
		updatedRoles := append(user.Roles, role.Name)
		userup := UserUpdate{Roles: updatedRoles}

		body, err := json.Marshal(userup)

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

		req, err := http.NewRequest(http.MethodPut, getGraylogBaseUrl()+"/api/users/"+username, bytes.NewBuffer(body))
		if err != nil {
			log.Fatal(err.Error())
		}

		req.Header.Add("Content-Type", "application/json")
		req.Header.Add("Accept", "application/json")
		req.SetBasicAuth(getGraylogSessionToken(), "session")

		resp, err := client.Do(req)
		if err != nil {
			log.Println(fmt.Sprintf("Error occured while calling Graylog for PermissionGrant on Stream. Error was: %s", err.Error()))
		}

		switch resp.StatusCode {
		case 204:
			success = true
		case 400:
			log.Println("Graylog communication for PermissionGrant on Stream failed due to permission denied for user.")
			success = false

		case 404:
			success = false
		}
	} else {
		success = true
	}

	return success
}

func TakePermissionForStream(namespaceName, username string) bool {
	success := false
	if !isGrayLogActive() {
		return success
	}

	client := http.DefaultClient

	user, err := getGraylogUser(username)
	if err != nil {
		log.Println(fmt.Sprintf("Error occured while calling Graylog for retrieval of User. Error was: %s", err.Error()))
		return success
	}

	if contained, index := contains(user.Roles, getRoleNameForNamespace(namespaceName)); contained == true {
		// remove role from Roles slice
		updatedRoles := append(user.Roles[:index], user.Roles[index+1:]...)
		userup := UserUpdate{Roles: updatedRoles}

		body, err := json.Marshal(userup)

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

		req, err := http.NewRequest(http.MethodPut, getGraylogBaseUrl()+"/api/users/"+username, bytes.NewBuffer(body))
		if err != nil {
			log.Fatal(err.Error())
		}

		req.Header.Add("Content-Type", "application/json")
		req.Header.Add("Accept", "application/json")
		req.SetBasicAuth(getGraylogSessionToken(), "session")

		resp, err := client.Do(req)
		if err != nil {
			log.Println(fmt.Sprintf("Error occured while calling Graylog for PermissionGrant on Stream. Error was: %s", err.Error()))
		}

		switch resp.StatusCode {
		case 204:
			success = true
		case 400:
			log.Println("Graylog communication for PermissionGrant on Stream failed due to permission denied for user.")
			success = false
		}
	} else {
		success = true
	}
	return success
}

<|endoftext|>


<commit_before>package etcd

import (
	"encoding/json"
	"fmt"
	"github.com/coreos/etcd/store"
	"io/ioutil"
	"net/http"
	"path"
	"strings"
)

type Options map[string]interface{}

var (
	// Making a map to make it easier to test existence
	validGetOptions = map[string]bool{
		"recursive":  true,
		"consistent": true,
		"sorted":     true,
		"wait":       true,
		"wait_index": true,
	}
)

func (c *Client) Get(key string) ([]*store.Response, error) {
	return c.GetWithOptions(key, nil)
}

func (c *Client) GetWithOptions(key string, options Options) ([]*store.Response, error) {
	logger.Debugf("get %s [%s]", key, c.cluster.Leader)

	p := path.Join("keys", key)
	if options != nil {
		p += "?"
		optionArr := make([]string, 0)
		for opKey, opVal := range options {
			if validGetOptions[opKey] {
				optionArr = append(optionArr, fmt.Sprintf("%v=%v", opKey, opVal))
			} else {
				return nil, fmt.Errorf("Invalid option: %v", opKey)
			}
		}
		p += strings.Join(optionArr, "&")
	}

	resp, err := c.sendRequest("GET", p, "")

	if err != nil {
		return nil, err
	}

	b, err := ioutil.ReadAll(resp.Body)

	resp.Body.Close()

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {

		return nil, handleError(b)
	}

	return convertGetResponse(b)
}

// GetTo gets the value of the key from a given machine address.
// If the given machine is not available it returns an error.
// Mainly use for testing purpose
func (c *Client) GetFrom(key string, addr string) ([]*store.Response, error) {
	httpPath := c.createHttpPath(addr, path.Join(version, "keys", key))

	resp, err := c.httpClient.Get(httpPath)

	if err != nil {
		return nil, err
	}

	b, err := ioutil.ReadAll(resp.Body)

	resp.Body.Close()

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, handleError(b)
	}

	return convertGetResponse(b)
}

// Convert byte stream to response.
func convertGetResponse(b []byte) ([]*store.Response, error) {

	var results []*store.Response
	var result *store.Response

	err := json.Unmarshal(b, &result)

	if err != nil {
		err = json.Unmarshal(b, &results)

		if err != nil {
			return nil, err
		}

	} else {
		results = make([]*store.Response, 1)
		results[0] = result
	}
	return results, nil
}
<commit_msg>Allows passing options to GetFrom<commit_after>package etcd

import (
	"encoding/json"
	"fmt"
	"github.com/coreos/etcd/store"
	"io/ioutil"
	"net/http"
	"path"
	"strings"
)

type Options map[string]interface{}

var (
	// Making a map to make it easier to test existence
	validGetOptions = map[string]bool{
		"recursive":  true,
		"consistent": true,
		"sorted":     true,
		"wait":       true,
		"wait_index": true,
	}
)

// Get the value of the given key
func (c *Client) Get(key string) ([]*store.Response, error) {
	return c.GetWithOptions(key, nil)
}

// The same with Get, but allows passing options
func (c *Client) GetWithOptions(key string, options Options) ([]*store.Response, error) {
	logger.Debugf("get %s [%s]", key, c.cluster.Leader)

	p := path.Join("keys", key)
	if options != nil {
		str, err := optionsToString(options)
		if err != nil {
			return nil, err
		}
		p += str
	}

	resp, err := c.sendRequest("GET", p, "")

	if err != nil {
		return nil, err
	}

	b, err := ioutil.ReadAll(resp.Body)

	resp.Body.Close()

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {

		return nil, handleError(b)
	}

	return convertGetResponse(b)
}

// GetTo gets the value of the key from a given machine address.
// If the given machine is not available it returns an error.
// Mainly use for testing purpose
func (c *Client) GetFrom(key string, addr string) ([]*store.Response, error) {
	return c.GetFromWithOptions(key, addr, nil)
}

// The same with GetFrom, but allows passing options
func (c *Client) GetFromWithOptions(key string, addr string, options Options) ([]*store.Response, error) {
	httpPath := c.createHttpPath(addr, path.Join(version, "keys", key))

	if options != nil {
		str, err := optionsToString(options)
		if err != nil {
			return nil, err
		}
		httpPath += str
	}

	resp, err := c.httpClient.Get(httpPath)

	if err != nil {
		return nil, err
	}

	b, err := ioutil.ReadAll(resp.Body)

	resp.Body.Close()

	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, handleError(b)
	}

	return convertGetResponse(b)
}

// Convert byte stream to response.
func convertGetResponse(b []byte) ([]*store.Response, error) {

	var results []*store.Response
	var result *store.Response

	err := json.Unmarshal(b, &result)

	if err != nil {
		err = json.Unmarshal(b, &results)

		if err != nil {
			return nil, err
		}

	} else {
		results = make([]*store.Response, 1)
		results[0] = result
	}
	return results, nil
}

// Convert options to a string of HTML parameters
func optionsToString(options Options) (string, error) {
	p := "?"
	optionArr := make([]string, 0)
	for opKey, opVal := range options {
		if validGetOptions[opKey] {
			optionArr = append(optionArr, fmt.Sprintf("%v=%v", opKey, opVal))
		} else {
			return "", fmt.Errorf("Invalid option: %v", opKey)
		}
	}
	p += strings.Join(optionArr, "&")
	return p, nil
}
<|endoftext|>


<commit_before>/*
 * (C) Copyright 2014 Nuxeo SA (http://nuxeo.com/) and contributors.
 *
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Apache License Version 2.0
 * which accompanies this distribution, and is available at
 * http://www.apache.org/licenses/
 *
 * See the Apache Licence for more details.
 *
 * Contributors:
 *     nuxeo.io Team
 */

package main

import (
	"github.com/coreos/go-etcd/etcd"
	"time"
	"github.com/golang/glog"
	"os/exec"
	"strconv"
	"os"
)

type EtcdCron struct {
	client   *etcd.Client
	config   *Config
}

func NewEtcdCron(config *Config) (*EtcdCron, error) {
	client, err := config.getEtcdClient()
	if err != nil {
		return nil, err
	}
	return &EtcdCron{client, config}, nil
}

func (etcdcron *EtcdCron) init() {
	etcdcron.start()
}

func (etcdcron *EtcdCron) start() {
	cronDuration, _ := strconv.Atoi(etcdcron.config.cronDuration)
	interval := time.Duration(cronDuration) * time.Minute
	ticker := time.NewTicker(interval)
	for {
		<-ticker.C
		// Check every 5 minutes all services lastAccess etcd date entry
		response, err := etcdcron.client.Get(etcdcron.config.servicePrefix, true, true)
		if err == nil {
			for _, serviceNode := range response.Node.Nodes {
				etcdcron.checkServiceAccess(serviceNode, response.Action)
			}
		}
		ticker = time.NewTicker(interval)
	}
}

func (etcdcron *EtcdCron) checkServiceAccess(node *etcd.Node, action string) {
	serviceName := etcdcron.config.getServiceForNode(node, etcdcron.config)

	// Get service's root node instead of changed node.
	serviceNode, _ := etcdcron.client.Get(etcdcron.config.servicePrefix+"/"+serviceName, true, true)

	for _, indexNode := range serviceNode.Node.Nodes {

		serviceIndex := etcdcron.config.getServiceIndexForNode(indexNode, etcdcron.config)
		serviceKey := etcdcron.config.servicePrefix + "/" + serviceName + "/" + serviceIndex
		lastAccessKey := serviceKey + "/lastAccess"
		statusKey := serviceKey + "/status"

		response, err := etcdcron.client.Get(serviceKey, true, true)

		if err == nil {

			service := &Service{}
			service.index = serviceIndex
			service.nodeKey = serviceKey
			service.name = "nxio."+serviceName+"."+serviceIndex+".service"

			for _, node := range response.Node.Nodes {
				switch node.Key {
				case statusKey:
					service.status = &Status{}
				for _, subNode := range node.Nodes {
					switch subNode.Key {
					case statusKey + "/alive":
						service.status.alive = subNode.Value
					case statusKey + "/current":
						service.status.current = subNode.Value
					case statusKey + "/expected":
						service.status.expected = subNode.Value
					}
				}
				case lastAccessKey:
					lastAccess := node.Value
					lastAccessTime, err := time.Parse(TIME_FORMAT, lastAccess)
					if err != nil {
						glog.Errorf("Error parsing last access date with service %s: %s", service.name, err)
						break
					}
					service.lastAccess = &lastAccessTime
				}
			}

			parameter, _ := strconv.Atoi(etcdcron.config.passiveLimitDuration)
			passiveLimitDuration := time.Duration(parameter) * time.Hour

			// Checking if the service should be passivated or not
			if service.lastAccess != nil && service.status != nil {
				if service.status.current == STARTED_STATUS && etcdcron.hasToBePassivated(service, passiveLimitDuration) {
					responseCurrent, error := etcdcron.client.Set(statusKey+"/current", PASSIVATED_STATUS, 0)
					if error != nil && responseCurrent == nil {
						glog.Errorf("Setting status current to 'passivated' has failed for Service "+service.name+": %s", err)
					}
					response, error := etcdcron.client.Set(statusKey+"/expected", PASSIVATED_STATUS, 0)
					if error != nil && response == nil {
						glog.Errorf("Setting status expected to 'passivated' has failed for Service "+service.name+": %s", err)
					}
					cmd := exec.Command("/usr/bin/fleetctl", "--endpoint="+etcdcron.config.etcdAddress, "stop", service.name)
					cmd.Stdin = os.Stdin
					cmd.Stdout = os.Stdout
					cmd.Stderr = os.Stderr
					err := cmd.Run()
					if err != nil {
						glog.Errorf("Service "+service.name+" passivation has failed: %s", err)
						break
					}
					glog.Infof("Service %s passivated", service.name)
				}
			}
		}
	}
}

func (etcdcron *EtcdCron) hasToBePassivated(service *Service, passiveLimitDuration time.Duration) bool {
	return time.Now().After(service.lastAccess.Add(passiveLimitDuration)) && service.status.current == STARTED_STATUS
}
<commit_msg>NXIO-361 change unit name<commit_after>/*
 * (C) Copyright 2014 Nuxeo SA (http://nuxeo.com/) and contributors.
 *
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Apache License Version 2.0
 * which accompanies this distribution, and is available at
 * http://www.apache.org/licenses/
 *
 * See the Apache Licence for more details.
 *
 * Contributors:
 *     nuxeo.io Team
 */

package main

import (
	"github.com/coreos/go-etcd/etcd"
	"time"
	"github.com/golang/glog"
	"os/exec"
	"strconv"
	"os"
)

type EtcdCron struct {
	client   *etcd.Client
	config   *Config
}

func NewEtcdCron(config *Config) (*EtcdCron, error) {
	client, err := config.getEtcdClient()
	if err != nil {
		return nil, err
	}
	return &EtcdCron{client, config}, nil
}

func (etcdcron *EtcdCron) init() {
	etcdcron.start()
}

func (etcdcron *EtcdCron) start() {
	cronDuration, _ := strconv.Atoi(etcdcron.config.cronDuration)
	interval := time.Duration(cronDuration) * time.Minute
	ticker := time.NewTicker(interval)
	for {
		<-ticker.C
		// Check every 5 minutes all services lastAccess etcd date entry
		response, err := etcdcron.client.Get(etcdcron.config.servicePrefix, true, true)
		if err == nil {
			for _, serviceNode := range response.Node.Nodes {
				etcdcron.checkServiceAccess(serviceNode, response.Action)
			}
		}
		ticker = time.NewTicker(interval)
	}
}

func (etcdcron *EtcdCron) checkServiceAccess(node *etcd.Node, action string) {
	serviceName := etcdcron.config.getServiceForNode(node, etcdcron.config)

	// Get service's root node instead of changed node.
	serviceNode, _ := etcdcron.client.Get(etcdcron.config.servicePrefix+"/"+serviceName, true, true)

	for _, indexNode := range serviceNode.Node.Nodes {

		serviceIndex := etcdcron.config.getServiceIndexForNode(indexNode, etcdcron.config)
		serviceKey := etcdcron.config.servicePrefix + "/" + serviceName + "/" + serviceIndex
		lastAccessKey := serviceKey + "/lastAccess"
		statusKey := serviceKey + "/status"

		response, err := etcdcron.client.Get(serviceKey, true, true)

		if err == nil {

			service := &Service{}
			service.index = serviceIndex
			service.nodeKey = serviceKey

			service.name = "nxio@" + strings.Split(serviceName, "_")[1] + ".service"

			for _, node := range response.Node.Nodes {
				switch node.Key {
				case statusKey:
					service.status = &Status{}
				for _, subNode := range node.Nodes {
					switch subNode.Key {
					case statusKey + "/alive":
						service.status.alive = subNode.Value
					case statusKey + "/current":
						service.status.current = subNode.Value
					case statusKey + "/expected":
						service.status.expected = subNode.Value
					}
				}
				case lastAccessKey:
					lastAccess := node.Value
					lastAccessTime, err := time.Parse(TIME_FORMAT, lastAccess)
					if err != nil {
						glog.Errorf("Error parsing last access date with service %s: %s", service.name, err)
						break
					}
					service.lastAccess = &lastAccessTime
				}
			}

			parameter, _ := strconv.Atoi(etcdcron.config.passiveLimitDuration)
			passiveLimitDuration := time.Duration(parameter) * time.Hour

			// Checking if the service should be passivated or not
			if service.lastAccess != nil && service.status != nil {
				if service.status.current == STARTED_STATUS && etcdcron.hasToBePassivated(service, passiveLimitDuration) {
					responseCurrent, error := etcdcron.client.Set(statusKey+"/current", PASSIVATED_STATUS, 0)
					if error != nil && responseCurrent == nil {
						glog.Errorf("Setting status current to 'passivated' has failed for Service "+service.name+": %s", err)
					}
					response, error := etcdcron.client.Set(statusKey+"/expected", PASSIVATED_STATUS, 0)
					if error != nil && response == nil {
						glog.Errorf("Setting status expected to 'passivated' has failed for Service "+service.name+": %s", err)
					}
					cmd := exec.Command("/usr/bin/fleetctl", "--endpoint="+etcdcron.config.etcdAddress, "stop", service.name)
					cmd.Stdin = os.Stdin
					cmd.Stdout = os.Stdout
					cmd.Stderr = os.Stderr
					err := cmd.Run()
					if err != nil {
						glog.Errorf("Service "+service.name+" passivation has failed: %s", err)
						break
					}
					glog.Infof("Service %s passivated", service.name)
				}
			}
		}
	}
}

func (etcdcron *EtcdCron) hasToBePassivated(service *Service, passiveLimitDuration time.Duration) bool {
	return time.Now().After(service.lastAccess.Add(passiveLimitDuration)) && service.status.current == STARTED_STATUS
}
<|endoftext|>