<commit_before>/*
Copyright 2016-2017 Mirantis

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 libvirttools

import (
	"encoding/json"
	"fmt"

	"github.com/Mirantis/virtlet/pkg/utils"
	"github.com/golang/glog"
	libvirt "github.com/libvirt/libvirt-go"
)

const (
	defaultCapacity     = 1024
	defaultCapacityUnit = "MB"
	poolTypeDir         = "dir"
)

type Volume struct {
	tool   StorageOperations
	Name   string
	volume *libvirt.StorageVol
}

type VirtletVolume struct {
	Name         string `json:"Name"`
	Format       string `json:"Format"`
	Capacity     int    `json:"Capacity"`
	CapacityUnit string `json:"CapacityUnit"`
}

func (vol *VirtletVolume) UnmarshalJSON(data []byte) error {
	// volAlias is needed to prevent recursive calls to UnmarshalJSON
	type volAlias VirtletVolume
	volWithDefaults := &volAlias{
		Format:       "qcow2",
		Capacity:     defaultCapacity,
		CapacityUnit: defaultCapacityUnit,
	}

	err := json.Unmarshal(data, volWithDefaults)

	if err == nil && volWithDefaults.Name == "" {
		return fmt.Errorf("Validation failed for volumes definition within pod's annotations: volume name is mandatory.")
	}

	*vol = VirtletVolume(*volWithDefaults)
	return err
}

func (v *Volume) Remove() error {
	return v.tool.RemoveVolume(v.volume)
}

func (v *Volume) GetPath() (string, error) {
	return v.tool.VolumeGetPath(v.volume)
}

type VolumeInfo struct {
	tool StorageOperations
	Name string
	Size uint64
}

func (v *Volume) Info() (*VolumeInfo, error) {
	return volumeInfo(v.tool, v.Name, v.volume)
}

func volumeInfo(tool StorageOperations, name string, volume *libvirt.StorageVol) (*VolumeInfo, error) {
	volInfo, err := tool.VolumeGetInfo(volume)
	if err != nil {
		return nil, err
	}
	return &VolumeInfo{Name: name, Size: volInfo.Capacity}, nil
}

type Pool struct {
	tool       StorageOperations
	pool       *libvirt.StoragePool
	volumesDir string
	poolType   string
}

type PoolSet map[string]*Pool

var DefaultPools PoolSet = PoolSet{
	"default": &Pool{volumesDir: "/var/lib/libvirt/images", poolType: poolTypeDir},
	"volumes": &Pool{volumesDir: "/var/lib/virtlet", poolType: poolTypeDir},
}

func generatePoolXML(name string, path string, poolType string) string {
	poolXML := `
<pool type="%s">
    <name>%s</name>
    <target>
	<path>%s</path>
    </target>
</pool>`
	return fmt.Sprintf(poolXML, poolType, name, path)
}

func createPool(tool StorageOperations, name string, path string, poolType string) (*Pool, error) {
	poolXML := generatePoolXML(name, path, poolType)

	glog.V(2).Infof("Creating storage pool (name: %s, path: %s)", name, path)
	pool, err := tool.CreateFromXML(poolXML)
	if err != nil {
		return nil, err
	}
	return &Pool{tool: tool, pool: pool, volumesDir: path, poolType: poolType}, nil
}

func LookupStoragePool(tool StorageOperations, name string) (*Pool, error) {
	poolInfo, exist := DefaultPools[name]
	if !exist {
		return nil, fmt.Errorf("pool with name '%s' is unknown", name)
	}

	pool, _ := tool.LookupByName(name)
	if pool == nil {
		return createPool(tool, name, poolInfo.volumesDir, poolInfo.poolType)
	}
	// TODO: reset libvirt error

	return &Pool{tool: tool, pool: pool, volumesDir: poolInfo.volumesDir, poolType: poolInfo.poolType}, nil
}

func (p *Pool) RemoveVolume(name string) error {
	vol, err := p.LookupVolume(name)
	if err != nil {
		return err
	}
	return vol.Remove()
}

func (p *Pool) CreateVolume(name, volXML string) (*Volume, error) {
	vol, err := p.tool.CreateVolFromXML(p.pool, volXML)
	if err != nil {
		return nil, err
	}
	return &Volume{tool: p.tool, Name: name, volume: vol}, nil
}

func (p *Pool) LookupVolume(name string) (*Volume, error) {
	vol, err := p.tool.LookupVolumeByName(p.pool, name)
	if err != nil {
		return nil, err
	}
	return &Volume{tool: p.tool, Name: name, volume: vol}, nil
}

func (p *Pool) ListVolumes() ([]*VolumeInfo, error) {
	volumes, err := p.tool.ListAllVolumes(p.pool)
	if err != nil {
		return nil, err
	}

	volumeInfos := make([]*VolumeInfo, 0, len(volumes))

	for _, volume := range volumes {
		name, err := p.tool.VolumeGetName(&volume)
		volInfo, err := volumeInfo(p.tool, name, &volume)
		if err != nil {
			return nil, err
		}

		volumeInfos = append(volumeInfos, volInfo)
	}

	return volumeInfos, nil
}

type StorageTool struct {
	name string
	tool StorageOperations
	pool *Pool
}

func NewStorageTool(conn *libvirt.Connect, poolName string) (*StorageTool, error) {
	tool := NewLibvirtStorageOperations(conn)
	pool, err := LookupStoragePool(tool, poolName)
	if err != nil {
		return nil, err
	}
	return &StorageTool{name: poolName, tool: tool, pool: pool}, nil
}

func (s *StorageTool) GenerateVolumeXML(shortName string, capacity int, capacityUnit string, path string) string {
	volXML := `
<volume>
    <name>%s</name>
    <allocation>0</allocation>
    <capacity unit="%s">%d</capacity>
    <target>
        <path>%s</path>
    </target>
</volume>`
	return fmt.Sprintf(volXML, shortName, capacityUnit, capacity, path)
}

func (s *StorageTool) CreateVolume(name string, capacity int, capacityUnit string) (*Volume, error) {
	volumeXML := `
<volume>
    <name>%s</name>
    <allocation>0</allocation>
    <capacity unit="%s">%d</capacity>
</volume>`
	volumeXML = fmt.Sprintf(volumeXML, name, capacityUnit, capacity)
	glog.V(2).Infof("Create volume using XML description: %s", volumeXML)
	return s.pool.CreateVolume(name, volumeXML)
}

func (s *StorageTool) CreateSnapshot(name string, capacity int, capacityUnit string, backingStorePath string) (*Volume, error) {
	snapshotXML := `
<volume type='file'>
    <name>%s</name>
    <allocation>0</allocation>
    <capacity unit="%s">%d</capacity>
    <target>
         <format type='qcow2'/>
    </target>
    <backingStore>
         <path>%s</path>
         <format type='qcow2'/>
     </backingStore>
</volume>`
	snapshotXML = fmt.Sprintf(snapshotXML, name, capacityUnit, capacity, backingStorePath)
	glog.V(2).Infof("Create volume using XML description: %s", snapshotXML)
	return s.pool.CreateVolume(name, snapshotXML)
}

func (s *StorageTool) LookupVolume(name string) (*Volume, error) {
	return s.pool.LookupVolume(name)
}

func (s *StorageTool) RemoveVolume(name string) error {
	return s.pool.RemoveVolume(name)
}

func (s *StorageTool) ListVolumes() ([]*VolumeInfo, error) {
	return s.pool.ListVolumes()
}

func (s *StorageTool) CleanAttachedVolumes(virtletVolsDesc string, containerId string) error {

	var vols []VirtletVolume
	if err := json.Unmarshal([]byte(virtletVolsDesc), &vols); err != nil {
		glog.Errorf("Error when unmarshalling json string with volumes description '%s' for container %s: %v", virtletVolsDesc, containerId, err)
	}

	for _, virtletVol := range vols {
		volName := containerId + "-" + virtletVol.Name
		if err := s.RemoveVolume(volName); err != nil {
			glog.Errorf("Error during removal of volume '%s' for container %s: %v", volName, containerId, err)
			return err
		}
	}

	return nil
}

func (s *StorageTool) CreateVolumesToBeAttached(virtletVolsDesc string, containerId string, letterInd int) ([]string, error) {
	var volumesXML []string
	var virtletVols []VirtletVolume
	if err := json.Unmarshal([]byte(virtletVolsDesc), &virtletVols); err != nil {
		glog.Errorf("Error when unmarshalling json string with volumes description '%s' for container %s: %v", virtletVolsDesc, containerId, err)
	}

	for _, virtletVol := range virtletVols {
		if letterInd == len(diskLetters) {
			fmt.Errorf("Had to omit creating and attaching of one ore more volumes. Limit on number is: %d", letterInd)
			break
		}

		volName := containerId + "-" + virtletVol.Name
		vol, err := s.CreateVolume(volName, defaultCapacity, defaultCapacityUnit)
		if err != nil {
			glog.Errorf("Error during creation of volume'%s' with virtlet description %s: %v", volName, virtletVol, err)
			return nil, err
		}

		path, err := vol.GetPath()
		if err != nil {
			return nil, err
		}

		err = utils.FormatDisk(path)
		if err != nil {
			return nil, err
		}

		volXML := fmt.Sprintf(volXMLTemplate, path, "vd"+diskLetters[letterInd])
		volumesXML = append(volumesXML, volXML)
	}

	return volumesXML, nil
}

func (s *StorageTool) PullImageToVolume(path, volumeName string) error {
	libvirtFilePath := fmt.Sprintf("/var/lib/libvirt/images/%s", volumeName)
	volXML := s.GenerateVolumeXML(volumeName, 5, "G", libvirtFilePath)

	return s.tool.PullImageToVolume(s.pool.pool, volumeName, path, volXML)
}
<commit_msg>libvirttools: Import image using it's real size<commit_after>/*
Copyright 2016-2017 Mirantis

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 libvirttools

import (
	"encoding/json"
	"fmt"
	"os"

	"github.com/Mirantis/virtlet/pkg/utils"
	"github.com/golang/glog"
	libvirt "github.com/libvirt/libvirt-go"
)

const (
	defaultCapacity     = 1024
	defaultCapacityUnit = "MB"
	poolTypeDir         = "dir"
)

type Volume struct {
	tool   StorageOperations
	Name   string
	volume *libvirt.StorageVol
}

type VirtletVolume struct {
	Name         string `json:"Name"`
	Format       string `json:"Format"`
	Capacity     int    `json:"Capacity"`
	CapacityUnit string `json:"CapacityUnit"`
}

func (vol *VirtletVolume) UnmarshalJSON(data []byte) error {
	// volAlias is needed to prevent recursive calls to UnmarshalJSON
	type volAlias VirtletVolume
	volWithDefaults := &volAlias{
		Format:       "qcow2",
		Capacity:     defaultCapacity,
		CapacityUnit: defaultCapacityUnit,
	}

	err := json.Unmarshal(data, volWithDefaults)

	if err == nil && volWithDefaults.Name == "" {
		return fmt.Errorf("Validation failed for volumes definition within pod's annotations: volume name is mandatory.")
	}

	*vol = VirtletVolume(*volWithDefaults)
	return err
}

func (v *Volume) Remove() error {
	return v.tool.RemoveVolume(v.volume)
}

func (v *Volume) GetPath() (string, error) {
	return v.tool.VolumeGetPath(v.volume)
}

type VolumeInfo struct {
	tool StorageOperations
	Name string
	Size uint64
}

func (v *Volume) Info() (*VolumeInfo, error) {
	return volumeInfo(v.tool, v.Name, v.volume)
}

func volumeInfo(tool StorageOperations, name string, volume *libvirt.StorageVol) (*VolumeInfo, error) {
	volInfo, err := tool.VolumeGetInfo(volume)
	if err != nil {
		return nil, err
	}
	return &VolumeInfo{Name: name, Size: volInfo.Capacity}, nil
}

type Pool struct {
	tool       StorageOperations
	pool       *libvirt.StoragePool
	volumesDir string
	poolType   string
}

type PoolSet map[string]*Pool

var DefaultPools PoolSet = PoolSet{
	"default": &Pool{volumesDir: "/var/lib/libvirt/images", poolType: poolTypeDir},
	"volumes": &Pool{volumesDir: "/var/lib/virtlet", poolType: poolTypeDir},
}

func generatePoolXML(name string, path string, poolType string) string {
	poolXML := `
<pool type="%s">
    <name>%s</name>
    <target>
	<path>%s</path>
    </target>
</pool>`
	return fmt.Sprintf(poolXML, poolType, name, path)
}

func createPool(tool StorageOperations, name string, path string, poolType string) (*Pool, error) {
	poolXML := generatePoolXML(name, path, poolType)

	glog.V(2).Infof("Creating storage pool (name: %s, path: %s)", name, path)
	pool, err := tool.CreateFromXML(poolXML)
	if err != nil {
		return nil, err
	}
	return &Pool{tool: tool, pool: pool, volumesDir: path, poolType: poolType}, nil
}

func LookupStoragePool(tool StorageOperations, name string) (*Pool, error) {
	poolInfo, exist := DefaultPools[name]
	if !exist {
		return nil, fmt.Errorf("pool with name '%s' is unknown", name)
	}

	pool, _ := tool.LookupByName(name)
	if pool == nil {
		return createPool(tool, name, poolInfo.volumesDir, poolInfo.poolType)
	}
	// TODO: reset libvirt error

	return &Pool{tool: tool, pool: pool, volumesDir: poolInfo.volumesDir, poolType: poolInfo.poolType}, nil
}

func (p *Pool) RemoveVolume(name string) error {
	vol, err := p.LookupVolume(name)
	if err != nil {
		return err
	}
	return vol.Remove()
}

func (p *Pool) CreateVolume(name, volXML string) (*Volume, error) {
	vol, err := p.tool.CreateVolFromXML(p.pool, volXML)
	if err != nil {
		return nil, err
	}
	return &Volume{tool: p.tool, Name: name, volume: vol}, nil
}

func (p *Pool) LookupVolume(name string) (*Volume, error) {
	vol, err := p.tool.LookupVolumeByName(p.pool, name)
	if err != nil {
		return nil, err
	}
	return &Volume{tool: p.tool, Name: name, volume: vol}, nil
}

func (p *Pool) ListVolumes() ([]*VolumeInfo, error) {
	volumes, err := p.tool.ListAllVolumes(p.pool)
	if err != nil {
		return nil, err
	}

	volumeInfos := make([]*VolumeInfo, 0, len(volumes))

	for _, volume := range volumes {
		name, err := p.tool.VolumeGetName(&volume)
		volInfo, err := volumeInfo(p.tool, name, &volume)
		if err != nil {
			return nil, err
		}

		volumeInfos = append(volumeInfos, volInfo)
	}

	return volumeInfos, nil
}

type StorageTool struct {
	name string
	tool StorageOperations
	pool *Pool
}

func NewStorageTool(conn *libvirt.Connect, poolName string) (*StorageTool, error) {
	tool := NewLibvirtStorageOperations(conn)
	pool, err := LookupStoragePool(tool, poolName)
	if err != nil {
		return nil, err
	}
	return &StorageTool{name: poolName, tool: tool, pool: pool}, nil
}

func (s *StorageTool) GenerateVolumeXML(shortName string, capacity int64, capacityUnit string, path string) string {
	volXML := `
<volume>
    <name>%s</name>
    <allocation>0</allocation>
    <capacity unit="%s">%d</capacity>
    <target>
        <path>%s</path>
    </target>
</volume>`
	return fmt.Sprintf(volXML, shortName, capacityUnit, capacity, path)
}

func (s *StorageTool) CreateVolume(name string, capacity int, capacityUnit string) (*Volume, error) {
	volumeXML := `
<volume>
    <name>%s</name>
    <allocation>0</allocation>
    <capacity unit="%s">%d</capacity>
</volume>`
	volumeXML = fmt.Sprintf(volumeXML, name, capacityUnit, capacity)
	glog.V(2).Infof("Create volume using XML description: %s", volumeXML)
	return s.pool.CreateVolume(name, volumeXML)
}

func (s *StorageTool) CreateSnapshot(name string, capacity int, capacityUnit string, backingStorePath string) (*Volume, error) {
	snapshotXML := `
<volume type='file'>
    <name>%s</name>
    <allocation>0</allocation>
    <capacity unit="%s">%d</capacity>
    <target>
         <format type='qcow2'/>
    </target>
    <backingStore>
         <path>%s</path>
         <format type='qcow2'/>
     </backingStore>
</volume>`
	snapshotXML = fmt.Sprintf(snapshotXML, name, capacityUnit, capacity, backingStorePath)
	glog.V(2).Infof("Create volume using XML description: %s", snapshotXML)
	return s.pool.CreateVolume(name, snapshotXML)
}

func (s *StorageTool) LookupVolume(name string) (*Volume, error) {
	return s.pool.LookupVolume(name)
}

func (s *StorageTool) RemoveVolume(name string) error {
	return s.pool.RemoveVolume(name)
}

func (s *StorageTool) ListVolumes() ([]*VolumeInfo, error) {
	return s.pool.ListVolumes()
}

func (s *StorageTool) CleanAttachedVolumes(virtletVolsDesc string, containerId string) error {

	var vols []VirtletVolume
	if err := json.Unmarshal([]byte(virtletVolsDesc), &vols); err != nil {
		glog.Errorf("Error when unmarshalling json string with volumes description '%s' for container %s: %v", virtletVolsDesc, containerId, err)
	}

	for _, virtletVol := range vols {
		volName := containerId + "-" + virtletVol.Name
		if err := s.RemoveVolume(volName); err != nil {
			glog.Errorf("Error during removal of volume '%s' for container %s: %v", volName, containerId, err)
			return err
		}
	}

	return nil
}

func (s *StorageTool) CreateVolumesToBeAttached(virtletVolsDesc string, containerId string, letterInd int) ([]string, error) {
	var volumesXML []string
	var virtletVols []VirtletVolume
	if err := json.Unmarshal([]byte(virtletVolsDesc), &virtletVols); err != nil {
		glog.Errorf("Error when unmarshalling json string with volumes description '%s' for container %s: %v", virtletVolsDesc, containerId, err)
	}

	for _, virtletVol := range virtletVols {
		if letterInd == len(diskLetters) {
			fmt.Errorf("Had to omit creating and attaching of one ore more volumes. Limit on number is: %d", letterInd)
			break
		}

		volName := containerId + "-" + virtletVol.Name
		vol, err := s.CreateVolume(volName, defaultCapacity, defaultCapacityUnit)
		if err != nil {
			glog.Errorf("Error during creation of volume'%s' with virtlet description %s: %v", volName, virtletVol, err)
			return nil, err
		}

		path, err := vol.GetPath()
		if err != nil {
			return nil, err
		}

		err = utils.FormatDisk(path)
		if err != nil {
			return nil, err
		}

		volXML := fmt.Sprintf(volXMLTemplate, path, "vd"+diskLetters[letterInd])
		volumesXML = append(volumesXML, volXML)
	}

	return volumesXML, nil
}

func (s *StorageTool) PullImageToVolume(path, volumeName string) error {
	imageSize, err := getFileSize(path)
	if err != nil {
		return err
	}
	libvirtFilePath := fmt.Sprintf("/var/lib/libvirt/images/%s", volumeName)
	volXML := s.GenerateVolumeXML(volumeName, imageSize, "B", libvirtFilePath)

	return s.tool.PullImageToVolume(s.pool.pool, volumeName, path, volXML)
}

func getFileSize(path string) (int64, error) {
	fileInfo, err := os.Stat(path)
	if err != nil {
		return 0, err
	}
	return fileInfo.Size(), nil
}
<|endoftext|>


<commit_before>package smokescreen

import (
	"errors"
	"fmt"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"os"
	"strings"
)

type EgressAclRule struct {
	Project    string
	Policy     ConfigEnforcementPolicy
	DomainGlob []string
}

type EgressAclConfig struct {
	Services map[string]EgressAclRule
	Default  *EgressAclRule
}

func (ew *EgressAclConfig) Decide(fromService string, toHost string) (EgressAclDecision, error) {

	var (
		found   bool
		service EgressAclRule
	)
	if fromService == "" {
		found = false
	} else {
		service, found = ew.Services[fromService]
	}

	if !found && ew.Default != nil {
		found = true
		service = *ew.Default
	}

	if !found {
		return 0, fmt.Errorf("unknown role role=%#v", fromService)
	}

	if service.Policy == ConfigEnforcementPolicyOpen {
		return EgressAclDecisionAllow, nil
	}

	matches := false
	for _, host := range service.DomainGlob {
		if len(host) > 0 && host[0] == '*' {
			postfix := host[1:]
			if strings.HasSuffix(toHost, postfix) {
				matches = true
				break
			}
		} else {
			if host == toHost {
				matches = true
				break
			}
		}
	}

	switch service.Policy {
	case ConfigEnforcementPolicyReport:
		if matches {
			return EgressAclDecisionAllow, nil
		} else {
			return EgressAclDecisionAllowAndReport, nil
		}
	case ConfigEnforcementPolicyEnforce:
		if matches {
			return EgressAclDecisionAllow, nil
		} else {
			return EgressAclDecisionDeny, nil
		}
	default:
		return 0, errors.New("unexpected state")
	}
}

func (ew *EgressAclConfig) Project(fromService string) (string, error) {
	service, found := ew.Services[fromService]

	if found {
		return service.Project, nil
	}

	if ew.Default != nil {
		return ew.Default.Project, nil
	}

	return "", fmt.Errorf("warn: No known project: role=%#v\n", fromService)
}

// Configuration

type ServiceRule struct {
	Name         string   `yaml:"name"`
	Project      string   `yaml:"project"`
	Action       string   `yaml:"action"`
	AllowedHosts []string `yaml:"allowed_domains"`
}

type YamlEgressAclConfiguration struct {
	Services []ServiceRule `yaml:"services"`
	Default  *ServiceRule  `yaml:"default"`
	Version  string        `yaml:"version"`
}

func LoadYamlAclFromFilePath(config *Config, aclPath string) (*EgressAclConfig, error) {
	file, err := os.Open(aclPath)

	if err != nil {
		return nil, err
	}
	defer file.Close()
	return LoadYamlAclFromReader(config, file)
}

func LoadYamlAclFromReader(config *Config, aclReader io.Reader) (*EgressAclConfig, error) {
	fail := func(err error) (*EgressAclConfig, error) { return nil, err }

	yamlConfig := YamlEgressAclConfiguration{}

	yamlFile, err := ioutil.ReadAll(aclReader)
	if err != nil {
		return fail(fmt.Errorf("could not load acl configuration"))
	}

	err = yaml.Unmarshal(yamlFile, &yamlConfig)

	return BuildAclFromYamlConfig(config, &yamlConfig)
}

func BuildAclFromYamlConfig(config *Config, yamlConfig *YamlEgressAclConfiguration) (*EgressAclConfig, error) {
	fail := func(err error) (*EgressAclConfig, error) { return nil, err }

	var err error

	if yamlConfig.Version != "v1" {
		return fail(fmt.Errorf("Expected version \"v1\" got %#v\n", yamlConfig.Version))
	}

	if err != nil {
		return nil, err
	}

	acl := EgressAclConfig{Services: make(map[string]EgressAclRule)}

	if yamlConfig.Services == nil {
		return nil, errors.New("Top level list 'services' is missing")
	}

	for _, v := range yamlConfig.Services {
		res, err := aclConfigToRule(&v, config.DisabledAclPolicyActions)
		if err != nil {
			config.Log.Error("gnored policy", err)
		} else {
			acl.Services[v.Name] = res
		}
	}

	if yamlConfig.Default != nil {
		res, err := aclConfigToRule(yamlConfig.Default, config.DisabledAclPolicyActions)
		if err != nil {
			config.Log.Error("gnored policy", err)
		} else {
			acl.Default = &res
		}
	}
	return &acl, nil
}

func aclConfigToRule(v *ServiceRule, disabledAclPolicyAction []string) (EgressAclRule, error) {
	var enforcementPolicy ConfigEnforcementPolicy

	// Validate policy action
	for _, disabledAction := range disabledAclPolicyAction {
		if disabledAction == v.Action {
			return EgressAclRule{}, fmt.Errorf("policy action %#v has been disabled but is used in rule for service %#v", disabledAction, v.Name)
		}
	}

	// Validate hosts

	for _, host := range v.AllowedHosts {
		if !strings.HasPrefix(host, "*.") && strings.HasPrefix(host, "*") {
			return EgressAclRule{}, fmt.Errorf("glob must represent a full prefxi (sub)domain")
		}

		// Check for stars elsewhere
		hostToCheck := host
		if strings.HasPrefix(hostToCheck, "*") {
			hostToCheck = hostToCheck[1:]
		}
		if strings.Contains(hostToCheck, "*") {
			return EgressAclRule{}, fmt.Errorf("globs are only supported as prefix")
		}
	}

	switch v.Action {
	case "open":
		enforcementPolicy = ConfigEnforcementPolicyOpen
	case "report":
		enforcementPolicy = ConfigEnforcementPolicyReport
	case "enforce":
		enforcementPolicy = ConfigEnforcementPolicyEnforce
	default:
		enforcementPolicy = 0
		return EgressAclRule{}, errors.New(fmt.Sprintf("Unknown action '%s' under '%s'.", v.Action, v.Name))
	}

	return EgressAclRule{
		Project:    v.Project,
		Policy:     enforcementPolicy,
		DomainGlob: v.AllowedHosts,
	}, nil
}
<commit_msg>Implement ValidateConfig on the dest yaml struct<commit_after>package smokescreen

import (
	"errors"
	"fmt"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"os"
	"strings"
)

type EgressAclRule struct {
	Project    string
	Policy     ConfigEnforcementPolicy
	DomainGlob []string
}

type EgressAclConfig struct {
	Services map[string]EgressAclRule
	Default  *EgressAclRule
}

func (ew *EgressAclConfig) Decide(fromService string, toHost string) (EgressAclDecision, error) {

	var (
		found   bool
		service EgressAclRule
	)
	if fromService == "" {
		found = false
	} else {
		service, found = ew.Services[fromService]
	}

	if !found && ew.Default != nil {
		found = true
		service = *ew.Default
	}

	if !found {
		return 0, fmt.Errorf("unknown role role=%#v", fromService)
	}

	if service.Policy == ConfigEnforcementPolicyOpen {
		return EgressAclDecisionAllow, nil
	}

	matches := false
	for _, host := range service.DomainGlob {
		if len(host) > 0 && host[0] == '*' {
			postfix := host[1:]
			if strings.HasSuffix(toHost, postfix) {
				matches = true
				break
			}
		} else {
			if host == toHost {
				matches = true
				break
			}
		}
	}

	switch service.Policy {
	case ConfigEnforcementPolicyReport:
		if matches {
			return EgressAclDecisionAllow, nil
		} else {
			return EgressAclDecisionAllowAndReport, nil
		}
	case ConfigEnforcementPolicyEnforce:
		if matches {
			return EgressAclDecisionAllow, nil
		} else {
			return EgressAclDecisionDeny, nil
		}
	default:
		return 0, errors.New("unexpected state")
	}
}

func (ew *EgressAclConfig) Project(fromService string) (string, error) {
	service, found := ew.Services[fromService]

	if found {
		return service.Project, nil
	}

	if ew.Default != nil {
		return ew.Default.Project, nil
	}

	return "", fmt.Errorf("warn: No known project: role=%#v\n", fromService)
}

// Configuration

type ServiceRule struct {
	Name         string   `yaml:"name"`
	Project      string   `yaml:"project"`
	Action       string   `yaml:"action"`
	AllowedHosts []string `yaml:"allowed_domains"`
}

type YamlEgressAclConfiguration struct {
	Services []ServiceRule `yaml:"services"`
	Default  *ServiceRule  `yaml:"default"`
	Version  string        `yaml:"version"`
}

func (yamlConf *YamlEgressAclConfiguration) ValidateConfig() error {
	return nil
}

func LoadYamlAclFromFilePath(config *Config, aclPath string) (*EgressAclConfig, error) {
	file, err := os.Open(aclPath)

	if err != nil {
		return nil, err
	}
	defer file.Close()
	return LoadYamlAclFromReader(config, file)
}

func LoadYamlAclFromReader(config *Config, aclReader io.Reader) (*EgressAclConfig, error) {
	fail := func(err error) (*EgressAclConfig, error) { return nil, err }

	yamlConfig := YamlEgressAclConfiguration{}

	yamlFile, err := ioutil.ReadAll(aclReader)
	if err != nil {
		return fail(fmt.Errorf("could not load acl configuration"))
	}

	err = yaml.Unmarshal(yamlFile, &yamlConfig)

	return BuildAclFromYamlConfig(config, &yamlConfig)
}

func BuildAclFromYamlConfig(config *Config, yamlConfig *YamlEgressAclConfiguration) (*EgressAclConfig, error) {
	fail := func(err error) (*EgressAclConfig, error) { return nil, err }

	var err error

	if yamlConfig.Version != "v1" {
		return fail(fmt.Errorf("Expected version \"v1\" got %#v\n", yamlConfig.Version))
	}

	if err != nil {
		return nil, err
	}

	acl := EgressAclConfig{Services: make(map[string]EgressAclRule)}

	if yamlConfig.Services == nil {
		return nil, errors.New("Top level list 'services' is missing")
	}

	for _, v := range yamlConfig.Services {
		res, err := aclConfigToRule(&v, config.DisabledAclPolicyActions)
		if err != nil {
			config.Log.Error("gnored policy", err)
		} else {
			acl.Services[v.Name] = res
		}
	}

	if yamlConfig.Default != nil {
		res, err := aclConfigToRule(yamlConfig.Default, config.DisabledAclPolicyActions)
		if err != nil {
			config.Log.Error("gnored policy", err)
		} else {
			acl.Default = &res
		}
	}
	return &acl, nil
}

func aclConfigToRule(v *ServiceRule, disabledAclPolicyAction []string) (EgressAclRule, error) {
	var enforcementPolicy ConfigEnforcementPolicy

	// Validate policy action
	for _, disabledAction := range disabledAclPolicyAction {
		if disabledAction == v.Action {
			return EgressAclRule{}, fmt.Errorf("policy action %#v has been disabled but is used in rule for service %#v", disabledAction, v.Name)
		}
	}

	// Validate hosts

	for _, host := range v.AllowedHosts {
		if !strings.HasPrefix(host, "*.") && strings.HasPrefix(host, "*") {
			return EgressAclRule{}, fmt.Errorf("glob must represent a full prefxi (sub)domain")
		}

		// Check for stars elsewhere
		hostToCheck := host
		if strings.HasPrefix(hostToCheck, "*") {
			hostToCheck = hostToCheck[1:]
		}
		if strings.Contains(hostToCheck, "*") {
			return EgressAclRule{}, fmt.Errorf("globs are only supported as prefix")
		}
	}

	switch v.Action {
	case "open":
		enforcementPolicy = ConfigEnforcementPolicyOpen
	case "report":
		enforcementPolicy = ConfigEnforcementPolicyReport
	case "enforce":
		enforcementPolicy = ConfigEnforcementPolicyEnforce
	default:
		enforcementPolicy = 0
		return EgressAclRule{}, errors.New(fmt.Sprintf("Unknown action '%s' under '%s'.", v.Action, v.Name))
	}

	return EgressAclRule{
		Project:    v.Project,
		Policy:     enforcementPolicy,
		DomainGlob: v.AllowedHosts,
	}, nil
}
<|endoftext|>


<commit_before>// Copyright 2020 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
//
// https://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 metric

import (
	"context"
	"log"
	"net"
	"os"
	"testing"

	"github.com/golang/protobuf/ptypes/empty"
	"google.golang.org/genproto/googleapis/monitoring/v3"
	metrics "google.golang.org/genproto/googleapis/api/metric"
	mockMetricServer "github.com/googleinterns/cloud-operations-api-mock/server/metric"
	"google.golang.org/grpc"
	"google.golang.org/grpc/test/bufconn"
	"google.golang.org/protobuf/proto"
)

const bufSize = 1024 * 1024

var (
	client monitoring.MetricServiceClient
	conn   *grpc.ClientConn
	ctx    context.Context
	lis    *bufconn.Listener
)

func setup() {
	// Setup the in-memory server.
	lis = bufconn.Listen(bufSize)
	grpcServer := grpc.NewServer()
	monitoring.RegisterMetricServiceServer(grpcServer, &mockMetricServer.MockMetricServer{})
	go func() {
		if err := grpcServer.Serve(lis); err != nil {
			log.Fatalf("server exited with error: %v", err)
		}
	}()

	// Setup the connection and client.
	ctx = context.Background()
	var err error
	conn, err = grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
	if err != nil {
		log.Fatalf("failed to dial bufnet: %v", err)
	}
	client = monitoring.NewMetricServiceClient(conn)
}

func tearDown() {
	conn.Close()
}

func TestMain(m *testing.M) {
	setup()
	retCode := m.Run()
	tearDown()
	os.Exit(retCode)
}

func bufDialer(context.Context, string) (net.Conn, error) {
	return lis.Dial()
}

func TestMockMetricServer_GetMetricDescriptor(t *testing.T) {
	cases := []struct {
		in *monitoring.GetMetricDescriptorRequest
		want *metrics.MetricDescriptor
	}{
		{
			&monitoring.GetMetricDescriptorRequest{
				Name: "test get metric descriptor",
			},
			&metrics.MetricDescriptor{},
		},
	}

	for _,c := range cases {
		response, err := client.GetMetricDescriptor(ctx, c.in)
		if err != nil {
			t.Fatalf("failed to call GetMetricDescriptor: %v", err)
		}

		if !proto.Equal(response, c.want) {
			t.Errorf("GetMetricDescriptor(%q) == %q, want %q", c.in, response, c.want)
		}
	}
}

func TestMockMetricServer_CreateMetricDescriptor(t *testing.T) {
	cases := []struct {
		in *monitoring.CreateMetricDescriptorRequest
		want *metrics.MetricDescriptor
	}{
		{
			&monitoring.CreateMetricDescriptorRequest{
				Name: "test create metric descriptor",
				MetricDescriptor: &metrics.MetricDescriptor{},
			},
			&metrics.MetricDescriptor{},
		},
	}

	for _,c := range cases {
		response, err := client.CreateMetricDescriptor(ctx, c.in)
		if err != nil {
			t.Fatalf("failed to call CreateMetricDescriptorRequest: %v", err)
		}

		if !proto.Equal(response, c.want) {
			t.Errorf("CreateMetricDescriptorRequest(%q) == %q, want %q", c.in, response, c.want)
		}
	}
}

func TestMockMetricServer_DeleteMetricDescriptor(t *testing.T) {
	cases := []struct {
		in *monitoring.DeleteMetricDescriptorRequest
		want *empty.Empty
	}{
		{
			&monitoring.DeleteMetricDescriptorRequest{
				Name: "test create metric descriptor",
			},
			&empty.Empty{},
		},
	}

	for _,c := range cases {
		response, err := client.DeleteMetricDescriptor(ctx, c.in)
		if err != nil {
			t.Fatalf("failed to call CDeleteMetricDescriptorRequest: %v", err)
		}

		if !proto.Equal(response, c.want) {
			t.Errorf("DeleteMetricDescriptorRequest(%q) == %q, want %q", c.in, response, c.want)
		}
	}
}<commit_msg>fix spelling/spacing nits, remove metrics package rename<commit_after>// Copyright 2020 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
//
// https://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 metric

import (
	"context"
	"log"
	"net"
	"os"
	"testing"

	"github.com/golang/protobuf/ptypes/empty"
	"google.golang.org/genproto/googleapis/monitoring/v3"
	"google.golang.org/genproto/googleapis/api/metric"
	"google.golang.org/grpc"
	"google.golang.org/grpc/test/bufconn"
	"google.golang.org/protobuf/proto"
)

const bufSize = 1024 * 1024

var (
	client monitoring.MetricServiceClient
	conn   *grpc.ClientConn
	ctx    context.Context
	lis    *bufconn.Listener
)

func setup() {
	// Setup the in-memory server.
	lis = bufconn.Listen(bufSize)
	grpcServer := grpc.NewServer()
	monitoring.RegisterMetricServiceServer(grpcServer, &MockMetricServer{})
	go func() {
		if err := grpcServer.Serve(lis); err != nil {
			log.Fatalf("server exited with error: %v", err)
		}
	}()

	// Setup the connection and client.
	ctx = context.Background()
	var err error
	conn, err = grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
	if err != nil {
		log.Fatalf("failed to dial bufnet: %v", err)
	}
	client = monitoring.NewMetricServiceClient(conn)
}

func tearDown() {
	conn.Close()
}

func TestMain(m *testing.M) {
	setup()
	retCode := m.Run()
	tearDown()
	os.Exit(retCode)
}

func bufDialer(context.Context, string) (net.Conn, error) {
	return lis.Dial()
}

func TestMockMetricServer_GetMetricDescriptor(t *testing.T) {
	cases := []struct {
		in *monitoring.GetMetricDescriptorRequest
		want *metric.MetricDescriptor
	}{
		{
			&monitoring.GetMetricDescriptorRequest{
				Name: "test get metric descriptor",
			},
			&metric.MetricDescriptor{},
		},
	}

	for _, c := range cases {
		response, err := client.GetMetricDescriptor(ctx, c.in)
		if err != nil {
			t.Fatalf("failed to call GetMetricDescriptor %v", err)
		}

		if !proto.Equal(response, c.want) {
			t.Errorf("GetMetricDescriptor(%q) == %q, want %q", c.in, response, c.want)
		}
	}
}

func TestMockMetricServer_CreateMetricDescriptor(t *testing.T) {
	cases := []struct {
		in *monitoring.CreateMetricDescriptorRequest
		want *metric.MetricDescriptor
	}{
		{
			&monitoring.CreateMetricDescriptorRequest{
				Name: "test create metric descriptor",
				MetricDescriptor: &metric.MetricDescriptor{},
			},
			&metric.MetricDescriptor{},
		},
	}

	for _, c := range cases {
		response, err := client.CreateMetricDescriptor(ctx, c.in)
		if err != nil {
			t.Fatalf("failed to call CreateMetricDescriptorRequest: %v", err)
		}

		if !proto.Equal(response, c.want) {
			t.Errorf("CreateMetricDescriptorRequest(%q) == %q, want %q", c.in, response, c.want)
		}
	}
}

func TestMockMetricServer_DeleteMetricDescriptor(t *testing.T) {
	cases := []struct {
		in *monitoring.DeleteMetricDescriptorRequest
		want *empty.Empty
	}{
		{
			&monitoring.DeleteMetricDescriptorRequest{
				Name: "test create metric descriptor",
			},
			&empty.Empty{},
		},
	}

	for _,c := range cases {
		response, err := client.DeleteMetricDescriptor(ctx, c.in)
		if err != nil {
			t.Fatalf("failed to call DeleteMetricDescriptorRequest: %v", err)
		}

		if !proto.Equal(response, c.want) {
			t.Errorf("DeleteMetricDescriptorRequest(%q) == %q, want %q", c.in, response, c.want)
		}
	}
}<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cinder

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/api/v1"
	"k8s.io/kubernetes/pkg/util/exec"
	"k8s.io/kubernetes/pkg/volume"
)

type CinderDiskUtil struct{}

// Attaches a disk specified by a volume.CinderPersistenDisk to the current kubelet.
// Mounts the disk to its global path.
func (util *CinderDiskUtil) AttachDisk(b *cinderVolumeMounter, globalPDPath string) error {
	options := []string{}
	if b.readOnly {
		options = append(options, "ro")
	}
	cloud, err := b.plugin.getCloudProvider()
	if err != nil {
		return err
	}
	instanceid, err := cloud.InstanceID()
	if err != nil {
		return err
	}
	diskid, err := cloud.AttachDisk(instanceid, b.pdName)
	if err != nil {
		return err
	}

	var devicePath string
	numTries := 0
	for {
		devicePath = cloud.GetDevicePath(diskid)
		probeAttachedVolume()

		_, err := os.Stat(devicePath)
		if err == nil {
			break
		}
		if err != nil && !os.IsNotExist(err) {
			return err
		}
		numTries++
		if numTries == 10 {
			return errors.New("Could not attach disk: Timeout after 60s")
		}
		time.Sleep(time.Second * 6)
	}
	notmnt, err := b.mounter.IsLikelyNotMountPoint(globalPDPath)
	if err != nil {
		if os.IsNotExist(err) {
			if err := os.MkdirAll(globalPDPath, 0750); err != nil {
				return err
			}
			notmnt = true
		} else {
			return err
		}
	}
	if notmnt {
		err = b.blockDeviceMounter.FormatAndMount(devicePath, globalPDPath, b.fsType, options)
		if err != nil {
			os.Remove(globalPDPath)
			return err
		}
		glog.V(2).Infof("Safe mount successful: %q\n", devicePath)
	}
	return nil
}

// Unmounts the device and detaches the disk from the kubelet's host machine.
func (util *CinderDiskUtil) DetachDisk(cd *cinderVolumeUnmounter) error {
	globalPDPath := makeGlobalPDName(cd.plugin.host, cd.pdName)
	if err := cd.mounter.Unmount(globalPDPath); err != nil {
		return err
	}
	if err := os.Remove(globalPDPath); err != nil {
		return err
	}
	glog.V(2).Infof("Successfully unmounted main device: %s\n", globalPDPath)

	cloud, err := cd.plugin.getCloudProvider()
	if err != nil {
		return err
	}
	instanceid, err := cloud.InstanceID()
	if err != nil {
		return err
	}
	if err = cloud.DetachDisk(instanceid, cd.pdName); err != nil {
		return err
	}
	glog.V(2).Infof("Successfully detached cinder volume %s", cd.pdName)
	return nil
}

func (util *CinderDiskUtil) DeleteVolume(cd *cinderVolumeDeleter) error {
	cloud, err := cd.plugin.getCloudProvider()
	if err != nil {
		return err
	}

	if err = cloud.DeleteVolume(cd.pdName); err != nil {
		// OpenStack cloud provider returns volume.tryAgainError when necessary,
		// no handling needed here.
		glog.V(2).Infof("Error deleting cinder volume %s: %v", cd.pdName, err)
		return err
	}
	glog.V(2).Infof("Successfully deleted cinder volume %s", cd.pdName)
	return nil
}

func (util *CinderDiskUtil) CreateVolume(c *cinderVolumeProvisioner) (volumeID string, volumeSizeGB int, err error) {
	cloud, err := c.plugin.getCloudProvider()
	if err != nil {
		return "", 0, err
	}

	capacity := c.options.PVC.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)]
	volSizeBytes := capacity.Value()
	// Cinder works with gigabytes, convert to GiB with rounding up
	volSizeGB := int(volume.RoundUpSize(volSizeBytes, 1024*1024*1024))
	name := volume.GenerateVolumeName(c.options.ClusterName, c.options.PVName, 255) // Cinder volume name can have up to 255 characters
	vtype := ""
	availability := ""
	// Apply ProvisionerParameters (case-insensitive). We leave validation of
	// the values to the cloud provider.
	for k, v := range c.options.Parameters {
		switch strings.ToLower(k) {
		case "type":
			vtype = v
		case "availability":
			availability = v
		default:
			return "", 0, fmt.Errorf("invalid option %q for volume plugin %s", k, c.plugin.GetPluginName())
		}
	}
	// TODO: implement PVC.Selector parsing
	if c.options.PVC.Spec.Selector != nil {
		return "", 0, fmt.Errorf("claim.Spec.Selector is not supported for dynamic provisioning on Cinder")
	}

	name, err = cloud.CreateVolume(name, volSizeGB, vtype, availability, c.options.CloudTags)
	if err != nil {
		glog.V(2).Infof("Error creating cinder volume: %v", err)
		return "", 0, err
	}
	glog.V(2).Infof("Successfully created cinder volume %s", name)
	return name, volSizeGB, nil
}

func probeAttachedVolume() error {
	// rescan scsi bus
	scsiHostRescan()

	executor := exec.New()
	args := []string{"trigger"}
	cmd := executor.Command("/usr/bin/udevadm", args...)
	_, err := cmd.CombinedOutput()
	if err != nil {
		glog.Errorf("error running udevadm trigger %v\n", err)
		return err
	}
	glog.V(4).Infof("Successfully probed all attachments")
	return nil
}

func scsiHostRescan() {
	scsi_path := "/sys/class/scsi_host/"
	if dirs, err := ioutil.ReadDir(scsi_path); err == nil {
		for _, f := range dirs {
			name := scsi_path + f.Name() + "/scan"
			data := []byte("- - -")
			ioutil.WriteFile(name, data, 0666)
		}
	}
}
<commit_msg>resolve udevadm from PATH<commit_after>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cinder

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/api/v1"
	"k8s.io/kubernetes/pkg/util/exec"
	"k8s.io/kubernetes/pkg/volume"
)

type CinderDiskUtil struct{}

// Attaches a disk specified by a volume.CinderPersistenDisk to the current kubelet.
// Mounts the disk to its global path.
func (util *CinderDiskUtil) AttachDisk(b *cinderVolumeMounter, globalPDPath string) error {
	options := []string{}
	if b.readOnly {
		options = append(options, "ro")
	}
	cloud, err := b.plugin.getCloudProvider()
	if err != nil {
		return err
	}
	instanceid, err := cloud.InstanceID()
	if err != nil {
		return err
	}
	diskid, err := cloud.AttachDisk(instanceid, b.pdName)
	if err != nil {
		return err
	}

	var devicePath string
	numTries := 0
	for {
		devicePath = cloud.GetDevicePath(diskid)
		probeAttachedVolume()

		_, err := os.Stat(devicePath)
		if err == nil {
			break
		}
		if err != nil && !os.IsNotExist(err) {
			return err
		}
		numTries++
		if numTries == 10 {
			return errors.New("Could not attach disk: Timeout after 60s")
		}
		time.Sleep(time.Second * 6)
	}
	notmnt, err := b.mounter.IsLikelyNotMountPoint(globalPDPath)
	if err != nil {
		if os.IsNotExist(err) {
			if err := os.MkdirAll(globalPDPath, 0750); err != nil {
				return err
			}
			notmnt = true
		} else {
			return err
		}
	}
	if notmnt {
		err = b.blockDeviceMounter.FormatAndMount(devicePath, globalPDPath, b.fsType, options)
		if err != nil {
			os.Remove(globalPDPath)
			return err
		}
		glog.V(2).Infof("Safe mount successful: %q\n", devicePath)
	}
	return nil
}

// Unmounts the device and detaches the disk from the kubelet's host machine.
func (util *CinderDiskUtil) DetachDisk(cd *cinderVolumeUnmounter) error {
	globalPDPath := makeGlobalPDName(cd.plugin.host, cd.pdName)
	if err := cd.mounter.Unmount(globalPDPath); err != nil {
		return err
	}
	if err := os.Remove(globalPDPath); err != nil {
		return err
	}
	glog.V(2).Infof("Successfully unmounted main device: %s\n", globalPDPath)

	cloud, err := cd.plugin.getCloudProvider()
	if err != nil {
		return err
	}
	instanceid, err := cloud.InstanceID()
	if err != nil {
		return err
	}
	if err = cloud.DetachDisk(instanceid, cd.pdName); err != nil {
		return err
	}
	glog.V(2).Infof("Successfully detached cinder volume %s", cd.pdName)
	return nil
}

func (util *CinderDiskUtil) DeleteVolume(cd *cinderVolumeDeleter) error {
	cloud, err := cd.plugin.getCloudProvider()
	if err != nil {
		return err
	}

	if err = cloud.DeleteVolume(cd.pdName); err != nil {
		// OpenStack cloud provider returns volume.tryAgainError when necessary,
		// no handling needed here.
		glog.V(2).Infof("Error deleting cinder volume %s: %v", cd.pdName, err)
		return err
	}
	glog.V(2).Infof("Successfully deleted cinder volume %s", cd.pdName)
	return nil
}

func (util *CinderDiskUtil) CreateVolume(c *cinderVolumeProvisioner) (volumeID string, volumeSizeGB int, err error) {
	cloud, err := c.plugin.getCloudProvider()
	if err != nil {
		return "", 0, err
	}

	capacity := c.options.PVC.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)]
	volSizeBytes := capacity.Value()
	// Cinder works with gigabytes, convert to GiB with rounding up
	volSizeGB := int(volume.RoundUpSize(volSizeBytes, 1024*1024*1024))
	name := volume.GenerateVolumeName(c.options.ClusterName, c.options.PVName, 255) // Cinder volume name can have up to 255 characters
	vtype := ""
	availability := ""
	// Apply ProvisionerParameters (case-insensitive). We leave validation of
	// the values to the cloud provider.
	for k, v := range c.options.Parameters {
		switch strings.ToLower(k) {
		case "type":
			vtype = v
		case "availability":
			availability = v
		default:
			return "", 0, fmt.Errorf("invalid option %q for volume plugin %s", k, c.plugin.GetPluginName())
		}
	}
	// TODO: implement PVC.Selector parsing
	if c.options.PVC.Spec.Selector != nil {
		return "", 0, fmt.Errorf("claim.Spec.Selector is not supported for dynamic provisioning on Cinder")
	}

	name, err = cloud.CreateVolume(name, volSizeGB, vtype, availability, c.options.CloudTags)
	if err != nil {
		glog.V(2).Infof("Error creating cinder volume: %v", err)
		return "", 0, err
	}
	glog.V(2).Infof("Successfully created cinder volume %s", name)
	return name, volSizeGB, nil
}

func probeAttachedVolume() error {
	// rescan scsi bus
	scsiHostRescan()

	executor := exec.New()
	args := []string{"trigger"}
	cmd := executor.Command("udevadm", args...)
	_, err := cmd.CombinedOutput()
	if err != nil {
		glog.Errorf("error running udevadm trigger %v\n", err)
		return err
	}
	glog.V(4).Infof("Successfully probed all attachments")
	return nil
}

func scsiHostRescan() {
	scsi_path := "/sys/class/scsi_host/"
	if dirs, err := ioutil.ReadDir(scsi_path); err == nil {
		for _, f := range dirs {
			name := scsi_path + f.Name() + "/scan"
			data := []byte("- - -")
			ioutil.WriteFile(name, data, 0666)
		}
	}
}
<|endoftext|>


<commit_before>package natsgw

import (
	"errors"
	"os/exec"
	"syscall"
	"testing"
	"time"

	"github.com/CyCoreSystems/ari"
	"github.com/CyCoreSystems/ari/client/mock"
	"github.com/CyCoreSystems/ari/client/nc"
	"github.com/golang/mock/gomock"
)

var ServerWaitDelay = 500 * time.Millisecond

func TestApplicationsSubscribeUnsubscribe(t *testing.T) {

	//TODO: embed nats?

	bin, err := exec.LookPath("gnatsd")
	if err != nil {
		t.Skip("No gnatsd binary in PATH, skipping")
	}

	cmd := exec.Command(bin, "-p", "4333")
	if err := cmd.Start(); err != nil {
		t.Errorf("Unable to run gnatsd: '%v'", err)
		return
	}

	defer func() {
		cmd.Process.Signal(syscall.SIGTERM)
	}()

	<-time.After(ServerWaitDelay)

	// test client

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockApplication := mock.NewMockApplication(ctrl)
	mockApplication.EXPECT().Subscribe("app1", "evt1").Return(nil)
	mockApplication.EXPECT().Subscribe("app2", "evt1").Return(errors.New("Application not found"))
	mockApplication.EXPECT().Unsubscribe("app1", "evt1").Return(nil)
	mockApplication.EXPECT().Unsubscribe("app2", "evt1").Return(errors.New("Application not found"))

	cl := &ari.Client{
		Application: mockApplication,
	}
	s, err := NewServer(cl, &Options{
		URL: "nats://127.0.0.1:4333",
	})

	failed := s == nil || err != nil
	if failed {
		t.Errorf("natsgw.NewServer(cl, nil) => {%v, %v}, expected {%v, %v}", s, err, "cl", "nil")
	}

	go s.Listen()

	natsClient, err := nc.New("nats://127.0.0.1:4333")

	failed = natsClient == nil || err != nil
	if failed {
		t.Errorf("nc.New(url) => {%v, %v}, expected {%v, %v}", natsClient, err, "cl", "nil")
	}

	{
		err = natsClient.Application.Subscribe("app1", "evt1")

		failed = err != nil
		if failed {
			t.Errorf("nc.Application.Subscribe(app1,evt1) => '%v', expected '%v'", err, "nil")
		}
	}

	{
		err = natsClient.Application.Subscribe("app2", "evt1")

		failed = err == nil || err.Error() != "Application not found"
		if failed {
			t.Errorf("nc.Application.Subscribe(app2,evt1) => '%v', expected '%v'", err, "nil")
		}
	}

	{
		err = natsClient.Application.Unsubscribe("app1", "evt1")

		failed = err != nil
		if failed {
			t.Errorf("nc.Application.Unsubscribe(app1,evt1) => '%v', expected '%v'", err, "nil")
		}
	}

	{
		err = natsClient.Application.Unsubscribe("app2", "evt1")

		failed = err == nil || err.Error() != "Application not found"
		if failed {
			t.Errorf("nc.Application.Unsubscribe(app2,evt1) => '%v', expected '%v'", err, "nil")
		}
	}

	s.Close()
}
func TestApplicationsData(t *testing.T) {

	//TODO: embed nats?

	bin, err := exec.LookPath("gnatsd")
	if err != nil {
		t.Skip("No gnatsd binary in PATH, skipping")
	}

	cmd := exec.Command(bin, "-p", "4333")
	if err := cmd.Start(); err != nil {
		t.Errorf("Unable to run gnatsd: '%v'", err)
		return
	}

	defer func() {
		cmd.Process.Signal(syscall.SIGTERM)
	}()

	<-time.After(ServerWaitDelay)

	// test clientiontruc

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var appData ari.ApplicationData

	mockApplication := mock.NewMockApplication(ctrl)
	mockApplication.EXPECT().Data("app1").Return(appData, nil)
	mockApplication.EXPECT().Data("app2").Return(appData, errors.New("Application not found"))

	cl := &ari.Client{
		Application: mockApplication,
	}
	s, err := NewServer(cl, &Options{
		URL: "nats://127.0.0.1:4333",
	})

	failed := s == nil || err != nil
	if failed {
		t.Errorf("natsgw.NewServer(cl, nil) => {%v, %v}, expected {%v, %v}", s, err, "cl", "nil")
	}

	go s.Listen()

	natsClient, err := nc.New("nats://127.0.0.1:4333")

	failed = natsClient == nil || err != nil
	if failed {
		t.Errorf("nc.New(url) => {%v, %v}, expected {%v, %v}", natsClient, err, "cl", "nil")
	}

	{
		appData, err := natsClient.Application.Data("app1")

		failed = err != nil
		if failed {
			t.Errorf("nc.Application.Data(app1) => '%v', '%v', expected '%v', '%v'", appData, err, "appData", nil)
		}
	}

	{
		appData, err := natsClient.Application.Data("app2")

		failed = err == nil || err.Error() != "Application not found"
		if failed {
			t.Errorf("nc.Application.Data(app2) => '%v', '%v', expected '%v', '%v'", appData, err, "appData", "Application not found")
		}
	}

	s.Close()
}

func TestApplicationsList(t *testing.T) {

	//TODO: embed nats?

	bin, err := exec.LookPath("gnatsd")
	if err != nil {
		t.Skip("No gnatsd binary in PATH, skipping")
	}

	cmd := exec.Command(bin, "-p", "4333")
	if err := cmd.Start(); err != nil {
		t.Errorf("Unable to run gnatsd: '%v'", err)
		return
	}

	defer func() {
		cmd.Process.Signal(syscall.SIGTERM)
	}()

	<-time.After(ServerWaitDelay)

	// test clientiontruc

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockApplication := mock.NewMockApplication(ctrl)
	mockApplication.EXPECT().List().Return([]*ari.ApplicationHandle{ari.NewApplicationHandle("app1", mockApplication), ari.NewApplicationHandle("app2", mockApplication)}, nil)

	cl := &ari.Client{
		Application: mockApplication,
	}
	s, err := NewServer(cl, &Options{
		URL: "nats://127.0.0.1:4333",
	})

	failed := s == nil || err != nil
	if failed {
		t.Errorf("natsgw.NewServer(cl, nil) => {%v, %v}, expected {%v, %v}", s, err, "cl", "nil")
	}

	go s.Listen()

	natsClient, err := nc.New("nats://127.0.0.1:4333")

	failed = natsClient == nil || err != nil
	if failed {
		t.Errorf("nc.New(url) => {%v, %v}, expected {%v, %v}", natsClient, err, "cl", "nil")
	}

	apps, err := natsClient.Application.List()

	failed = len(apps) != 2 || err != nil
	if failed {
		t.Errorf("nc.Application.List() => {%v, %v}, expected {%v, %v}", apps, err, "[app1,app2]", "nil")
	}

	s.Close()
}

func TestApplicationsListError(t *testing.T) {

	//TODO: embed nats?

	bin, err := exec.LookPath("gnatsd")
	if err != nil {
		t.Skip("No gnatsd binary in PATH, skipping")
	}

	cmd := exec.Command(bin, "-p", "4333")
	if err := cmd.Start(); err != nil {
		t.Errorf("Unable to run gnatsd: '%v'", err)
		return
	}

	defer func() {
		cmd.Process.Signal(syscall.SIGTERM)
	}()

	<-time.After(ServerWaitDelay)

	// test client

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockApplication := mock.NewMockApplication(ctrl)
	mockApplication.EXPECT().List().Return(nil, errors.New("Error getting application list"))

	cl := &ari.Client{
		Application: mockApplication,
	}
	s, err := NewServer(cl, &Options{
		URL: "nats://127.0.0.1:4333",
	})

	failed := s == nil || err != nil
	if failed {
		t.Errorf("natsgw.NewServer(cl, nil) => {%v, %v}, expected {%v, %v}", s, err, "cl", "nil")
	}

	go s.Listen()

	natsClient, err := nc.New("nats://127.0.0.1:4333")

	failed = natsClient == nil || err != nil
	if failed {
		t.Errorf("nc.New(url) => {%v, %v}, expected {%v, %v}", natsClient, err, "cl", "nil")
	}

	apps, err := natsClient.Application.List()

	failed = len(apps) != 0 || err == nil
	if failed {
		t.Errorf("nc.Application.List() => {%v, %v}, expected {%v, %v}", apps, err, "[]", "err")
	}

	s.Close()
}
<commit_msg>v3 - natsgw - Move test server close to defer<commit_after>package natsgw

import (
	"errors"
	"os/exec"
	"syscall"
	"testing"
	"time"

	"github.com/CyCoreSystems/ari"
	"github.com/CyCoreSystems/ari/client/mock"
	"github.com/CyCoreSystems/ari/client/nc"
	"github.com/golang/mock/gomock"
)

var ServerWaitDelay = 500 * time.Millisecond

func TestApplicationsSubscribeUnsubscribe(t *testing.T) {

	//TODO: embed nats?

	bin, err := exec.LookPath("gnatsd")
	if err != nil {
		t.Skip("No gnatsd binary in PATH, skipping")
	}

	cmd := exec.Command(bin, "-p", "4333")
	if err := cmd.Start(); err != nil {
		t.Errorf("Unable to run gnatsd: '%v'", err)
		return
	}

	defer func() {
		cmd.Process.Signal(syscall.SIGTERM)
	}()

	<-time.After(ServerWaitDelay)

	// test client

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockApplication := mock.NewMockApplication(ctrl)
	mockApplication.EXPECT().Subscribe("app1", "evt1").Return(nil)
	mockApplication.EXPECT().Subscribe("app2", "evt1").Return(errors.New("Application not found"))
	mockApplication.EXPECT().Unsubscribe("app1", "evt1").Return(nil)
	mockApplication.EXPECT().Unsubscribe("app2", "evt1").Return(errors.New("Application not found"))

	cl := &ari.Client{
		Application: mockApplication,
	}
	s, err := NewServer(cl, &Options{
		URL: "nats://127.0.0.1:4333",
	})

	failed := s == nil || err != nil
	if failed {
		t.Errorf("natsgw.NewServer(cl, nil) => {%v, %v}, expected {%v, %v}", s, err, "cl", "nil")
	}

	go s.Listen()
	defer s.Close()

	natsClient, err := nc.New("nats://127.0.0.1:4333")

	failed = natsClient == nil || err != nil
	if failed {
		t.Errorf("nc.New(url) => {%v, %v}, expected {%v, %v}", natsClient, err, "cl", "nil")
	}

	{
		err = natsClient.Application.Subscribe("app1", "evt1")

		failed = err != nil
		if failed {
			t.Errorf("nc.Application.Subscribe(app1,evt1) => '%v', expected '%v'", err, "nil")
		}
	}

	{
		err = natsClient.Application.Subscribe("app2", "evt1")

		failed = err == nil || err.Error() != "Application not found"
		if failed {
			t.Errorf("nc.Application.Subscribe(app2,evt1) => '%v', expected '%v'", err, "nil")
		}
	}

	{
		err = natsClient.Application.Unsubscribe("app1", "evt1")

		failed = err != nil
		if failed {
			t.Errorf("nc.Application.Unsubscribe(app1,evt1) => '%v', expected '%v'", err, "nil")
		}
	}

	{
		err = natsClient.Application.Unsubscribe("app2", "evt1")

		failed = err == nil || err.Error() != "Application not found"
		if failed {
			t.Errorf("nc.Application.Unsubscribe(app2,evt1) => '%v', expected '%v'", err, "nil")
		}
	}

}
func TestApplicationsData(t *testing.T) {

	//TODO: embed nats?

	bin, err := exec.LookPath("gnatsd")
	if err != nil {
		t.Skip("No gnatsd binary in PATH, skipping")
	}

	cmd := exec.Command(bin, "-p", "4333")
	if err := cmd.Start(); err != nil {
		t.Errorf("Unable to run gnatsd: '%v'", err)
		return
	}

	defer func() {
		cmd.Process.Signal(syscall.SIGTERM)
	}()

	<-time.After(ServerWaitDelay)

	// test clientiontruc

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var appData ari.ApplicationData

	mockApplication := mock.NewMockApplication(ctrl)
	mockApplication.EXPECT().Data("app1").Return(appData, nil)
	mockApplication.EXPECT().Data("app2").Return(appData, errors.New("Application not found"))

	cl := &ari.Client{
		Application: mockApplication,
	}
	s, err := NewServer(cl, &Options{
		URL: "nats://127.0.0.1:4333",
	})

	failed := s == nil || err != nil
	if failed {
		t.Errorf("natsgw.NewServer(cl, nil) => {%v, %v}, expected {%v, %v}", s, err, "cl", "nil")
	}

	go s.Listen()
	defer s.Close()

	natsClient, err := nc.New("nats://127.0.0.1:4333")

	failed = natsClient == nil || err != nil
	if failed {
		t.Errorf("nc.New(url) => {%v, %v}, expected {%v, %v}", natsClient, err, "cl", "nil")
	}

	{
		appData, err := natsClient.Application.Data("app1")

		failed = err != nil
		if failed {
			t.Errorf("nc.Application.Data(app1) => '%v', '%v', expected '%v', '%v'", appData, err, "appData", nil)
		}
	}

	{
		appData, err := natsClient.Application.Data("app2")

		failed = err == nil || err.Error() != "Application not found"
		if failed {
			t.Errorf("nc.Application.Data(app2) => '%v', '%v', expected '%v', '%v'", appData, err, "appData", "Application not found")
		}
	}

}

func TestApplicationsList(t *testing.T) {

	//TODO: embed nats?

	bin, err := exec.LookPath("gnatsd")
	if err != nil {
		t.Skip("No gnatsd binary in PATH, skipping")
	}

	cmd := exec.Command(bin, "-p", "4333")
	if err := cmd.Start(); err != nil {
		t.Errorf("Unable to run gnatsd: '%v'", err)
		return
	}

	defer func() {
		cmd.Process.Signal(syscall.SIGTERM)
	}()

	<-time.After(ServerWaitDelay)

	// test clientiontruc

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockApplication := mock.NewMockApplication(ctrl)
	mockApplication.EXPECT().List().Return([]*ari.ApplicationHandle{ari.NewApplicationHandle("app1", mockApplication), ari.NewApplicationHandle("app2", mockApplication)}, nil)

	cl := &ari.Client{
		Application: mockApplication,
	}
	s, err := NewServer(cl, &Options{
		URL: "nats://127.0.0.1:4333",
	})

	failed := s == nil || err != nil
	if failed {
		t.Errorf("natsgw.NewServer(cl, nil) => {%v, %v}, expected {%v, %v}", s, err, "cl", "nil")
	}

	go s.Listen()
	defer s.Close()

	natsClient, err := nc.New("nats://127.0.0.1:4333")

	failed = natsClient == nil || err != nil
	if failed {
		t.Errorf("nc.New(url) => {%v, %v}, expected {%v, %v}", natsClient, err, "cl", "nil")
	}

	apps, err := natsClient.Application.List()

	failed = len(apps) != 2 || err != nil
	if failed {
		t.Errorf("nc.Application.List() => {%v, %v}, expected {%v, %v}", apps, err, "[app1,app2]", "nil")
	}

}

func TestApplicationsListError(t *testing.T) {

	//TODO: embed nats?

	bin, err := exec.LookPath("gnatsd")
	if err != nil {
		t.Skip("No gnatsd binary in PATH, skipping")
	}

	cmd := exec.Command(bin, "-p", "4333")
	if err := cmd.Start(); err != nil {
		t.Errorf("Unable to run gnatsd: '%v'", err)
		return
	}

	defer func() {
		cmd.Process.Signal(syscall.SIGTERM)
	}()

	<-time.After(ServerWaitDelay)

	// test client

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockApplication := mock.NewMockApplication(ctrl)
	mockApplication.EXPECT().List().Return(nil, errors.New("Error getting application list"))

	cl := &ari.Client{
		Application: mockApplication,
	}
	s, err := NewServer(cl, &Options{
		URL: "nats://127.0.0.1:4333",
	})

	failed := s == nil || err != nil
	if failed {
		t.Errorf("natsgw.NewServer(cl, nil) => {%v, %v}, expected {%v, %v}", s, err, "cl", "nil")
	}

	go s.Listen()
	defer s.Close()

	natsClient, err := nc.New("nats://127.0.0.1:4333")

	failed = natsClient == nil || err != nil
	if failed {
		t.Errorf("nc.New(url) => {%v, %v}, expected {%v, %v}", natsClient, err, "cl", "nil")
	}

	apps, err := natsClient.Application.List()

	failed = len(apps) != 0 || err == nil
	if failed {
		t.Errorf("nc.Application.List() => {%v, %v}, expected {%v, %v}", apps, err, "[]", "err")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"github.com/codegangsta/cli"
	kval "github.com/kval-access-language/kval-boltdb"
	"github.com/olekukonko/tablewriter"
	"log"
	"os"
	"strings"
)

var instructionLine = "> Enter bucket to explore (CTRL-Q to quit, CTRL-B to go back, ENTER to reset):\n\n"

func main() {
	var file string

	cli.AppHelpTemplate = `NAME:
  {{.Name}} - {{.Usage}}

VERSION:
  {{.Version}}

USAGE:
  {{.HelpName}} {{if .VisibleFlags}}[global options]{{end}}

GLOBAL OPTIONS:
  {{range .VisibleFlags}}{{.}}
  {{end}}
AUTHOR:
  {{range .Authors}}{{ . }}{{end}}
COPYRIGHT:
  {{.Copyright}}
`
	app := cli.NewApp()
	app.Name = "bolter"
	app.Usage = "view boltdb file interactively in your terminal"
	app.Version = "2.0.0-kval-fork"
	app.Author = "Originally by Hasit Mistry. Interactive mode: Ross Spencer"
	app.Copyright = "(c) 2016 Hasit Mistry"
	//app.Email = ""
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:        "file, f",
			Usage:       "boltdb `FILE` to view",
			Destination: &file,
		},
	}
	app.Action = func(c *cli.Context) error {
		if file == "" {
			cli.ShowAppHelp(c)
			return nil
		}

		var i impl
		i = impl{fmt: &tableFormatter{}}
		if _, err := os.Stat(file); os.IsNotExist(err) {
			log.Fatal(err)
			return err
		}
		i.initDB(file)
		defer kval.Disconnect(i.kb)

		i.readInput()

		return nil
	}
	app.Run(os.Args)
}

func (i *impl) readInput() {
	i.listBuckets()
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		bucket := scanner.Text()
		fmt.Fprintln(os.Stdout, "")
		switch bucket {
		case "\x11":
			return
		case "\x02":
			// TODO: Change KVAL to get first record...
			if !strings.Contains(i.loc, "") || !strings.Contains(i.loc, ">>") {
				fmt.Println("> Going back...")
				i.loc = ""
				i.listBuckets()
			} else {
				i.listBucketItems(bucket, true)
			}
		case "":
			i.listBuckets()
		default:
			i.listBucketItems(bucket, false)
		}
		bucket = ""
	}
}

type formatter interface {
	DumpBuckets([]bucket)
	DumpBucketItems(string, []item)
}

type impl struct {
	kb     kval.Kvalboltdb
	fmt    formatter
	bucket string
	loc    string // where we are in the structure
	cur    string
	root   bool // are we on root
}

type item struct {
	Key   string
	Value string
}

type bucket struct {
	Name string
}

func (i *impl) initDB(file string) {
	var err error
	// Connect to KVAL
	i.kb, err = kval.Connect(file)
	if err != nil {
		log.Fatal(err)
	}
}

func (i *impl) updateLoc(bucket string, goBack bool) string {

	// we've probably an invalid value and want to display
	// ourselves again...
	if bucket == i.cur {
		i.loc = bucket
		return i.loc
	}

	// handle goback
	if goBack {
		s := strings.Split(i.loc, ">>")
		i.loc = strings.Join(s[:len(s)-1], ">>")
		i.bucket = strings.Trim(s[len(s)-2], " ")
		return i.loc
	}

	// handle location on merit...
	if i.loc == "" {
		i.loc = bucket
		i.bucket = bucket
	} else {
		i.loc = i.loc + " >> " + bucket
		i.bucket = bucket
	}
	return i.loc
}

func (i *impl) listBucketItems(bucket string, goBack bool) {
	items := []item{}
	getItems := i.updateLoc(bucket, goBack)
	fmt.Fprintf(os.Stderr, "Query: "+getItems+"\n\n")
	if getItems != "" {
		res, err := kval.Query(i.kb, "GET "+getItems)
		if err != nil {
			if err.Error() != "Cannot GOTO bucket, bucket not found" {
				log.Fatal(err)
			} else {
				fmt.Fprintf(os.Stderr, "> Bucket not found\n")
				if i.root == true {
					i.listBuckets()
					return
				}
				i.listBucketItems(i.loc, true)
			}
		}
		if len(res.Result) == 0 {
			fmt.Fprintf(os.Stderr, "Invalid request.\n\n")
			i.listBucketItems(i.cur, false)
			return
		}

		for k, v := range res.Result {
			if v == kval.Nestedbucket {
				k = k + "*"
				v = ""
			}
			items = append(items, item{Key: string(k), Value: string(v)})
		}
		i.fmt.DumpBucketItems(i.bucket, items)
		i.root = false
		i.cur = getItems
		fmt.Fprint(os.Stdout, instructionLine)
	}
}

func (i *impl) listBuckets() {
	i.root = true
	i.loc = ""

	buckets := []bucket{}

	res, err := kval.Query(i.kb, "GET _")
	if err != nil {
		log.Fatal(err)
	}
	for k, _ := range res.Result {
		buckets = append(buckets, bucket{Name: string(k) + "*"})
	}

	fmt.Fprint(os.Stdout, "DB Layout:\n\n")
	i.fmt.DumpBuckets(buckets)
	fmt.Fprint(os.Stdout, instructionLine)
}

type tableFormatter struct{}

func (tf tableFormatter) DumpBuckets(buckets []bucket) {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"Buckets"})
	for _, b := range buckets {
		row := []string{b.Name}
		table.Append(row)
	}
	table.Render()
	fmt.Println()
}

func (tf tableFormatter) DumpBucketItems(bucket string, items []item) {
	fmt.Printf("Bucket: %s\n", bucket)
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"Key", "Value"})
	for _, item := range items {
		row := []string{item.Key, item.Value}
		table.Append(row)
	}
	table.Render()
	fmt.Println()
}
<commit_msg>[bugfix] ctrl-q didn't work on linux. ctrl-x instead.<commit_after>package main

import (
	"bufio"
	"fmt"
	"github.com/codegangsta/cli"
	kval "github.com/kval-access-language/kval-boltdb"
	"github.com/olekukonko/tablewriter"
	"log"
	"os"
	"strings"
)

var instructionLine = "> Enter bucket to explore (CTRL-X to quit, CTRL-B to go back, ENTER to reset):\n\n"

func main() {
	var file string

	cli.AppHelpTemplate = `NAME:
  {{.Name}} - {{.Usage}}

VERSION:
  {{.Version}}

USAGE:
  {{.HelpName}} {{if .VisibleFlags}}[global options]{{end}}

GLOBAL OPTIONS:
  {{range .VisibleFlags}}{{.}}
  {{end}}
AUTHOR:
  {{range .Authors}}{{ . }}{{end}}
COPYRIGHT:
  {{.Copyright}}
`
	app := cli.NewApp()
	app.Name = "bolter"
	app.Usage = "view boltdb file interactively in your terminal"
	app.Version = "2.0.0-kval-fork"
	app.Author = "Originally by Hasit Mistry. Interactive mode: Ross Spencer"
	app.Copyright = "(c) 2016 Hasit Mistry"
	//app.Email = ""
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:        "file, f",
			Usage:       "boltdb `FILE` to view",
			Destination: &file,
		},
	}
	app.Action = func(c *cli.Context) error {
		if file == "" {
			cli.ShowAppHelp(c)
			return nil
		}

		var i impl
		i = impl{fmt: &tableFormatter{}}
		if _, err := os.Stat(file); os.IsNotExist(err) {
			log.Fatal(err)
			return err
		}
		i.initDB(file)
		defer kval.Disconnect(i.kb)

		i.readInput()

		return nil
	}
	app.Run(os.Args)
}

func (i *impl) readInput() {
	i.listBuckets()
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		bucket := scanner.Text()
		fmt.Fprintln(os.Stdout, "")
		switch bucket {
		case "\x18":
			return
		case "\x02":
			// TODO: Change KVAL to get first record...
			if !strings.Contains(i.loc, "") || !strings.Contains(i.loc, ">>") {
				fmt.Println("> Going back...")
				i.loc = ""
				i.listBuckets()
			} else {
				i.listBucketItems(bucket, true)
			}
		case "":
			i.listBuckets()
		default:
			i.listBucketItems(bucket, false)
		}
		bucket = ""
	}
}

type formatter interface {
	DumpBuckets([]bucket)
	DumpBucketItems(string, []item)
}

type impl struct {
	kb     kval.Kvalboltdb
	fmt    formatter
	bucket string
	loc    string // where we are in the structure
	cur    string
	root   bool // are we on root
}

type item struct {
	Key   string
	Value string
}

type bucket struct {
	Name string
}

func (i *impl) initDB(file string) {
	var err error
	// Connect to KVAL
	i.kb, err = kval.Connect(file)
	if err != nil {
		log.Fatal(err)
	}
}

func (i *impl) updateLoc(bucket string, goBack bool) string {

	// we've probably an invalid value and want to display
	// ourselves again...
	if bucket == i.cur {
		i.loc = bucket
		return i.loc
	}

	// handle goback
	if goBack {
		s := strings.Split(i.loc, ">>")
		i.loc = strings.Join(s[:len(s)-1], ">>")
		i.bucket = strings.Trim(s[len(s)-2], " ")
		return i.loc
	}

	// handle location on merit...
	if i.loc == "" {
		i.loc = bucket
		i.bucket = bucket
	} else {
		i.loc = i.loc + " >> " + bucket
		i.bucket = bucket
	}
	return i.loc
}

func (i *impl) listBucketItems(bucket string, goBack bool) {
	items := []item{}
	getItems := i.updateLoc(bucket, goBack)
	fmt.Fprintf(os.Stderr, "Query: "+getItems+"\n\n")
	if getItems != "" {
		res, err := kval.Query(i.kb, "GET "+getItems)
		if err != nil {
			if err.Error() != "Cannot GOTO bucket, bucket not found" {
				log.Fatal(err)
			} else {
				fmt.Fprintf(os.Stderr, "> Bucket not found\n")
				if i.root == true {
					i.listBuckets()
					return
				}
				i.listBucketItems(i.loc, true)
			}
		}
		if len(res.Result) == 0 {
			fmt.Fprintf(os.Stderr, "Invalid request.\n\n")
			i.listBucketItems(i.cur, false)
			return
		}

		for k, v := range res.Result {
			if v == kval.Nestedbucket {
				k = k + "*"
				v = ""
			}
			items = append(items, item{Key: string(k), Value: string(v)})
		}
		i.fmt.DumpBucketItems(i.bucket, items)
		i.root = false
		i.cur = getItems
		fmt.Fprint(os.Stdout, instructionLine)
	}
}

func (i *impl) listBuckets() {
	i.root = true
	i.loc = ""

	buckets := []bucket{}

	res, err := kval.Query(i.kb, "GET _")
	if err != nil {
		log.Fatal(err)
	}
	for k, _ := range res.Result {
		buckets = append(buckets, bucket{Name: string(k) + "*"})
	}

	fmt.Fprint(os.Stdout, "DB Layout:\n\n")
	i.fmt.DumpBuckets(buckets)
	fmt.Fprint(os.Stdout, instructionLine)
}

type tableFormatter struct{}

func (tf tableFormatter) DumpBuckets(buckets []bucket) {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"Buckets"})
	for _, b := range buckets {
		row := []string{b.Name}
		table.Append(row)
	}
	table.Render()
	fmt.Println()
}

func (tf tableFormatter) DumpBucketItems(bucket string, items []item) {
	fmt.Printf("Bucket: %s\n", bucket)
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"Key", "Value"})
	for _, item := range items {
		row := []string{item.Key, item.Value}
		table.Append(row)
	}
	table.Render()
	fmt.Println()
}
<|endoftext|>


<commit_before>package smux

import (
	"encoding/binary"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	siasync "github.com/NebulousLabs/Sia/sync" // TODO: Replace with github.com/NebulousLabs/trymutex
	"github.com/pkg/errors"
)

const (
	defaultAcceptBacklog = 1024
)

var (
	errBrokenPipe      = errors.New("broken pipe")
	errGoAway          = errors.New("stream id overflows, should start a new connection")
	errInvalidProtocol = errors.New("invalid protocol version")
	errLargeFrame      = errors.New("frame is too large to send")
)

// Session defines a multiplexed connection for streams
type Session struct {
	conn        net.Conn
	dataWasRead int32            // used to determine if KeepAlive has failed
	sendMu      siasync.TryMutex // ensures only one thread sends at a time

	config           *Config
	nextStreamID     uint32 // next stream identifier
	nextStreamIDLock sync.Mutex

	bucket       int32         // token bucket
	bucketNotify chan struct{} // used for waiting for tokens

	streams    map[uint32]*Stream // all streams in this session
	streamLock sync.Mutex         // locks streams

	die       chan struct{} // flag session has died
	dieLock   sync.Mutex
	chAccepts chan *Stream

	goAway int32 // flag id exhausted

	deadline atomic.Value
}

func newSession(config *Config, conn net.Conn, client bool) *Session {
	s := new(Session)
	s.die = make(chan struct{})
	s.conn = conn
	s.config = config
	s.streams = make(map[uint32]*Stream)
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.bucket = int32(config.MaxReceiveBuffer)
	s.bucketNotify = make(chan struct{}, 1)

	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 0
	}

	go s.recvLoop()
	// keepaliveSend and keepaliveTimeout need to be separate threads, because
	// the keepaliveSend can block, and especially if the underlying conn has no
	// deadline or a very long deadline, we may not check the keepaliveTimeout
	// for an extended period of time and potentially even end in a deadlock.
	go s.keepAliveSend()
	go s.keepAliveTimeout()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, errBrokenPipe
	}

	// generate stream id
	s.nextStreamIDLock.Lock()
	if s.goAway > 0 {
		s.nextStreamIDLock.Unlock()
		return nil, errGoAway
	}

	s.nextStreamID += 2
	sid := s.nextStreamID
	if sid == sid%2 { // stream-id overflows
		s.goAway = 1
		s.nextStreamIDLock.Unlock()
		return nil, errGoAway
	}
	s.nextStreamIDLock.Unlock()

	stream := newStream(sid, s.config.MaxFrameSize, s)

	if _, err := s.writeFrame(newFrame(cmdSYN, sid), time.Now().Add(s.config.WriteTimeout)); err != nil {
		return nil, errors.Wrap(err, "writeFrame")
	}

	s.streamLock.Lock()
	s.streams[sid] = stream
	s.streamLock.Unlock()
	return stream, nil
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	var deadline <-chan time.Time
	if d, ok := s.deadline.Load().(time.Time); ok && !d.IsZero() {
		timer := time.NewTimer(d.Sub(time.Now()))
		defer timer.Stop()
		deadline = timer.C
	}
	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-deadline:
		return nil, errTimeout
	case <-s.die:
		return nil, errBrokenPipe
	}
}

// Close is used to close the session and all streams.
func (s *Session) Close() (err error) {
	s.dieLock.Lock()

	select {
	case <-s.die:
		s.dieLock.Unlock()
		return errBrokenPipe
	default:
		close(s.die)
		s.dieLock.Unlock()
		s.streamLock.Lock()
		for k := range s.streams {
			s.streams[k].sessionClose()
		}
		s.streamLock.Unlock()
		s.notifyBucket()
		return s.conn.Close()
	}
}

// notifyBucket notifies recvLoop that bucket is available
func (s *Session) notifyBucket() {
	select {
	case s.bucketNotify <- struct{}{}:
	default:
	}
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	if s.IsClosed() {
		return 0
	}
	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	return len(s.streams)
}

// SetDeadline sets a deadline used by Accept* calls.
// A zero time value disables the deadline.
func (s *Session) SetDeadline(t time.Time) error {
	s.deadline.Store(t)
	return nil
}

// notify the session that a stream has closed
func (s *Session) streamClosed(sid uint32) {
	s.streamLock.Lock()
	if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket
		if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
			s.notifyBucket()
		}
	}
	delete(s.streams, sid)
	s.streamLock.Unlock()
}

// returnTokens is called by stream to return token after read
func (s *Session) returnTokens(n int) {
	if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
		s.notifyBucket()
	}
}

// session read a frame from underlying connection
// it's data is pointed to the input buffer
func (s *Session) readFrame(buffer []byte) (f Frame, err error) {
	// Ensure that reading a frame follows the global timeout.
	s.conn.SetReadDeadline(time.Now().Add(s.config.ReadTimeout))
	defer s.conn.SetReadDeadline(time.Time{})

	if _, err := io.ReadFull(s.conn, buffer[:headerSize]); err != nil {
		return f, errors.Wrap(err, "readFrame")
	}

	dec := rawHeader(buffer)
	if dec.Version() != version {
		return f, errInvalidProtocol
	}

	f.ver = dec.Version()
	f.cmd = dec.Cmd()
	f.sid = dec.StreamID()
	if length := dec.Length(); length > 0 {
		if _, err := io.ReadFull(s.conn, buffer[headerSize:headerSize+length]); err != nil {
			return f, errors.Wrap(err, "readFrame")
		}
		f.data = buffer[headerSize : headerSize+length]
	}
	return f, nil
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	buffer := make([]byte, (1<<16)+headerSize)
	for {
		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
			<-s.bucketNotify
		}

		if f, err := s.readFrame(buffer); err == nil {
			atomic.StoreInt32(&s.dataWasRead, 1)

			switch f.cmd {
			case cmdNOP:
			case cmdSYN:
				s.streamLock.Lock()
				if _, ok := s.streams[f.sid]; !ok {
					stream := newStream(f.sid, s.config.MaxFrameSize, s)
					s.streams[f.sid] = stream
					select {
					case s.chAccepts <- stream:
					case <-s.die:
					}
				}
				s.streamLock.Unlock()
			case cmdFIN:
				s.streamLock.Lock()
				if stream, ok := s.streams[f.sid]; ok {
					stream.markRST()
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			case cmdPSH:
				s.streamLock.Lock()
				if stream, ok := s.streams[f.sid]; ok {
					atomic.AddInt32(&s.bucket, -int32(len(f.data)))
					stream.pushBytes(f.data)
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			default:
				s.Close()
				return
			}
		} else {
			s.Close()
			return
		}
	}
}

// keepAliveSend will periodically send a keepalive meesage to the remote peer.
func (s *Session) keepAliveSend() {
	// Use 'time.Ticker' to ensure that if the keepalive sending falls behind
	// schedule, unnecessary sleeping does not occur.
	ticker := time.NewTicker(s.config.KeepAliveInterval)
	defer ticker.Stop()
	for {
		select {
		case <-s.die:
			return
		case <-ticker.C:
			s.writeFrame(newFrame(cmdNOP, 0), time.Now().Add(s.config.WriteTimeout))
			s.notifyBucket() // force a signal to the recvLoop
		}
	}
}

// keepAliveTimeout will periodically check that some sort of message has been
// sent by the remote peer, closing the session if not.
func (s *Session) keepAliveTimeout() {
	ticker := time.NewTicker(s.config.KeepAliveTimeout)
	defer ticker.Stop()
	for {
		select {
		case <-s.die:
			return
		case <-ticker.C:
			if !atomic.CompareAndSwapInt32(&s.dataWasRead, 1, 0) {
				s.Close()
				return
			}
		}
	}
}

// writeFrame writes the frame to the underlying connection
// and returns the number of bytes written if successful
func (s *Session) writeFrame(frame Frame, timeout time.Time) (int, error) {
	// Verify the frame data size.
	if len(frame.data) > 1<<16 {
		return 0, errLargeFrame
	}

	// Ensure that the configured WriteTimeout is the maximum amount of time
	// that we can wait to send a single frame.
	latestTimeout := time.Now().Add(s.config.WriteTimeout)
	if timeout.IsZero() || timeout.After(latestTimeout) {
		timeout = latestTimeout
	}

	// Determine how much time remains in the timeout, wait for up to that long
	// to grab the sendMu.
	currentTime := time.Now()
	if !timeout.After(currentTime) {
		return 0, errTimeout
	}
	remaining := timeout.Sub(currentTime)
	if !s.sendMu.TryLockTimed(remaining) {
		return 0, errTimeout
	}
	defer s.sendMu.Unlock()

	// Check again that the stream has not been killed.
	select {
	case <-s.die:
		return 0, errBrokenPipe
	default:
	}

	// Prepare the write data.
	buf := make([]byte, headerSize+len(frame.data))
	buf[0] = frame.ver
	buf[1] = frame.cmd
	binary.LittleEndian.PutUint16(buf[2:], uint16(len(frame.data)))
	binary.LittleEndian.PutUint32(buf[4:], frame.sid)
	copy(buf[headerSize:], frame.data)

	// Write the data using the provided writeTimeout.
	s.conn.SetWriteDeadline(timeout)
	n, err := s.conn.Write(buf[:headerSize+len(frame.data)])
	s.conn.SetWriteDeadline(time.Time{})
	n -= headerSize
	if n < 0 {
		n = 0
	}
	return n, err
}
<commit_msg>fix typo<commit_after>package smux

import (
	"encoding/binary"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	siasync "github.com/NebulousLabs/Sia/sync" // TODO: Replace with github.com/NebulousLabs/trymutex
	"github.com/pkg/errors"
)

const (
	defaultAcceptBacklog = 1024
)

var (
	errBrokenPipe      = errors.New("broken pipe")
	errGoAway          = errors.New("stream id overflows, should start a new connection")
	errInvalidProtocol = errors.New("invalid protocol version")
	errLargeFrame      = errors.New("frame is too large to send")
)

// Session defines a multiplexed connection for streams
type Session struct {
	conn        net.Conn
	dataWasRead int32            // used to determine if KeepAlive has failed
	sendMu      siasync.TryMutex // ensures only one thread sends at a time

	config           *Config
	nextStreamID     uint32 // next stream identifier
	nextStreamIDLock sync.Mutex

	bucket       int32         // token bucket
	bucketNotify chan struct{} // used for waiting for tokens

	streams    map[uint32]*Stream // all streams in this session
	streamLock sync.Mutex         // locks streams

	die       chan struct{} // flag session has died
	dieLock   sync.Mutex
	chAccepts chan *Stream

	goAway int32 // flag id exhausted

	deadline atomic.Value
}

func newSession(config *Config, conn net.Conn, client bool) *Session {
	s := new(Session)
	s.die = make(chan struct{})
	s.conn = conn
	s.config = config
	s.streams = make(map[uint32]*Stream)
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.bucket = int32(config.MaxReceiveBuffer)
	s.bucketNotify = make(chan struct{}, 1)

	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 0
	}

	go s.recvLoop()
	// keepaliveSend and keepaliveTimeout need to be separate threads, because
	// the keepaliveSend can block, and especially if the underlying conn has no
	// deadline or a very long deadline, we may not check the keepaliveTimeout
	// for an extended period of time and potentially even end in a deadlock.
	go s.keepAliveSend()
	go s.keepAliveTimeout()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, errBrokenPipe
	}

	// generate stream id
	s.nextStreamIDLock.Lock()
	if s.goAway > 0 {
		s.nextStreamIDLock.Unlock()
		return nil, errGoAway
	}

	s.nextStreamID += 2
	sid := s.nextStreamID
	if sid == sid%2 { // stream-id overflows
		s.goAway = 1
		s.nextStreamIDLock.Unlock()
		return nil, errGoAway
	}
	s.nextStreamIDLock.Unlock()

	stream := newStream(sid, s.config.MaxFrameSize, s)

	if _, err := s.writeFrame(newFrame(cmdSYN, sid), time.Now().Add(s.config.WriteTimeout)); err != nil {
		return nil, errors.Wrap(err, "writeFrame")
	}

	s.streamLock.Lock()
	s.streams[sid] = stream
	s.streamLock.Unlock()
	return stream, nil
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	var deadline <-chan time.Time
	if d, ok := s.deadline.Load().(time.Time); ok && !d.IsZero() {
		timer := time.NewTimer(d.Sub(time.Now()))
		defer timer.Stop()
		deadline = timer.C
	}
	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-deadline:
		return nil, errTimeout
	case <-s.die:
		return nil, errBrokenPipe
	}
}

// Close is used to close the session and all streams.
func (s *Session) Close() (err error) {
	s.dieLock.Lock()

	select {
	case <-s.die:
		s.dieLock.Unlock()
		return errBrokenPipe
	default:
		close(s.die)
		s.dieLock.Unlock()
		s.streamLock.Lock()
		for k := range s.streams {
			s.streams[k].sessionClose()
		}
		s.streamLock.Unlock()
		s.notifyBucket()
		return s.conn.Close()
	}
}

// notifyBucket notifies recvLoop that bucket is available
func (s *Session) notifyBucket() {
	select {
	case s.bucketNotify <- struct{}{}:
	default:
	}
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	if s.IsClosed() {
		return 0
	}
	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	return len(s.streams)
}

// SetDeadline sets a deadline used by Accept* calls.
// A zero time value disables the deadline.
func (s *Session) SetDeadline(t time.Time) error {
	s.deadline.Store(t)
	return nil
}

// notify the session that a stream has closed
func (s *Session) streamClosed(sid uint32) {
	s.streamLock.Lock()
	if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket
		if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
			s.notifyBucket()
		}
	}
	delete(s.streams, sid)
	s.streamLock.Unlock()
}

// returnTokens is called by stream to return token after read
func (s *Session) returnTokens(n int) {
	if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
		s.notifyBucket()
	}
}

// session read a frame from underlying connection
// it's data is pointed to the input buffer
func (s *Session) readFrame(buffer []byte) (f Frame, err error) {
	// Ensure that reading a frame follows the global timeout.
	s.conn.SetReadDeadline(time.Now().Add(s.config.ReadTimeout))
	defer s.conn.SetReadDeadline(time.Time{})

	if _, err := io.ReadFull(s.conn, buffer[:headerSize]); err != nil {
		return f, errors.Wrap(err, "readFrame")
	}

	dec := rawHeader(buffer)
	if dec.Version() != version {
		return f, errInvalidProtocol
	}

	f.ver = dec.Version()
	f.cmd = dec.Cmd()
	f.sid = dec.StreamID()
	if length := dec.Length(); length > 0 {
		if _, err := io.ReadFull(s.conn, buffer[headerSize:headerSize+length]); err != nil {
			return f, errors.Wrap(err, "readFrame")
		}
		f.data = buffer[headerSize : headerSize+length]
	}
	return f, nil
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	buffer := make([]byte, (1<<16)+headerSize)
	for {
		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
			<-s.bucketNotify
		}

		if f, err := s.readFrame(buffer); err == nil {
			atomic.StoreInt32(&s.dataWasRead, 1)

			switch f.cmd {
			case cmdNOP:
			case cmdSYN:
				s.streamLock.Lock()
				if _, ok := s.streams[f.sid]; !ok {
					stream := newStream(f.sid, s.config.MaxFrameSize, s)
					s.streams[f.sid] = stream
					select {
					case s.chAccepts <- stream:
					case <-s.die:
					}
				}
				s.streamLock.Unlock()
			case cmdFIN:
				s.streamLock.Lock()
				if stream, ok := s.streams[f.sid]; ok {
					stream.markRST()
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			case cmdPSH:
				s.streamLock.Lock()
				if stream, ok := s.streams[f.sid]; ok {
					atomic.AddInt32(&s.bucket, -int32(len(f.data)))
					stream.pushBytes(f.data)
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			default:
				s.Close()
				return
			}
		} else {
			s.Close()
			return
		}
	}
}

// keepAliveSend will periodically send a keepalive message to the remote peer.
func (s *Session) keepAliveSend() {
	ticker := time.NewTicker(s.config.KeepAliveInterval)
	defer ticker.Stop()
	for {
		select {
		case <-s.die:
			return
		case <-ticker.C:
			s.writeFrame(newFrame(cmdNOP, 0), time.Now().Add(s.config.WriteTimeout))
			s.notifyBucket() // force a signal to the recvLoop
		}
	}
}

// keepAliveTimeout will periodically check that some sort of message has been
// sent by the remote peer, closing the session if not.
func (s *Session) keepAliveTimeout() {
	ticker := time.NewTicker(s.config.KeepAliveTimeout)
	defer ticker.Stop()
	for {
		select {
		case <-s.die:
			return
		case <-ticker.C:
			if !atomic.CompareAndSwapInt32(&s.dataWasRead, 1, 0) {
				s.Close()
				return
			}
		}
	}
}

// writeFrame writes the frame to the underlying connection
// and returns the number of bytes written if successful
func (s *Session) writeFrame(frame Frame, timeout time.Time) (int, error) {
	// Verify the frame data size.
	if len(frame.data) > 1<<16 {
		return 0, errLargeFrame
	}

	// Ensure that the configured WriteTimeout is the maximum amount of time
	// that we can wait to send a single frame.
	latestTimeout := time.Now().Add(s.config.WriteTimeout)
	if timeout.IsZero() || timeout.After(latestTimeout) {
		timeout = latestTimeout
	}

	// Determine how much time remains in the timeout, wait for up to that long
	// to grab the sendMu.
	currentTime := time.Now()
	if !timeout.After(currentTime) {
		return 0, errTimeout
	}
	remaining := timeout.Sub(currentTime)
	if !s.sendMu.TryLockTimed(remaining) {
		return 0, errTimeout
	}
	defer s.sendMu.Unlock()

	// Check again that the stream has not been killed.
	select {
	case <-s.die:
		return 0, errBrokenPipe
	default:
	}

	// Prepare the write data.
	buf := make([]byte, headerSize+len(frame.data))
	buf[0] = frame.ver
	buf[1] = frame.cmd
	binary.LittleEndian.PutUint16(buf[2:], uint16(len(frame.data)))
	binary.LittleEndian.PutUint32(buf[4:], frame.sid)
	copy(buf[headerSize:], frame.data)

	// Write the data using the provided writeTimeout.
	s.conn.SetWriteDeadline(timeout)
	n, err := s.conn.Write(buf[:headerSize+len(frame.data)])
	s.conn.SetWriteDeadline(time.Time{})
	n -= headerSize
	if n < 0 {
		n = 0
	}
	return n, err
}
<|endoftext|>