<commit_before>package depot

import (
	"errors"

	"github.com/cloudfoundry-incubator/executor/api"
	"github.com/cloudfoundry-incubator/executor/registry"
	"github.com/cloudfoundry-incubator/executor/sequence"
	"github.com/cloudfoundry-incubator/executor/transformer"
	"github.com/cloudfoundry-incubator/garden/warden"
	"github.com/cloudfoundry/gosteno"
)

var (
	ContainerGuidNotAvailable      = errors.New("container guid not available")
	InsufficientResourcesAvailable = errors.New("container guid not available")
	ContainerNotFound              = errors.New("container not found")
	StepsInvalid                   = errors.New("steps invalid")
	LimitsInvalid                  = errors.New("container limits invalid")
)

type client struct {
	containerOwnerName    string
	containerMaxCPUShares uint64
	wardenClient          warden.Client
	registry              registry.Registry
	transformer           *transformer.Transformer
	runActions            chan<- DepotRunAction
	logger                *gosteno.Logger
}

func NewClient(
	containerOwnerName string,
	containerMaxCPUShares uint64,
	wardenClient warden.Client,
	registry registry.Registry,
	transformer *transformer.Transformer,
	runActions chan<- DepotRunAction,
	logger *gosteno.Logger,
) api.Client {
	return &client{
		containerOwnerName:    containerOwnerName,
		containerMaxCPUShares: containerMaxCPUShares,
		wardenClient:          wardenClient,
		registry:              registry,
		transformer:           transformer,
		runActions:            runActions,
		logger:                logger,
	}
}

func (c *client) InitializeContainer(guid string, request api.ContainerInitializationRequest) (api.Container, error) {
	if request.CpuPercent > 100 || request.CpuPercent < 0 {
		return api.Container{}, LimitsInvalid
	}

	reg, err := c.registry.FindByGuid(guid)
	if err != nil {
		c.logger.Infod(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.not-found")
		return api.Container{}, ContainerNotFound
	}

	containerClient, err := c.wardenClient.Create(warden.ContainerSpec{
		Properties: warden.Properties{
			"owner": c.containerOwnerName,
		},
	})
	if err != nil {
		c.logger.Errord(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.create-failed")
		return api.Container{}, err
	}

	err = c.limitContainerDiskAndMemory(reg, containerClient)
	if err != nil {
		c.logger.Errord(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.limit-disk-and-memory-failed")
		return api.Container{}, err
	}

	err = c.limitContainerCPU(request, containerClient)
	if err != nil {
		c.logger.Errord(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.limit-cpu-failed")
		return api.Container{}, err
	}

	portMapping, err := c.mapPorts(request, containerClient)
	if err != nil {
		c.logger.Errord(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.port-mapping-failed")
		return api.Container{}, err
	}

	request.Ports = portMapping

	reg, err = c.registry.Create(reg.Guid, containerClient.Handle(), request)
	if err != nil {
		c.logger.Errord(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.registry-failed")
		return api.Container{}, err
	}

	return reg, nil
}

func (c *client) limitContainerDiskAndMemory(reg api.Container, containerClient warden.Container) error {
	if reg.MemoryMB != 0 {
		err := containerClient.LimitMemory(warden.MemoryLimits{
			LimitInBytes: uint64(reg.MemoryMB * 1024 * 1024),
		})
		if err != nil {
			return err
		}
	}

	if reg.DiskMB != 0 {
		err := containerClient.LimitDisk(warden.DiskLimits{
			ByteHard: uint64(reg.DiskMB * 1024 * 1024),
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *client) limitContainerCPU(request api.ContainerInitializationRequest, containerClient warden.Container) error {
	if request.CpuPercent != 0 {
		err := containerClient.LimitCPU(warden.CPULimits{
			LimitInShares: uint64(float64(c.containerMaxCPUShares) * float64(request.CpuPercent) / 100.0),
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *client) mapPorts(request api.ContainerInitializationRequest, containerClient warden.Container) ([]api.PortMapping, error) {
	var result []api.PortMapping
	for _, mapping := range request.Ports {
		hostPort, containerPort, err := containerClient.NetIn(mapping.HostPort, mapping.ContainerPort)
		if err != nil {
			return nil, err
		}

		result = append(result, api.PortMapping{
			HostPort:      hostPort,
			ContainerPort: containerPort,
		})
	}

	return result, nil
}

func (c *client) AllocateContainer(guid string, request api.ContainerAllocationRequest) (api.Container, error) {
	container, err := c.registry.Reserve(guid, request)
	if err == registry.ErrContainerAlreadyExists {
		c.logger.Warnd(map[string]interface{}{
			"error": err.Error(),
			"guid":  guid,
		}, "executor.allocate-container.container-already-exists")
		return api.Container{}, ContainerGuidNotAvailable
	}

	if err != nil {
		c.logger.Warnd(map[string]interface{}{
			"error": err.Error(),
		}, "executor.allocate-container.full")
		return api.Container{}, InsufficientResourcesAvailable
	}

	return container, nil
}

func (c *client) GetContainer(guid string) (api.Container, error) {
	container, err := c.registry.FindByGuid(guid)
	if err != nil {
		c.logger.Infod(map[string]interface{}{
			"error": err.Error(),
		}, "executor.get-container.not-found")
		return api.Container{}, ContainerNotFound
	}
	return container, nil
}

func (c *client) Run(guid string, request api.ContainerRunRequest) error {
	registration, err := c.registry.FindByGuid(guid)
	if err != nil {
		c.logger.Infod(map[string]interface{}{
			"error": err.Error(),
		}, "executor.run-actions.container-not-found")
		return ContainerNotFound
	}

	container, err := c.wardenClient.Lookup(registration.ContainerHandle)
	if err != nil {
		c.logger.Infod(map[string]interface{}{
			"error": err.Error(),
		}, "executor.run-actions.lookup-failed")
		return err
	}

	var result string
	steps, err := c.transformer.StepsFor(registration.Log, request.Actions, container, &result)
	if err != nil {
		c.logger.Warnd(map[string]interface{}{
			"error": err.Error(),
		}, "executor.run-actions.steps-invalid")
		return StepsInvalid
	}

	c.runActions <- DepotRunAction{
		CompleteURL:  request.CompleteURL,
		Registration: registration,
		Sequence:     sequence.New(steps),
		Result:       &result,
	}

	return nil
}

func (c *client) ListContainers() ([]api.Container, error) {
	return c.registry.GetAllContainers(), nil
}

func (c *client) DeleteContainer(guid string) error {
	registration, err := c.registry.FindByGuid(guid)
	if err != nil {
		return handleDeleteError(err, c.logger)
	}

	//TODO once wardenClient has an ErrNotFound error code, use it
	//to determine if we should delete from registry
	if registration.ContainerHandle != "" {
		err = c.wardenClient.Destroy(registration.ContainerHandle)
		if err != nil {
			return handleDeleteError(err, c.logger)
		}
	}

	err = c.registry.Delete(guid)
	if err != nil {
		return handleDeleteError(err, c.logger)
	}

	return nil
}

func (c *client) RemainingResources() (api.ExecutorResources, error) {
	cap := c.registry.CurrentCapacity()

	return api.ExecutorResources{
		MemoryMB:   cap.MemoryMB,
		DiskMB:     cap.DiskMB,
		Containers: cap.Containers,
	}, nil
}

func (c *client) Ping() error {
	return c.wardenClient.Ping()
}

func (c *client) TotalResources() (api.ExecutorResources, error) {
	totalCapacity := c.registry.TotalCapacity()

	resources := api.ExecutorResources{
		MemoryMB:   totalCapacity.MemoryMB,
		DiskMB:     totalCapacity.DiskMB,
		Containers: totalCapacity.Containers,
	}
	return resources, nil
}

func handleDeleteError(err error, logger *gosteno.Logger) error {
	if err == registry.ErrContainerNotFound {
		logger.Infod(map[string]interface{}{
			"error": err.Error(),
		}, "executor.delete-container.not-found")
		return ContainerNotFound
	}

	logger.Errord(map[string]interface{}{
		"error": err.Error(),
	}, "executor.delete-container.failed")
	return err
}
<commit_msg>fix error string for insufficient resources<commit_after>package depot

import (
	"errors"

	"github.com/cloudfoundry-incubator/executor/api"
	"github.com/cloudfoundry-incubator/executor/registry"
	"github.com/cloudfoundry-incubator/executor/sequence"
	"github.com/cloudfoundry-incubator/executor/transformer"
	"github.com/cloudfoundry-incubator/garden/warden"
	"github.com/cloudfoundry/gosteno"
)

var (
	ContainerGuidNotAvailable      = errors.New("container guid not available")
	InsufficientResourcesAvailable = errors.New("insufficient resources available")
	ContainerNotFound              = errors.New("container not found")
	StepsInvalid                   = errors.New("steps invalid")
	LimitsInvalid                  = errors.New("container limits invalid")
)

type client struct {
	containerOwnerName    string
	containerMaxCPUShares uint64
	wardenClient          warden.Client
	registry              registry.Registry
	transformer           *transformer.Transformer
	runActions            chan<- DepotRunAction
	logger                *gosteno.Logger
}

func NewClient(
	containerOwnerName string,
	containerMaxCPUShares uint64,
	wardenClient warden.Client,
	registry registry.Registry,
	transformer *transformer.Transformer,
	runActions chan<- DepotRunAction,
	logger *gosteno.Logger,
) api.Client {
	return &client{
		containerOwnerName:    containerOwnerName,
		containerMaxCPUShares: containerMaxCPUShares,
		wardenClient:          wardenClient,
		registry:              registry,
		transformer:           transformer,
		runActions:            runActions,
		logger:                logger,
	}
}

func (c *client) InitializeContainer(guid string, request api.ContainerInitializationRequest) (api.Container, error) {
	if request.CpuPercent > 100 || request.CpuPercent < 0 {
		return api.Container{}, LimitsInvalid
	}

	reg, err := c.registry.FindByGuid(guid)
	if err != nil {
		c.logger.Infod(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.not-found")
		return api.Container{}, ContainerNotFound
	}

	containerClient, err := c.wardenClient.Create(warden.ContainerSpec{
		Properties: warden.Properties{
			"owner": c.containerOwnerName,
		},
	})
	if err != nil {
		c.logger.Errord(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.create-failed")
		return api.Container{}, err
	}

	err = c.limitContainerDiskAndMemory(reg, containerClient)
	if err != nil {
		c.logger.Errord(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.limit-disk-and-memory-failed")
		return api.Container{}, err
	}

	err = c.limitContainerCPU(request, containerClient)
	if err != nil {
		c.logger.Errord(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.limit-cpu-failed")
		return api.Container{}, err
	}

	portMapping, err := c.mapPorts(request, containerClient)
	if err != nil {
		c.logger.Errord(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.port-mapping-failed")
		return api.Container{}, err
	}

	request.Ports = portMapping

	reg, err = c.registry.Create(reg.Guid, containerClient.Handle(), request)
	if err != nil {
		c.logger.Errord(map[string]interface{}{
			"error": err.Error(),
		}, "executor.init-container.registry-failed")
		return api.Container{}, err
	}

	return reg, nil
}

func (c *client) limitContainerDiskAndMemory(reg api.Container, containerClient warden.Container) error {
	if reg.MemoryMB != 0 {
		err := containerClient.LimitMemory(warden.MemoryLimits{
			LimitInBytes: uint64(reg.MemoryMB * 1024 * 1024),
		})
		if err != nil {
			return err
		}
	}

	if reg.DiskMB != 0 {
		err := containerClient.LimitDisk(warden.DiskLimits{
			ByteHard: uint64(reg.DiskMB * 1024 * 1024),
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *client) limitContainerCPU(request api.ContainerInitializationRequest, containerClient warden.Container) error {
	if request.CpuPercent != 0 {
		err := containerClient.LimitCPU(warden.CPULimits{
			LimitInShares: uint64(float64(c.containerMaxCPUShares) * float64(request.CpuPercent) / 100.0),
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *client) mapPorts(request api.ContainerInitializationRequest, containerClient warden.Container) ([]api.PortMapping, error) {
	var result []api.PortMapping
	for _, mapping := range request.Ports {
		hostPort, containerPort, err := containerClient.NetIn(mapping.HostPort, mapping.ContainerPort)
		if err != nil {
			return nil, err
		}

		result = append(result, api.PortMapping{
			HostPort:      hostPort,
			ContainerPort: containerPort,
		})
	}

	return result, nil
}

func (c *client) AllocateContainer(guid string, request api.ContainerAllocationRequest) (api.Container, error) {
	container, err := c.registry.Reserve(guid, request)
	if err == registry.ErrContainerAlreadyExists {
		c.logger.Warnd(map[string]interface{}{
			"error": err.Error(),
			"guid":  guid,
		}, "executor.allocate-container.container-already-exists")
		return api.Container{}, ContainerGuidNotAvailable
	}

	if err != nil {
		c.logger.Warnd(map[string]interface{}{
			"error": err.Error(),
		}, "executor.allocate-container.full")
		return api.Container{}, InsufficientResourcesAvailable
	}

	return container, nil
}

func (c *client) GetContainer(guid string) (api.Container, error) {
	container, err := c.registry.FindByGuid(guid)
	if err != nil {
		c.logger.Infod(map[string]interface{}{
			"error": err.Error(),
		}, "executor.get-container.not-found")
		return api.Container{}, ContainerNotFound
	}
	return container, nil
}

func (c *client) Run(guid string, request api.ContainerRunRequest) error {
	registration, err := c.registry.FindByGuid(guid)
	if err != nil {
		c.logger.Infod(map[string]interface{}{
			"error": err.Error(),
		}, "executor.run-actions.container-not-found")
		return ContainerNotFound
	}

	container, err := c.wardenClient.Lookup(registration.ContainerHandle)
	if err != nil {
		c.logger.Infod(map[string]interface{}{
			"error": err.Error(),
		}, "executor.run-actions.lookup-failed")
		return err
	}

	var result string
	steps, err := c.transformer.StepsFor(registration.Log, request.Actions, container, &result)
	if err != nil {
		c.logger.Warnd(map[string]interface{}{
			"error": err.Error(),
		}, "executor.run-actions.steps-invalid")
		return StepsInvalid
	}

	c.runActions <- DepotRunAction{
		CompleteURL:  request.CompleteURL,
		Registration: registration,
		Sequence:     sequence.New(steps),
		Result:       &result,
	}

	return nil
}

func (c *client) ListContainers() ([]api.Container, error) {
	return c.registry.GetAllContainers(), nil
}

func (c *client) DeleteContainer(guid string) error {
	registration, err := c.registry.FindByGuid(guid)
	if err != nil {
		return handleDeleteError(err, c.logger)
	}

	//TODO once wardenClient has an ErrNotFound error code, use it
	//to determine if we should delete from registry
	if registration.ContainerHandle != "" {
		err = c.wardenClient.Destroy(registration.ContainerHandle)
		if err != nil {
			return handleDeleteError(err, c.logger)
		}
	}

	err = c.registry.Delete(guid)
	if err != nil {
		return handleDeleteError(err, c.logger)
	}

	return nil
}

func (c *client) RemainingResources() (api.ExecutorResources, error) {
	cap := c.registry.CurrentCapacity()

	return api.ExecutorResources{
		MemoryMB:   cap.MemoryMB,
		DiskMB:     cap.DiskMB,
		Containers: cap.Containers,
	}, nil
}

func (c *client) Ping() error {
	return c.wardenClient.Ping()
}

func (c *client) TotalResources() (api.ExecutorResources, error) {
	totalCapacity := c.registry.TotalCapacity()

	resources := api.ExecutorResources{
		MemoryMB:   totalCapacity.MemoryMB,
		DiskMB:     totalCapacity.DiskMB,
		Containers: totalCapacity.Containers,
	}
	return resources, nil
}

func handleDeleteError(err error, logger *gosteno.Logger) error {
	if err == registry.ErrContainerNotFound {
		logger.Infod(map[string]interface{}{
			"error": err.Error(),
		}, "executor.delete-container.not-found")
		return ContainerNotFound
	}

	logger.Errord(map[string]interface{}{
		"error": err.Error(),
	}, "executor.delete-container.failed")
	return err
}
<|endoftext|>


<commit_before>// Copyright 2017 CoreOS, Inc.
//
// 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 general

import (
	"github.com/coreos/ignition/tests/register"
	"github.com/coreos/ignition/tests/types"
)

func init() {
	register.Register(register.PositiveTest, ReformatFilesystemAndWriteFile())
	register.Register(register.PositiveTest, ReplaceConfigWithRemoteConfigHTTP())
	register.Register(register.PositiveTest, AppendConfigWithRemoteConfigHTTP())
	register.Register(register.PositiveTest, ReplaceConfigWithRemoteConfigTFTP())
	register.Register(register.PositiveTest, AppendConfigWithRemoteConfigTFTP())
	register.Register(register.PositiveTest, ReplaceConfigWithRemoteConfigOEM())
	register.Register(register.PositiveTest, AppendConfigWithRemoteConfigOEM())
	register.Register(register.PositiveTest, VersionOnlyConfig())
	register.Register(register.PositiveTest, EmptyUserdata())
}

func ReformatFilesystemAndWriteFile() types.Test {
	name := "Reformat Filesystem to ext4 & drop file in /ignition/test"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	mntDevices := []types.MntDevice{
		{
			Label:        "EFI-SYSTEM",
			Substitution: "$DEVICE",
		},
	}
	config := `{
		"ignition": {"version": "2.0.0"},
		"storage": {
			"filesystems": [{
				"mount": {
					"device": "$DEVICE",
					"format": "ext4",
					"create": {
						"force": true
					}},
				 "name": "test"}],
			"files": [{
				"filesystem": "test",
				"path": "/ignition/test",
				"contents": {"source": "data:,asdf"}
			}]}
	}`

	out[0].Partitions.GetPartition("EFI-SYSTEM").FilesystemType = "ext4"
	out[0].Partitions.GetPartition("EFI-SYSTEM").Files = []types.File{
		{
			Node: types.Node{
				Name:      "test",
				Directory: "ignition",
			},
			Contents: "asdf",
		},
	}

	return types.Test{
		Name:       name,
		In:         in,
		Out:        out,
		MntDevices: mntDevices,
		Config:     config,
	}
}

func ReplaceConfigWithRemoteConfigHTTP() types.Test {
	name := "Replacing the Config with a Remote Config from HTTP"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
	  "ignition": {
	    "version": "2.0.0",
	    "config": {
	      "replace": {
	        "source": "http://127.0.0.1:8080/config",
			"verification": { "hash": "sha512-41d9a1593dd4cbcacc966dce574523ffe3780ec2710716fab28b46f0f24d20b5ec49f307a9e9d331af958e508f472f32135c740d1214c5f02fc36016b538e7ff" }
	      }
	    }
	  }
	}`
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func ReplaceConfigWithRemoteConfigTFTP() types.Test {
	name := "Replacing the Config with a Remote Config from TFTP"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
          "ignition": {
            "version": "2.1.0",
            "config": {
              "replace": {
                "source": "tftp://127.0.0.1:69/config",
                        "verification": { "hash": "sha512-fa00083efe3f00eb984e6dc27cc8673585cce4319e39099ce014103619ae7ab7dc3555e51401c7df472bdd125c552e528f54d717b8147129c99836d3dedc9760" }
              }
            }
          }
        }`
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func ReplaceConfigWithRemoteConfigOEM() types.Test {
	name := "Replacing the Config with a Remote Config from OEM"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
          "ignition": {
            "version": "2.1.0",
            "config": {
              "replace": {
                "source": "oem:///config",
                        "verification": { "hash": "sha512-73395ffef4b1aefac56b6406f7aed307199d960cc8ad9317e0e8b6497a64f879b33fd59eca533f5f139aa4237f7d81de08c6f7f17db9dd2c072e9ecccb0fed42" }
              }
            }
          }
        }`
	in[0].Partitions.AddFiles("OEM", []types.File{
		{
			Node: types.Node{
				Name: "config",
			},
			Contents: `{
	"ignition": { "version": "2.1.0" },
	"storage": {
		"files": [{
		  "filesystem": "root",
		  "path": "/foo/bar",
		  "contents": { "source": "data:,example%20file%0A" }
		}]
	}
}`,
		},
	})
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func AppendConfigWithRemoteConfigHTTP() types.Test {
	name := "Appending to the Config with a Remote Config from HTTP"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
	  "ignition": {
	    "version": "2.0.0",
	    "config": {
	      "append": [{
	        "source": "http://127.0.0.1:8080/config",
			"verification": { "hash": "sha512-41d9a1593dd4cbcacc966dce574523ffe3780ec2710716fab28b46f0f24d20b5ec49f307a9e9d331af958e508f472f32135c740d1214c5f02fc36016b538e7ff" }
	      }]
	    }
	  },
      "storage": {
        "files": [{
          "filesystem": "root",
          "path": "/foo/bar2",
          "contents": { "source": "data:,another%20example%20file%0A" }
        }]
      }
	}`
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
		{
			Node: types.Node{
				Name:      "bar2",
				Directory: "foo",
			},
			Contents: "another example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func AppendConfigWithRemoteConfigTFTP() types.Test {
	name := "Appending to the Config with a Remote Config from TFTP"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
          "ignition": {
            "version": "2.1.0",
            "config": {
              "append": [{
                "source": "tftp://127.0.0.1:69/config",
                        "verification": { "hash": "sha512-fa00083efe3f00eb984e6dc27cc8673585cce4319e39099ce014103619ae7ab7dc3555e51401c7df472bdd125c552e528f54d717b8147129c99836d3dedc9760" }
              }]
            }
          },
      "storage": {
        "files": [{
          "filesystem": "root",
          "path": "/foo/bar2",
          "contents": { "source": "data:,another%20example%20file%0A" }
        }]
      }
        }`
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
		{
			Node: types.Node{
				Name:      "bar2",
				Directory: "foo",
			},
			Contents: "another example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func AppendConfigWithRemoteConfigOEM() types.Test {
	name := "Appending to the Config with a Remote Config from OEM"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
          "ignition": {
            "version": "2.1.0",
            "config": {
              "append": [{
                "source": "oem:///config",
                        "verification": { "hash": "sha512-73395ffef4b1aefac56b6406f7aed307199d960cc8ad9317e0e8b6497a64f879b33fd59eca533f5f139aa4237f7d81de08c6f7f17db9dd2c072e9ecccb0fed42" }
              }]
            }
          },
      "storage": {
        "files": [{
          "filesystem": "root",
          "path": "/foo/bar2",
          "contents": { "source": "data:,another%20example%20file%0A" }
        }]
      }
        }`
	in[0].Partitions.AddFiles("OEM", []types.File{
		{
			Node: types.Node{
				Name: "config",
			},
			Contents: `{
	"ignition": { "version": "2.1.0" },
	"storage": {
		"files": [{
		  "filesystem": "root",
		  "path": "/foo/bar",
		  "contents": { "source": "data:,example%20file%0A" }
		}]
	}
}`,
		},
	})
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
		{
			Node: types.Node{
				Name:      "bar2",
				Directory: "foo",
			},
			Contents: "another example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func VersionOnlyConfig() types.Test {
	name := "Version Only Config"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
		"ignition": {"version": "2.1.0"}
	}`

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func EmptyUserdata() types.Test {
	name := "Empty Userdata"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := ``

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}
<commit_msg>tests: add tests for data URI's in append/replace<commit_after>// Copyright 2017 CoreOS, Inc.
//
// 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 general

import (
	"github.com/coreos/ignition/tests/register"
	"github.com/coreos/ignition/tests/types"
)

func init() {
	// TODO: Add S3 tests
	register.Register(register.PositiveTest, ReformatFilesystemAndWriteFile())
	register.Register(register.PositiveTest, ReplaceConfigWithRemoteConfigHTTP())
	register.Register(register.PositiveTest, AppendConfigWithRemoteConfigHTTP())
	register.Register(register.PositiveTest, ReplaceConfigWithRemoteConfigTFTP())
	register.Register(register.PositiveTest, AppendConfigWithRemoteConfigTFTP())
	register.Register(register.PositiveTest, ReplaceConfigWithRemoteConfigOEM())
	register.Register(register.PositiveTest, AppendConfigWithRemoteConfigOEM())
	register.Register(register.PositiveTest, ReplaceConfigWithRemoteConfigData())
	register.Register(register.PositiveTest, AppendConfigWithRemoteConfigData())
	register.Register(register.PositiveTest, VersionOnlyConfig())
	register.Register(register.PositiveTest, EmptyUserdata())
}

func ReformatFilesystemAndWriteFile() types.Test {
	name := "Reformat Filesystem to ext4 & drop file in /ignition/test"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	mntDevices := []types.MntDevice{
		{
			Label:        "EFI-SYSTEM",
			Substitution: "$DEVICE",
		},
	}
	config := `{
		"ignition": {"version": "2.0.0"},
		"storage": {
			"filesystems": [{
				"mount": {
					"device": "$DEVICE",
					"format": "ext4",
					"create": {
						"force": true
					}},
				 "name": "test"}],
			"files": [{
				"filesystem": "test",
				"path": "/ignition/test",
				"contents": {"source": "data:,asdf"}
			}]}
	}`

	out[0].Partitions.GetPartition("EFI-SYSTEM").FilesystemType = "ext4"
	out[0].Partitions.GetPartition("EFI-SYSTEM").Files = []types.File{
		{
			Node: types.Node{
				Name:      "test",
				Directory: "ignition",
			},
			Contents: "asdf",
		},
	}

	return types.Test{
		Name:       name,
		In:         in,
		Out:        out,
		MntDevices: mntDevices,
		Config:     config,
	}
}

func ReplaceConfigWithRemoteConfigHTTP() types.Test {
	name := "Replacing the Config with a Remote Config from HTTP"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
	  "ignition": {
	    "version": "2.0.0",
	    "config": {
	      "replace": {
	        "source": "http://127.0.0.1:8080/config",
			"verification": { "hash": "sha512-41d9a1593dd4cbcacc966dce574523ffe3780ec2710716fab28b46f0f24d20b5ec49f307a9e9d331af958e508f472f32135c740d1214c5f02fc36016b538e7ff" }
	      }
	    }
	  }
	}`
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func ReplaceConfigWithRemoteConfigTFTP() types.Test {
	name := "Replacing the Config with a Remote Config from TFTP"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
          "ignition": {
            "version": "2.1.0",
            "config": {
              "replace": {
                "source": "tftp://127.0.0.1:69/config",
                        "verification": { "hash": "sha512-fa00083efe3f00eb984e6dc27cc8673585cce4319e39099ce014103619ae7ab7dc3555e51401c7df472bdd125c552e528f54d717b8147129c99836d3dedc9760" }
              }
            }
          }
        }`
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func ReplaceConfigWithRemoteConfigOEM() types.Test {
	name := "Replacing the Config with a Remote Config from OEM"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
          "ignition": {
            "version": "2.1.0",
            "config": {
              "replace": {
                "source": "oem:///config",
                        "verification": { "hash": "sha512-73395ffef4b1aefac56b6406f7aed307199d960cc8ad9317e0e8b6497a64f879b33fd59eca533f5f139aa4237f7d81de08c6f7f17db9dd2c072e9ecccb0fed42" }
              }
            }
          }
        }`
	in[0].Partitions.AddFiles("OEM", []types.File{
		{
			Node: types.Node{
				Name: "config",
			},
			Contents: `{
	"ignition": { "version": "2.1.0" },
	"storage": {
		"files": [{
		  "filesystem": "root",
		  "path": "/foo/bar",
		  "contents": { "source": "data:,example%20file%0A" }
		}]
	}
}`,
		},
	})
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func AppendConfigWithRemoteConfigHTTP() types.Test {
	name := "Appending to the Config with a Remote Config from HTTP"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
	  "ignition": {
	    "version": "2.0.0",
	    "config": {
	      "append": [{
	        "source": "http://127.0.0.1:8080/config",
			"verification": { "hash": "sha512-41d9a1593dd4cbcacc966dce574523ffe3780ec2710716fab28b46f0f24d20b5ec49f307a9e9d331af958e508f472f32135c740d1214c5f02fc36016b538e7ff" }
	      }]
	    }
	  },
      "storage": {
        "files": [{
          "filesystem": "root",
          "path": "/foo/bar2",
          "contents": { "source": "data:,another%20example%20file%0A" }
        }]
      }
	}`
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
		{
			Node: types.Node{
				Name:      "bar2",
				Directory: "foo",
			},
			Contents: "another example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func AppendConfigWithRemoteConfigTFTP() types.Test {
	name := "Appending to the Config with a Remote Config from TFTP"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
          "ignition": {
            "version": "2.1.0",
            "config": {
              "append": [{
                "source": "tftp://127.0.0.1:69/config",
                        "verification": { "hash": "sha512-fa00083efe3f00eb984e6dc27cc8673585cce4319e39099ce014103619ae7ab7dc3555e51401c7df472bdd125c552e528f54d717b8147129c99836d3dedc9760" }
              }]
            }
          },
      "storage": {
        "files": [{
          "filesystem": "root",
          "path": "/foo/bar2",
          "contents": { "source": "data:,another%20example%20file%0A" }
        }]
      }
        }`
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
		{
			Node: types.Node{
				Name:      "bar2",
				Directory: "foo",
			},
			Contents: "another example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func AppendConfigWithRemoteConfigOEM() types.Test {
	name := "Appending to the Config with a Remote Config from OEM"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
          "ignition": {
            "version": "2.1.0",
            "config": {
              "append": [{
                "source": "oem:///config",
                        "verification": { "hash": "sha512-73395ffef4b1aefac56b6406f7aed307199d960cc8ad9317e0e8b6497a64f879b33fd59eca533f5f139aa4237f7d81de08c6f7f17db9dd2c072e9ecccb0fed42" }
              }]
            }
          },
      "storage": {
        "files": [{
          "filesystem": "root",
          "path": "/foo/bar2",
          "contents": { "source": "data:,another%20example%20file%0A" }
        }]
      }
        }`
	in[0].Partitions.AddFiles("OEM", []types.File{
		{
			Node: types.Node{
				Name: "config",
			},
			Contents: `{
	"ignition": { "version": "2.1.0" },
	"storage": {
		"files": [{
		  "filesystem": "root",
		  "path": "/foo/bar",
		  "contents": { "source": "data:,example%20file%0A" }
		}]
	}
}`,
		},
	})
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "example file\n",
		},
		{
			Node: types.Node{
				Name:      "bar2",
				Directory: "foo",
			},
			Contents: "another example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func ReplaceConfigWithRemoteConfigData() types.Test {
	name := "Replacing the Config with a Remote Config from Data"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
          "ignition": {
            "version": "2.1.0",
            "config": {
              "replace": {
				  "source": "data:,%7B%22ignition%22%3A%7B%22version%22%3A%20%222.1.0%22%7D%2C%22storage%22%3A%20%7B%22files%22%3A%20%5B%7B%22filesystem%22%3A%20%22root%22%2C%22path%22%3A%20%22%2Ffoo%2Fbar%22%2C%22contents%22%3A%7B%22source%22%3A%22data%3A%2Canother%2520example%2520file%250A%22%7D%7D%5D%7D%7D%0A"
              }
            }
          }
        }`
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "another example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func AppendConfigWithRemoteConfigData() types.Test {
	name := "Appending to the Config with a Remote Config from Data"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
          "ignition": {
            "version": "2.1.0",
            "config": {
              "append": [{
				  "source": "data:,%7B%22ignition%22%3A%7B%22version%22%3A%20%222.1.0%22%7D%2C%22storage%22%3A%20%7B%22files%22%3A%20%5B%7B%22filesystem%22%3A%20%22root%22%2C%22path%22%3A%20%22%2Ffoo%2Fbar%22%2C%22contents%22%3A%7B%22source%22%3A%22data%3A%2Canother%2520example%2520file%250A%22%7D%7D%5D%7D%7D%0A"
              }]
            }
          }
        }`
	out[0].Partitions.AddFiles("ROOT", []types.File{
		{
			Node: types.Node{
				Name:      "bar",
				Directory: "foo",
			},
			Contents: "another example file\n",
		},
	})

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func VersionOnlyConfig() types.Test {
	name := "Version Only Config"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := `{
		"ignition": {"version": "2.1.0"}
	}`

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}

func EmptyUserdata() types.Test {
	name := "Empty Userdata"
	in := types.GetBaseDisk()
	out := types.GetBaseDisk()
	config := ``

	return types.Test{
		Name:   name,
		In:     in,
		Out:    out,
		Config: config,
	}
}
<|endoftext|>


<commit_before>package xorm

import (
	"crypto/tls"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/go-xorm/core"
)

// func init() {
// 	RegisterDialect("mysql", &mysql{})
// }

type mysql struct {
	core.Base
	net               string
	addr              string
	params            map[string]string
	loc               *time.Location
	timeout           time.Duration
	tls               *tls.Config
	allowAllFiles     bool
	allowOldPasswords bool
	clientFoundRows   bool
}

func (db *mysql) Init(d *core.DB, uri *core.Uri, drivername, dataSourceName string) error {
	return db.Base.Init(d, db, uri, drivername, dataSourceName)
}

func (db *mysql) SqlType(c *core.Column) string {
	var res string
	switch t := c.SQLType.Name; t {
	case core.Bool:
		res = core.TinyInt
		c.Length = 1
	case core.Serial:
		c.IsAutoIncrement = true
		c.IsPrimaryKey = true
		c.Nullable = false
		res = core.Int
	case core.BigSerial:
		c.IsAutoIncrement = true
		c.IsPrimaryKey = true
		c.Nullable = false
		res = core.BigInt
	case core.Bytea:
		res = core.Blob
	case core.TimeStampz:
		res = core.Char
		c.Length = 64
	case core.Enum: //mysql enum
		res = core.Enum
		res += "("
		for v, k := range c.EnumOptions {
			if k > 0 {
				res += fmt.Sprintf(",'%v'", v)
			} else {
				res += fmt.Sprintf("'%v'", v)
			}
		}
		res += ")"
	case core.Set: //mysql set
		res = core.Set
		res += "("
		for v, k := range c.SetOptions {
			if k > 0 {
				res += fmt.Sprintf(",'%v'", v)
			} else {
				res += fmt.Sprintf("'%v'", v)
			}
		}
		res += ")"
	default:
		res = t
	}

	var hasLen1 bool = (c.Length > 0)
	var hasLen2 bool = (c.Length2 > 0)

	if hasLen2 {
		res += "(" + strconv.Itoa(c.Length) + "," + strconv.Itoa(c.Length2) + ")"
	} else if hasLen1 {
		res += "(" + strconv.Itoa(c.Length) + ")"
	}
	return res
}

func (db *mysql) SupportInsertMany() bool {
	return true
}

func (db *mysql) QuoteStr() string {
	return "`"
}

func (db *mysql) SupportEngine() bool {
	return true
}

func (db *mysql) AutoIncrStr() string {
	return "AUTO_INCREMENT"
}

func (db *mysql) SupportCharset() bool {
	return true
}

func (db *mysql) IndexOnTable() bool {
	return true
}

func (db *mysql) IndexCheckSql(tableName, idxName string) (string, []interface{}) {
	args := []interface{}{db.DbName, tableName, idxName}
	sql := "SELECT `INDEX_NAME` FROM `INFORMATION_SCHEMA`.`STATISTICS`"
	sql += " WHERE `TABLE_SCHEMA` = ? AND `TABLE_NAME` = ? AND `INDEX_NAME`=?"
	return sql, args
}

/*func (db *mysql) ColumnCheckSql(tableName, colName string) (string, []interface{}) {
	args := []interface{}{db.DbName, tableName, colName}
	sql := "SELECT `COLUMN_NAME` FROM `INFORMATION_SCHEMA`.`COLUMNS` WHERE `TABLE_SCHEMA` = ? AND `TABLE_NAME` = ? AND `COLUMN_NAME` = ?"
	return sql, args
}*/

func (db *mysql) TableCheckSql(tableName string) (string, []interface{}) {
	args := []interface{}{db.DbName, tableName}
	sql := "SELECT `TABLE_NAME` from `INFORMATION_SCHEMA`.`TABLES` WHERE `TABLE_SCHEMA`=? and `TABLE_NAME`=?"
	return sql, args
}

func (db *mysql) GetColumns(tableName string) ([]string, map[string]*core.Column, error) {
	args := []interface{}{db.DbName, tableName}
	s := "SELECT `COLUMN_NAME`, `IS_NULLABLE`, `COLUMN_DEFAULT`, `COLUMN_TYPE`," +
		" `COLUMN_KEY`, `EXTRA` FROM `INFORMATION_SCHEMA`.`COLUMNS` WHERE `TABLE_SCHEMA` = ? AND `TABLE_NAME` = ?"

	rows, err := db.DB().Query(s, args...)
	if err != nil {
		return nil, nil, err
	}
	defer rows.Close()

	cols := make(map[string]*core.Column)
	colSeq := make([]string, 0)
	for rows.Next() {
		col := new(core.Column)
		col.Indexes = make(map[string]bool)

		var columnName, isNullable, colType, colKey, extra string
		var colDefault *string
		err = rows.Scan(&columnName, &isNullable, &colDefault, &colType, &colKey, &extra)
		if err != nil {
			return nil, nil, err
		}
		//fmt.Println(columnName, isNullable, colType, colKey, extra, colDefault)
		col.Name = strings.Trim(columnName, "` ")
		if "YES" == isNullable {
			col.Nullable = true
		}

		if colDefault != nil {
			col.Default = *colDefault
			if col.Default == "" {
				col.DefaultIsEmpty = true
			}
		}

		cts := strings.Split(colType, "(")
		colName := cts[0]
		colType = strings.ToUpper(colName)
		var len1, len2 int
		if len(cts) == 2 {
			idx := strings.Index(cts[1], ")")
			if colType == core.Enum && cts[1][0] == '\'' { //enum
				options := strings.Split(cts[1][0:idx], ",")
				col.EnumOptions = make(map[string]int)
				for k, v := range options {
					v = strings.TrimSpace(v)
					v = strings.Trim(v, "'")
					col.EnumOptions[v] = k
				}
			} else if colType == core.Set && cts[1][0] == '\'' {
				options := strings.Split(cts[1][0:idx], ",")
				col.SetOptions = make(map[string]int)
				for k, v := range options {
					v = strings.TrimSpace(v)
					v = strings.Trim(v, "'")
					col.SetOptions[v] = k
				}
			} else {
				lens := strings.Split(cts[1][0:idx], ",")
				len1, err = strconv.Atoi(strings.TrimSpace(lens[0]))
				if err != nil {
					return nil, nil, err
				}
				if len(lens) == 2 {
					len2, err = strconv.Atoi(lens[1])
					if err != nil {
						return nil, nil, err
					}
				}
			}
		}
		if colType == "FLOAT UNSIGNED" {
			colType = "FLOAT"
		}
		col.Length = len1
		col.Length2 = len2
		if _, ok := core.SqlTypes[colType]; ok {
			col.SQLType = core.SQLType{colType, len1, len2}
		} else {
			return nil, nil, errors.New(fmt.Sprintf("unkonw colType %v", colType))
		}

		if colKey == "PRI" {
			col.IsPrimaryKey = true
		}
		if colKey == "UNI" {
			//col.is
		}

		if extra == "auto_increment" {
			col.IsAutoIncrement = true
		}

		if col.SQLType.IsText() {
			if col.Default != "" {
				col.Default = "'" + col.Default + "'"
			} else {
				if col.DefaultIsEmpty {
					col.Default = "''"
				}
			}
		}
		cols[col.Name] = col
		colSeq = append(colSeq, col.Name)
	}
	return colSeq, cols, nil
}

func (db *mysql) GetTables() ([]*core.Table, error) {
	args := []interface{}{db.DbName}
	s := "SELECT `TABLE_NAME`, `ENGINE`, `TABLE_ROWS`, `AUTO_INCREMENT` from `INFORMATION_SCHEMA`.`TABLES` WHERE `TABLE_SCHEMA`=?"

	rows, err := db.DB().Query(s, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	tables := make([]*core.Table, 0)
	for rows.Next() {
		table := core.NewEmptyTable()
		var name, engine, tableRows string
		var autoIncr *string
		err = rows.Scan(&name, &engine, &tableRows, &autoIncr)
		if err != nil {
			return nil, err
		}

		table.Name = name
		tables = append(tables, table)
	}
	return tables, nil
}

func (db *mysql) GetIndexes(tableName string) (map[string]*core.Index, error) {
	args := []interface{}{db.DbName, tableName}
	s := "SELECT `INDEX_NAME`, `NON_UNIQUE`, `COLUMN_NAME` FROM `INFORMATION_SCHEMA`.`STATISTICS` WHERE `TABLE_SCHEMA` = ? AND `TABLE_NAME` = ?"

	rows, err := db.DB().Query(s, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	indexes := make(map[string]*core.Index, 0)
	for rows.Next() {
		var indexType int
		var indexName, colName, nonUnique string
		err = rows.Scan(&indexName, &nonUnique, &colName)
		if err != nil {
			return nil, err
		}

		if indexName == "PRIMARY" {
			continue
		}

		if "YES" == nonUnique || nonUnique == "1" {
			indexType = core.IndexType
		} else {
			indexType = core.UniqueType
		}

		colName = strings.Trim(colName, "` ")

		if strings.HasPrefix(indexName, "IDX_"+tableName) || strings.HasPrefix(indexName, "UQE_"+tableName) {
			indexName = indexName[5+len(tableName) : len(indexName)]
		}

		var index *core.Index
		var ok bool
		if index, ok = indexes[indexName]; !ok {
			index = new(core.Index)
			index.Type = indexType
			index.Name = indexName
			indexes[indexName] = index
		}
		index.AddColumn(colName)
	}
	return indexes, nil
}

func (db *mysql) Filters() []core.Filter {
	return []core.Filter{&core.IdFilter{}}
}
<commit_msg>fix bug<commit_after>package xorm

import (
	"crypto/tls"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/go-xorm/core"
)

// func init() {
// 	RegisterDialect("mysql", &mysql{})
// }

type mysql struct {
	core.Base
	net               string
	addr              string
	params            map[string]string
	loc               *time.Location
	timeout           time.Duration
	tls               *tls.Config
	allowAllFiles     bool
	allowOldPasswords bool
	clientFoundRows   bool
}

func (db *mysql) Init(d *core.DB, uri *core.Uri, drivername, dataSourceName string) error {
	return db.Base.Init(d, db, uri, drivername, dataSourceName)
}

func (db *mysql) SqlType(c *core.Column) string {
	var res string
	switch t := c.SQLType.Name; t {
	case core.Bool:
		res = core.TinyInt
		c.Length = 1
	case core.Serial:
		c.IsAutoIncrement = true
		c.IsPrimaryKey = true
		c.Nullable = false
		res = core.Int
	case core.BigSerial:
		c.IsAutoIncrement = true
		c.IsPrimaryKey = true
		c.Nullable = false
		res = core.BigInt
	case core.Bytea:
		res = core.Blob
	case core.TimeStampz:
		res = core.Char
		c.Length = 64
	case core.Enum: //mysql enum
		res = core.Enum
		res += "("
		opts := ""
		for v, _ := range c.EnumOptions {
			opts += fmt.Sprintf(",'%v'", v)
		}
		res += strings.TrimLeft(opts,",")
		res += ")"
	case core.Set: //mysql set
		res = core.Set
		res += "("
		opts := ""
		for v, _ := range c.SetOptions {
			opts += fmt.Sprintf(",'%v'", v)
		}
		res += strings.TrimLeft(opts,",")
		res += ")"
	default:
		res = t
	}

	var hasLen1 bool = (c.Length > 0)
	var hasLen2 bool = (c.Length2 > 0)

	if hasLen2 {
		res += "(" + strconv.Itoa(c.Length) + "," + strconv.Itoa(c.Length2) + ")"
	} else if hasLen1 {
		res += "(" + strconv.Itoa(c.Length) + ")"
	}
	return res
}

func (db *mysql) SupportInsertMany() bool {
	return true
}

func (db *mysql) QuoteStr() string {
	return "`"
}

func (db *mysql) SupportEngine() bool {
	return true
}

func (db *mysql) AutoIncrStr() string {
	return "AUTO_INCREMENT"
}

func (db *mysql) SupportCharset() bool {
	return true
}

func (db *mysql) IndexOnTable() bool {
	return true
}

func (db *mysql) IndexCheckSql(tableName, idxName string) (string, []interface{}) {
	args := []interface{}{db.DbName, tableName, idxName}
	sql := "SELECT `INDEX_NAME` FROM `INFORMATION_SCHEMA`.`STATISTICS`"
	sql += " WHERE `TABLE_SCHEMA` = ? AND `TABLE_NAME` = ? AND `INDEX_NAME`=?"
	return sql, args
}

/*func (db *mysql) ColumnCheckSql(tableName, colName string) (string, []interface{}) {
	args := []interface{}{db.DbName, tableName, colName}
	sql := "SELECT `COLUMN_NAME` FROM `INFORMATION_SCHEMA`.`COLUMNS` WHERE `TABLE_SCHEMA` = ? AND `TABLE_NAME` = ? AND `COLUMN_NAME` = ?"
	return sql, args
}*/

func (db *mysql) TableCheckSql(tableName string) (string, []interface{}) {
	args := []interface{}{db.DbName, tableName}
	sql := "SELECT `TABLE_NAME` from `INFORMATION_SCHEMA`.`TABLES` WHERE `TABLE_SCHEMA`=? and `TABLE_NAME`=?"
	return sql, args
}

func (db *mysql) GetColumns(tableName string) ([]string, map[string]*core.Column, error) {
	args := []interface{}{db.DbName, tableName}
	s := "SELECT `COLUMN_NAME`, `IS_NULLABLE`, `COLUMN_DEFAULT`, `COLUMN_TYPE`," +
		" `COLUMN_KEY`, `EXTRA` FROM `INFORMATION_SCHEMA`.`COLUMNS` WHERE `TABLE_SCHEMA` = ? AND `TABLE_NAME` = ?"

	rows, err := db.DB().Query(s, args...)
	if err != nil {
		return nil, nil, err
	}
	defer rows.Close()

	cols := make(map[string]*core.Column)
	colSeq := make([]string, 0)
	for rows.Next() {
		col := new(core.Column)
		col.Indexes = make(map[string]bool)

		var columnName, isNullable, colType, colKey, extra string
		var colDefault *string
		err = rows.Scan(&columnName, &isNullable, &colDefault, &colType, &colKey, &extra)
		if err != nil {
			return nil, nil, err
		}
		//fmt.Println(columnName, isNullable, colType, colKey, extra, colDefault)
		col.Name = strings.Trim(columnName, "` ")
		if "YES" == isNullable {
			col.Nullable = true
		}

		if colDefault != nil {
			col.Default = *colDefault
			if col.Default == "" {
				col.DefaultIsEmpty = true
			}
		}

		cts := strings.Split(colType, "(")
		colName := cts[0]
		colType = strings.ToUpper(colName)
		var len1, len2 int
		if len(cts) == 2 {
			idx := strings.Index(cts[1], ")")
			if colType == core.Enum && cts[1][0] == '\'' { //enum
				options := strings.Split(cts[1][0:idx], ",")
				col.EnumOptions = make(map[string]int)
				for k, v := range options {
					v = strings.TrimSpace(v)
					v = strings.Trim(v, "'")
					col.EnumOptions[v] = k
				}
			} else if colType == core.Set && cts[1][0] == '\'' {
				options := strings.Split(cts[1][0:idx], ",")
				col.SetOptions = make(map[string]int)
				for k, v := range options {
					v = strings.TrimSpace(v)
					v = strings.Trim(v, "'")
					col.SetOptions[v] = k
				}
			} else {
				lens := strings.Split(cts[1][0:idx], ",")
				len1, err = strconv.Atoi(strings.TrimSpace(lens[0]))
				if err != nil {
					return nil, nil, err
				}
				if len(lens) == 2 {
					len2, err = strconv.Atoi(lens[1])
					if err != nil {
						return nil, nil, err
					}
				}
			}
		}
		if colType == "FLOAT UNSIGNED" {
			colType = "FLOAT"
		}
		col.Length = len1
		col.Length2 = len2
		if _, ok := core.SqlTypes[colType]; ok {
			col.SQLType = core.SQLType{colType, len1, len2}
		} else {
			return nil, nil, errors.New(fmt.Sprintf("unkonw colType %v", colType))
		}

		if colKey == "PRI" {
			col.IsPrimaryKey = true
		}
		if colKey == "UNI" {
			//col.is
		}

		if extra == "auto_increment" {
			col.IsAutoIncrement = true
		}

		if col.SQLType.IsText() {
			if col.Default != "" {
				col.Default = "'" + col.Default + "'"
			} else {
				if col.DefaultIsEmpty {
					col.Default = "''"
				}
			}
		}
		cols[col.Name] = col
		colSeq = append(colSeq, col.Name)
	}
	return colSeq, cols, nil
}

func (db *mysql) GetTables() ([]*core.Table, error) {
	args := []interface{}{db.DbName}
	s := "SELECT `TABLE_NAME`, `ENGINE`, `TABLE_ROWS`, `AUTO_INCREMENT` from `INFORMATION_SCHEMA`.`TABLES` WHERE `TABLE_SCHEMA`=?"

	rows, err := db.DB().Query(s, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	tables := make([]*core.Table, 0)
	for rows.Next() {
		table := core.NewEmptyTable()
		var name, engine, tableRows string
		var autoIncr *string
		err = rows.Scan(&name, &engine, &tableRows, &autoIncr)
		if err != nil {
			return nil, err
		}

		table.Name = name
		tables = append(tables, table)
	}
	return tables, nil
}

func (db *mysql) GetIndexes(tableName string) (map[string]*core.Index, error) {
	args := []interface{}{db.DbName, tableName}
	s := "SELECT `INDEX_NAME`, `NON_UNIQUE`, `COLUMN_NAME` FROM `INFORMATION_SCHEMA`.`STATISTICS` WHERE `TABLE_SCHEMA` = ? AND `TABLE_NAME` = ?"

	rows, err := db.DB().Query(s, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	indexes := make(map[string]*core.Index, 0)
	for rows.Next() {
		var indexType int
		var indexName, colName, nonUnique string
		err = rows.Scan(&indexName, &nonUnique, &colName)
		if err != nil {
			return nil, err
		}

		if indexName == "PRIMARY" {
			continue
		}

		if "YES" == nonUnique || nonUnique == "1" {
			indexType = core.IndexType
		} else {
			indexType = core.UniqueType
		}

		colName = strings.Trim(colName, "` ")

		if strings.HasPrefix(indexName, "IDX_"+tableName) || strings.HasPrefix(indexName, "UQE_"+tableName) {
			indexName = indexName[5+len(tableName) : len(indexName)]
		}

		var index *core.Index
		var ok bool
		if index, ok = indexes[indexName]; !ok {
			index = new(core.Index)
			index.Type = indexType
			index.Name = indexName
			indexes[indexName] = index
		}
		index.AddColumn(colName)
	}
	return indexes, nil
}

func (db *mysql) Filters() []core.Filter {
	return []core.Filter{&core.IdFilter{}}
}
<|endoftext|>


<commit_before><commit_msg>Use a 3 minutes timeout for /package/describe<commit_after><|endoftext|>


<commit_before>package knx

import (
	"context"
	"testing"
	"time"
)

var clientConfig = ClientConfig{
	2 * time.Second,
	defaultResendInterval,
	2 * time.Second,
	2 * time.Second,
}

func TestConnHandle_RequestConnection(t *testing.T) {
	ctx := context.Background()

	// Socket was closed before anything could be done.
	t.Run("SendFails", func (t *testing.T) {
		conn := connHandle{makeDummySocket(), clientConfig, 0}
		conn.sock.Close()

		err := conn.requestConnection(ctx)
		if err == nil {
			t.Fatal("Should not succeed")
		}
	})

	// Context is done.
	t.Run("CancelledContext", func (t *testing.T) {
		sock := makeDummySocket()
		defer sock.Close()

		conn := connHandle{sock, clientConfig, 0}

		ctx, cancel := context.WithCancel(ctx)
		cancel()

		err := conn.requestConnection(ctx)
		if err != ctx.Err() {
			t.Fatalf("Expected error %v, got %v", ctx.Err(), err)
		}
	})

	// Socket is closed before first resend.
	t.Run("ResendFails", func (t *testing.T) {
		sock := makeDummySocket()

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}
			gw.ignore()

			sock.closeOut()
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			config := DefaultClientConfig
			config.ResendInterval = 1

			conn := connHandle{sock, config, 0}

			err := conn.requestConnection(ctx)
			if err == nil {
				t.Fatal("Should not succeed")
			}
		})
	})

	// The gateway responds to the connection request.
	t.Run("Resend", func (t *testing.T) {
		sock := makeDummySocket()

		const channel uint8 = 1

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			gw.ignore()

			msg := gw.receive()
			if req, ok := msg.(*ConnectionRequest); ok {
				gw.send(&ConnectionResponse{channel, ConnResOk, req.Control})
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			config := DefaultClientConfig
			config.ResendInterval = 1

			conn := connHandle{sock, config, 0}

			err := conn.requestConnection(ctx)
			if err != nil {
				t.Fatal(err)
			}

			if conn.channel != channel {
				t.Error("Mismatching channel")
			}
		})
	})

	// Inbound channel is closed.
	t.Run("InboundClosed", func (t *testing.T) {
		sock := makeDummySocket()
		sock.closeIn()
		defer sock.Close()

		conn := connHandle{sock, clientConfig, 0}

		err := conn.requestConnection(ctx)
		if err == nil {
			t.Fatal("Should not succeed")
		}
	})

	// The gateway responds to the connection request.
	t.Run("Ok", func (t *testing.T) {
		sock := makeDummySocket()

		const channel uint8 = 1

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if req, ok := msg.(*ConnectionRequest); ok {
				gw.send(&ConnectionResponse{channel, ConnResOk, req.Control})
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			conn := connHandle{sock, clientConfig, 0}

			err := conn.requestConnection(ctx)
			if err != nil {
				t.Fatal(err)
			}

			if conn.channel != channel {
				t.Error("Mismatching channel")
			}
		})
	})

	// The gateway is only busy for the first attempt.
	t.Run("Busy", func (t *testing.T) {
		sock := makeDummySocket()

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if req, ok := msg.(*ConnectionRequest); ok {
				gw.send(&ConnectionResponse{0, ConnResBusy, req.Control})
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}

			msg = gw.receive()
			if req, ok := msg.(*ConnectionRequest); ok {
				gw.send(&ConnectionResponse{1, ConnResOk, req.Control})
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			config := DefaultClientConfig
			config.ResendInterval = 1

			conn := connHandle{sock, config, 0}

			err := conn.requestConnection(ctx)
			if err != nil {
				t.Fatal(err)
			}
		})
	})

	// The gateway doesn't supported the requested connection type.
	t.Run("Unsupported", func (t *testing.T) {
		sock := makeDummySocket()

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if req, ok := msg.(*ConnectionRequest); ok {
				gw.send(&ConnectionResponse{0, ConnResUnsupportedType, req.Control})
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			conn := connHandle{sock, clientConfig, 0}

			err := conn.requestConnection(ctx)
			if err != ConnResUnsupportedType {
				t.Fatalf("Expected error %v, got %v", ConnResUnsupportedType, err)
			}
		})
	})
}

func TestConnHandle_requestConnectionState(t *testing.T) {
	ctx := context.Background()

	t.Run("SendFails", func (t *testing.T) {
		sock := makeDummySocket()
		sock.Close()

		conn := connHandle{sock, clientConfig, 1}

		err := conn.requestConnectionState(ctx, make(chan ConnState))
		if err == nil {
			t.Fatal("Should not succeed")
		}
	})

	t.Run("CancelledContext", func (t *testing.T) {
		sock := makeDummySocket()
		defer sock.Close()

		conn := connHandle{sock, clientConfig, 1}

		ctx, cancel := context.WithCancel(ctx)
		cancel()

		err := conn.requestConnectionState(ctx, make(chan ConnState))
		if err != ctx.Err() {
			t.Fatal("Expected error %v, got %v", ctx.Err(), err)
		}
	})

	t.Run("ResendFails", func (t *testing.T) {
		sock := makeDummySocket()

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}
			gw.ignore()

			sock.closeOut()
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			config := DefaultClientConfig
			config.ResendInterval = 1

			conn := connHandle{sock, config, 1}

			err := conn.requestConnectionState(ctx, make(chan ConnState))
			if err == nil {
				t.Fatal("Should not succeed")
			}
		})
	})

	t.Run("Resend", func (t *testing.T) {
		sock := makeDummySocket()

		const channel uint8 = 1
		heartbeat := make(chan ConnState)

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			gw.ignore()

			msg := gw.receive()
			if req, ok := msg.(*ConnectionStateRequest); ok {
				if req.Channel != channel {
					t.Error("Mismatching channel")
				}

				if req.Status != 0 {
					t.Error("Non-null request status")
				}

				heartbeat <- ConnStateNormal
			} else {
				t.Fatal("Unexpected type %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			config := DefaultClientConfig
			config.ResendInterval = 1

			conn := connHandle{sock, config, channel}

			err := conn.requestConnectionState(ctx, heartbeat)
			if err != nil {
				t.Fatal(err)
			}
		})
	})

	t.Run("InboundClosed", func (t *testing.T) {
		sock := makeDummySocket()

		heartbeat := make(chan ConnState)
		close(heartbeat)

		conn := connHandle{sock, clientConfig, 1}

		err := conn.requestConnectionState(ctx, heartbeat)
		if err == nil {
			t.Fatal("Should not succeed")
		}
	})

	t.Run("Ok", func (t *testing.T) {
		sock := makeDummySocket()

		const channel uint8 = 1
		heartbeat := make(chan ConnState)

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if req, ok := msg.(*ConnectionStateRequest); ok {
				if req.Channel != channel {
					t.Error("Mismatching channel")
				}

				if req.Status != 0 {
					t.Error("Non-null request status")
				}

				heartbeat <- ConnStateNormal
			} else {
				t.Fatal("Unexpected type %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			conn := connHandle{sock, clientConfig, channel}

			err := conn.requestConnectionState(ctx, heartbeat)
			if err != nil {
				t.Fatal(err)
			}
		})
	})

	t.Run("Inactive", func (t *testing.T) {
		sock := makeDummySocket()

		const channel uint8 = 1
		heartbeat := make(chan ConnState)

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if req, ok := msg.(*ConnectionStateRequest); ok {
				if req.Channel != channel {
					t.Error("Mismatching channel")
				}

				if req.Status != 0 {
					t.Error("Non-null request status")
				}

				heartbeat <- ConnStateInactive
			} else {
				t.Fatal("Unexpected type %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			conn := connHandle{sock, clientConfig, channel}

			err := conn.requestConnectionState(ctx, heartbeat)
			if err != ConnStateInactive {
				t.Fatal(err)
			}
		})
	})
}

func TestConnHandle_handleTunnelRequest(t *testing.T) {
	ctx := context.Background()

	// Proper tunnel request from gateway.
	t.Run("Ok", func (t *testing.T) {
		sock := makeDummySocket()
		inbound := make(chan []byte)

		const (
			channel uint8 = 1
			seqNumber uint8 = 0
		)

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if res, ok := msg.(*TunnelResponse); ok {
				if res.Channel != channel {
					t.Error("Mismatching channels")
				}

				if res.SeqNumber != seqNumber {
					t.Error("Mismatching sequence numbers")
				}

				if res.Status != 0 {
					t.Error("Invalid response status")
				}
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}
		})

		t.Run("Worker", func (t *testing.T) {
			t.Parallel()

			var seqNo uint8 = seqNumber
			handle := &connHandle{sock, clientConfig, 1}

			req := &TunnelRequest{channel, seqNumber, []byte{}}
			err := handle.handleTunnelRequest(ctx, req, &seqNo, inbound)
			if err != nil {
				t.Fatal(err)
			}

			if seqNo != req.SeqNumber + 1 {
				t.Error("Sequence number was not increased")
			}
		})

		t.Run("Client", func (t *testing.T) {
			t.Parallel()

			select {
			case <-ctx.Done():
				t.Fatalf("While waiting for inbound packet: %v", ctx.Err())

			case _, open := <-inbound:
				if !open {
					t.Fatal("Inbound channel was closed")
				}
			}
		})
	})

	// Out-of-sequence tunnel request from the gateway.
	t.Run("OutOfSequence", func (t *testing.T) {
		sock := makeDummySocket()
		inbound := make(chan []byte)

		const (
			channel uint8 = 1
			seqNumber uint8 = 1
		)

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if res, ok := msg.(*TunnelResponse); ok {
				if res.Channel != channel {
					t.Error("Mismatching channels")
				}

				if res.SeqNumber != seqNumber {
					t.Error("Mismatching sequence numbers")
				}

				if res.Status != 0 {
					t.Error("Invalid response status")
				}
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}
		})

		t.Run("Worker", func (t *testing.T) {
			t.Parallel()

			var seqNo uint8 = seqNumber - 1
			handle := &connHandle{sock, clientConfig, channel}

			req := &TunnelRequest{channel, seqNumber, []byte{}}
			err := handle.handleTunnelRequest(ctx, req, &seqNo, inbound)
			if err != nil {
				t.Fatal(err)
			}

			if seqNo != seqNumber - 1 {
				t.Error("Sequence number was changed by an out-of-sequence tunnel request")
			}
		})
	})

	// Tunnel request on incorrect channel.
	t.Run("WrongChannel", func (t *testing.T) {
		sock := makeDummySocket()
		defer sock.Close()

		inbound := make(chan []byte)

		const (
			channel uint8 = 1
			seqNumber uint8 = 1
		)

		var seqNo uint8 = 0
		handle := &connHandle{sock, clientConfig, channel + 1}

		req := &TunnelRequest{channel, seqNumber, []byte{}}
		err := handle.handleTunnelRequest(ctx, req, &seqNo, inbound)
		if err == nil {
			t.Fatal("Tunnel request with wrong channel has been successful")
		}

		if seqNo != 0 {
			t.Error("Sequence number was changed by an invalid tunnel request")
		}
	})
}
<commit_msg>Refine tests for connHandle handleTunnelRequest<commit_after>package knx

import (
	"context"
	"testing"
	"time"
)

var clientConfig = ClientConfig{
	2 * time.Second,
	defaultResendInterval,
	2 * time.Second,
	2 * time.Second,
}

func TestConnHandle_RequestConnection(t *testing.T) {
	ctx := context.Background()

	// Socket was closed before anything could be done.
	t.Run("SendFails", func (t *testing.T) {
		conn := connHandle{makeDummySocket(), clientConfig, 0}
		conn.sock.Close()

		err := conn.requestConnection(ctx)
		if err == nil {
			t.Fatal("Should not succeed")
		}
	})

	// Context is done.
	t.Run("CancelledContext", func (t *testing.T) {
		sock := makeDummySocket()
		defer sock.Close()

		conn := connHandle{sock, clientConfig, 0}

		ctx, cancel := context.WithCancel(ctx)
		cancel()

		err := conn.requestConnection(ctx)
		if err != ctx.Err() {
			t.Fatalf("Expected error %v, got %v", ctx.Err(), err)
		}
	})

	// Socket is closed before first resend.
	t.Run("ResendFails", func (t *testing.T) {
		sock := makeDummySocket()

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}
			gw.ignore()

			sock.closeOut()
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			config := DefaultClientConfig
			config.ResendInterval = 1

			conn := connHandle{sock, config, 0}

			err := conn.requestConnection(ctx)
			if err == nil {
				t.Fatal("Should not succeed")
			}
		})
	})

	// The gateway responds to the connection request.
	t.Run("Resend", func (t *testing.T) {
		sock := makeDummySocket()

		const channel uint8 = 1

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			gw.ignore()

			msg := gw.receive()
			if req, ok := msg.(*ConnectionRequest); ok {
				gw.send(&ConnectionResponse{channel, ConnResOk, req.Control})
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			config := DefaultClientConfig
			config.ResendInterval = 1

			conn := connHandle{sock, config, 0}

			err := conn.requestConnection(ctx)
			if err != nil {
				t.Fatal(err)
			}

			if conn.channel != channel {
				t.Error("Mismatching channel")
			}
		})
	})

	// Inbound channel is closed.
	t.Run("InboundClosed", func (t *testing.T) {
		sock := makeDummySocket()
		sock.closeIn()
		defer sock.Close()

		conn := connHandle{sock, clientConfig, 0}

		err := conn.requestConnection(ctx)
		if err == nil {
			t.Fatal("Should not succeed")
		}
	})

	// The gateway responds to the connection request.
	t.Run("Ok", func (t *testing.T) {
		sock := makeDummySocket()

		const channel uint8 = 1

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if req, ok := msg.(*ConnectionRequest); ok {
				gw.send(&ConnectionResponse{channel, ConnResOk, req.Control})
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			conn := connHandle{sock, clientConfig, 0}

			err := conn.requestConnection(ctx)
			if err != nil {
				t.Fatal(err)
			}

			if conn.channel != channel {
				t.Error("Mismatching channel")
			}
		})
	})

	// The gateway is only busy for the first attempt.
	t.Run("Busy", func (t *testing.T) {
		sock := makeDummySocket()

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if req, ok := msg.(*ConnectionRequest); ok {
				gw.send(&ConnectionResponse{0, ConnResBusy, req.Control})
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}

			msg = gw.receive()
			if req, ok := msg.(*ConnectionRequest); ok {
				gw.send(&ConnectionResponse{1, ConnResOk, req.Control})
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			config := DefaultClientConfig
			config.ResendInterval = 1

			conn := connHandle{sock, config, 0}

			err := conn.requestConnection(ctx)
			if err != nil {
				t.Fatal(err)
			}
		})
	})

	// The gateway doesn't supported the requested connection type.
	t.Run("Unsupported", func (t *testing.T) {
		sock := makeDummySocket()

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if req, ok := msg.(*ConnectionRequest); ok {
				gw.send(&ConnectionResponse{0, ConnResUnsupportedType, req.Control})
			} else {
				t.Fatalf("Unexpected incoming message type: %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			conn := connHandle{sock, clientConfig, 0}

			err := conn.requestConnection(ctx)
			if err != ConnResUnsupportedType {
				t.Fatalf("Expected error %v, got %v", ConnResUnsupportedType, err)
			}
		})
	})
}

func TestConnHandle_requestConnectionState(t *testing.T) {
	ctx := context.Background()

	t.Run("SendFails", func (t *testing.T) {
		sock := makeDummySocket()
		sock.Close()

		conn := connHandle{sock, clientConfig, 1}

		err := conn.requestConnectionState(ctx, make(chan ConnState))
		if err == nil {
			t.Fatal("Should not succeed")
		}
	})

	t.Run("CancelledContext", func (t *testing.T) {
		sock := makeDummySocket()
		defer sock.Close()

		conn := connHandle{sock, clientConfig, 1}

		ctx, cancel := context.WithCancel(ctx)
		cancel()

		err := conn.requestConnectionState(ctx, make(chan ConnState))
		if err != ctx.Err() {
			t.Fatal("Expected error %v, got %v", ctx.Err(), err)
		}
	})

	t.Run("ResendFails", func (t *testing.T) {
		sock := makeDummySocket()

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}
			gw.ignore()

			sock.closeOut()
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			config := DefaultClientConfig
			config.ResendInterval = 1

			conn := connHandle{sock, config, 1}

			err := conn.requestConnectionState(ctx, make(chan ConnState))
			if err == nil {
				t.Fatal("Should not succeed")
			}
		})
	})

	t.Run("Resend", func (t *testing.T) {
		sock := makeDummySocket()

		const channel uint8 = 1
		heartbeat := make(chan ConnState)

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			gw.ignore()

			msg := gw.receive()
			if req, ok := msg.(*ConnectionStateRequest); ok {
				if req.Channel != channel {
					t.Error("Mismatching channel")
				}

				if req.Status != 0 {
					t.Error("Non-null request status")
				}

				heartbeat <- ConnStateNormal
			} else {
				t.Fatal("Unexpected type %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			config := DefaultClientConfig
			config.ResendInterval = 1

			conn := connHandle{sock, config, channel}

			err := conn.requestConnectionState(ctx, heartbeat)
			if err != nil {
				t.Fatal(err)
			}
		})
	})

	t.Run("InboundClosed", func (t *testing.T) {
		sock := makeDummySocket()

		heartbeat := make(chan ConnState)
		close(heartbeat)

		conn := connHandle{sock, clientConfig, 1}

		err := conn.requestConnectionState(ctx, heartbeat)
		if err == nil {
			t.Fatal("Should not succeed")
		}
	})

	t.Run("Ok", func (t *testing.T) {
		sock := makeDummySocket()

		const channel uint8 = 1
		heartbeat := make(chan ConnState)

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if req, ok := msg.(*ConnectionStateRequest); ok {
				if req.Channel != channel {
					t.Error("Mismatching channel")
				}

				if req.Status != 0 {
					t.Error("Non-null request status")
				}

				heartbeat <- ConnStateNormal
			} else {
				t.Fatal("Unexpected type %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			conn := connHandle{sock, clientConfig, channel}

			err := conn.requestConnectionState(ctx, heartbeat)
			if err != nil {
				t.Fatal(err)
			}
		})
	})

	t.Run("Inactive", func (t *testing.T) {
		sock := makeDummySocket()

		const channel uint8 = 1
		heartbeat := make(chan ConnState)

		t.Run("Gateway", func (t *testing.T) {
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if req, ok := msg.(*ConnectionStateRequest); ok {
				if req.Channel != channel {
					t.Error("Mismatching channel")
				}

				if req.Status != 0 {
					t.Error("Non-null request status")
				}

				heartbeat <- ConnStateInactive
			} else {
				t.Fatal("Unexpected type %T", msg)
			}
		})

		t.Run("Client", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			conn := connHandle{sock, clientConfig, channel}

			err := conn.requestConnectionState(ctx, heartbeat)
			if err != ConnStateInactive {
				t.Fatal(err)
			}
		})
	})
}

func TestConnHandle_handleTunnelRequest(t *testing.T) {
	ctx := context.Background()

	t.Run("InvalidChannel", func (t *testing.T) {
		sock := makeDummySocket()
		defer sock.Close()

		var seqNumber uint8 = 0

		conn := connHandle{sock, clientConfig, 1}
		req := &TunnelRequest{2, 0, []byte{}}

		err := conn.handleTunnelRequest(ctx, req, &seqNumber, make(chan []byte))
		if err == nil {
			t.Fatal("Should not succeed")
		}
	})

	t.Run("InvalidSeqNumber", func (t *testing.T) {
		sock := makeDummySocket()

		const (
			channel       uint8 = 1
			sendSeqNumber uint8 = 0
		)

		t.Run("Gateway", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if res, ok := msg.(*TunnelResponse); ok {
				if res.Channel != channel {
					t.Error("Mismatching channel")
				}

				if res.SeqNumber != sendSeqNumber {
					t.Error("Mismatching sequence number")
				}

				if res.Status != 0 {
					t.Error("Invalid response status")
				}
			} else {
				t.Fatal("Unexpected type %T", msg)
			}
		})

		t.Run("Worker", func (t *testing.T) {
			t.Parallel()

			var seqNumber uint8 = sendSeqNumber + 1

			conn := connHandle{sock, clientConfig, channel}
			req := &TunnelRequest{channel, sendSeqNumber, []byte{}}

			err := conn.handleTunnelRequest(ctx, req, &seqNumber, make(chan []byte))
			if err != nil {
				t.Fatal(err)
			}

			if seqNumber != sendSeqNumber + 1 {
				t.Error("Sequence number was modified")
			}
		})
	})

	t.Run("Ok", func (t *testing.T) {
		sock := makeDummySocket()
		inbound := make(chan []byte)

		const (
			channel       uint8 = 1
			sendSeqNumber uint8 = 0
		)

		t.Run("Gateway", func (t *testing.T) {
			defer sock.Close()
			t.Parallel()

			gw := gatewayHelper{ctx, sock, t}

			msg := gw.receive()
			if res, ok := msg.(*TunnelResponse); ok {
				if res.Channel != channel {
					t.Error("Mismatching channel")
				}

				if res.SeqNumber != sendSeqNumber {
					t.Error("Mismatching sequence number")
				}

				if res.Status != 0 {
					t.Error("Invalid response status")
				}
			} else {
				t.Fatal("Unexpected type %T", msg)
			}
		})

		t.Run("Worker", func (t *testing.T) {
			t.Parallel()

			var seqNumber uint8 = sendSeqNumber

			conn := connHandle{sock, clientConfig, channel}
			req := &TunnelRequest{channel, sendSeqNumber, []byte{}}

			err := conn.handleTunnelRequest(ctx, req, &seqNumber, inbound)
			if err != nil {
				t.Fatal(err)
			}

			if seqNumber != sendSeqNumber + 1 {
				t.Error("Sequence number has not been increased")
			}
		})

		t.Run("Client", func (t *testing.T) {
			t.Parallel()

			<-inbound
		})
	})
}
<|endoftext|>


<commit_before>package bandcamp

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"reflect"
	"time"

	"github.com/mjibson/mog/codec"
	"github.com/mjibson/mog/codec/mpa"
	"github.com/mjibson/mog/protocol"
	"golang.org/x/oauth2"
)

func init() {
	protocol.Register("bandcamp", []string{"URL"}, New, reflect.TypeOf(&Bandcamp{}))
	gob.Register(new(Bandcamp))
}

func New(params []string, token *oauth2.Token) (protocol.Instance, error) {
	if len(params) != 1 {
		return nil, fmt.Errorf("expected one parameter")
	}
	b := &Bandcamp{
		URL: params[0],
	}
	if _, err := b.Refresh(); err != nil {
		return nil, err
	}
	return b, nil
}

type Bandcamp struct {
	URL    string
	Songs  protocol.SongList
	Tracks map[codec.ID]*track
}

func (b *Bandcamp) Key() string {
	return b.URL
}

func (b *Bandcamp) Info(id codec.ID) (*codec.SongInfo, error) {
	t := b.Songs[id]
	if t == nil {
		return nil, fmt.Errorf("could not find %v", id)
	}
	return t, nil
}

func (b *Bandcamp) GetSong(id codec.ID) (codec.Song, error) {
	t := b.Tracks[id]
	if t == nil {
		return nil, fmt.Errorf("missing %v", id)
	}
	return mpa.NewSong(func() (io.ReadCloser, int64, error) {
		log.Println("BANDCAMP", id)
		res, err := http.Get(t.File.Mp3_128)
		if err != nil {
			return nil, 0, err
		}
		return res.Body, 0, nil
	})
}

func (b *Bandcamp) List() (protocol.SongList, error) {
	if len(b.Songs) == 0 {
		return b.Refresh()
	}
	return b.Songs, nil
}

type track struct {
	Duration float64 `json:"duration"`
	File     struct {
		Mp3_128 string `json:"mp3-128"`
	} `json:"file"`
	ID       int64  `json:"id"`
	Title    string `json:"title"`
	TrackNum int    `json:"track_num"`
}

type album struct {
	Artist string `json:"artist"`
	Title  string `json:"title"`
}

func (b *Bandcamp) Refresh() (protocol.SongList, error) {
	resp, err := http.Get(b.URL)
	if err != nil {
		return nil, err
	}
	y, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return nil, err
	}
	var tracks []*track
	var a album
	{
		const start = "    trackinfo : ["
		i := bytes.Index(y, []byte(start))
		if i < 0 {
			return nil, fmt.Errorf("not a bandcamp site: %v", b.URL)
		}
		e := bytes.Index(y[i:], []byte("],\n"))
		if e < 0 {
			return nil, fmt.Errorf("not a bandcamp site: %v", b.URL)
		}
		if err := json.Unmarshal(y[i+len(start)-1:e+1+i], &tracks); err != nil {
			return nil, err
		}
	}
	{
		const start = "    current: {"
		i := bytes.Index(y, []byte(start))
		if i < 0 {
			return nil, fmt.Errorf("not a bandcamp site: %v", b.URL)
		}
		e := bytes.Index(y[i:], []byte("},\n"))
		if e < 0 {
			return nil, fmt.Errorf("not a bandcamp site: %v", b.URL)
		}
		if err := json.Unmarshal(y[i+len(start)-1:e+1+i], &a); err != nil {
			return nil, err
		}
	}
	var artURL string
	func() {
		const start = `    artThumbURL: "`
		i := bytes.Index(y, []byte(start))
		if i < 0 {
			return
		}
		e := bytes.Index(y[i:], []byte("\",\n"))
		if e < 0 {
			return
		}
		if err := json.Unmarshal(y[i+len(start)-1:e+1+i], &artURL); err != nil {
			return
		}
	}()
	tracklist := make(map[codec.ID]*track)
	songs := make(protocol.SongList)
	for _, t := range tracks {
		if t.File.Mp3_128 == "" {
			continue
		}
		id := codec.Int64(t.ID)
		tracklist[id] = t
		songs[id] = &codec.SongInfo{
			Time:     time.Duration(t.Duration) * time.Second,
			Album:    a.Title,
			Artist:   a.Artist,
			Title:    t.Title,
			Track:    float64(t.TrackNum),
			ImageURL: artURL,
		}
	}
	b.Songs = songs
	b.Tracks = tracklist
	return songs, err
}
<commit_msg>Fix bandcamp URLs<commit_after>package bandcamp

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"reflect"
	"strings"
	"time"

	"github.com/mjibson/mog/codec"
	"github.com/mjibson/mog/codec/mpa"
	"github.com/mjibson/mog/protocol"
	"golang.org/x/oauth2"
)

func init() {
	protocol.Register("bandcamp", []string{"URL"}, New, reflect.TypeOf(&Bandcamp{}))
	gob.Register(new(Bandcamp))
}

func New(params []string, token *oauth2.Token) (protocol.Instance, error) {
	if len(params) != 1 {
		return nil, fmt.Errorf("expected one parameter")
	}
	b := &Bandcamp{
		URL: params[0],
	}
	if _, err := b.Refresh(); err != nil {
		return nil, err
	}
	return b, nil
}

type Bandcamp struct {
	URL    string
	Songs  protocol.SongList
	Tracks map[codec.ID]*track
}

func (b *Bandcamp) Key() string {
	return b.URL
}

func (b *Bandcamp) Info(id codec.ID) (*codec.SongInfo, error) {
	t := b.Songs[id]
	if t == nil {
		return nil, fmt.Errorf("could not find %v", id)
	}
	return t, nil
}

func (b *Bandcamp) GetSong(id codec.ID) (codec.Song, error) {
	t := b.Tracks[id]
	if t == nil {
		return nil, fmt.Errorf("missing %v", id)
	}
	return mpa.NewSong(func() (io.ReadCloser, int64, error) {
		log.Println("BANDCAMP", id)
		res, err := http.Get(t.File.Mp3_128)
		if err != nil {
			return nil, 0, err
		}
		return res.Body, 0, nil
	})
}

func (b *Bandcamp) List() (protocol.SongList, error) {
	if len(b.Songs) == 0 {
		return b.Refresh()
	}
	return b.Songs, nil
}

type track struct {
	Duration float64 `json:"duration"`
	File     struct {
		Mp3_128 string `json:"mp3-128"`
	} `json:"file"`
	ID       int64  `json:"id"`
	Title    string `json:"title"`
	TrackNum int    `json:"track_num"`
}

type album struct {
	Artist string `json:"artist"`
	Title  string `json:"title"`
}

func (b *Bandcamp) Refresh() (protocol.SongList, error) {
	resp, err := http.Get(b.URL)
	if err != nil {
		return nil, err
	}
	y, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return nil, err
	}
	var tracks []*track
	var a album
	{
		const start = "    trackinfo : ["
		i := bytes.Index(y, []byte(start))
		if i < 0 {
			return nil, fmt.Errorf("not a bandcamp site: %v", b.URL)
		}
		e := bytes.Index(y[i:], []byte("],\n"))
		if e < 0 {
			return nil, fmt.Errorf("not a bandcamp site: %v", b.URL)
		}
		if err := json.Unmarshal(y[i+len(start)-1:e+1+i], &tracks); err != nil {
			return nil, err
		}
	}
	{
		const start = "    current: {"
		i := bytes.Index(y, []byte(start))
		if i < 0 {
			return nil, fmt.Errorf("not a bandcamp site: %v", b.URL)
		}
		e := bytes.Index(y[i:], []byte("},\n"))
		if e < 0 {
			return nil, fmt.Errorf("not a bandcamp site: %v", b.URL)
		}
		if err := json.Unmarshal(y[i+len(start)-1:e+1+i], &a); err != nil {
			return nil, err
		}
	}
	var artURL string
	func() {
		const start = `    artThumbURL: "`
		i := bytes.Index(y, []byte(start))
		if i < 0 {
			return
		}
		e := bytes.Index(y[i:], []byte("\",\n"))
		if e < 0 {
			return
		}
		if err := json.Unmarshal(y[i+len(start)-1:e+1+i], &artURL); err != nil {
			return
		}
	}()
	tracklist := make(map[codec.ID]*track)
	songs := make(protocol.SongList)
	for _, t := range tracks {
		if t.File.Mp3_128 == "" {
			continue
		}
		if strings.HasPrefix(t.File.Mp3_128, "//") {
			t.File.Mp3_128 = "https:" + t.File.Mp3_128
		}
		id := codec.Int64(t.ID)
		tracklist[id] = t
		songs[id] = &codec.SongInfo{
			Time:     time.Duration(t.Duration) * time.Second,
			Album:    a.Title,
			Artist:   a.Artist,
			Title:    t.Title,
			Track:    float64(t.TrackNum),
			ImageURL: artURL,
		}
	}
	b.Songs = songs
	b.Tracks = tracklist
	return songs, err
}
<|endoftext|>


<commit_before>// Copyright 2019 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 scorecard

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"cloud.google.com/go/storage"
	"github.com/forseti-security/config-validator/pkg/api/validator"
	cvasset "github.com/forseti-security/config-validator/pkg/asset"
	"github.com/pkg/errors"
)

func getDataFromReader(config *ScoringConfig, reader io.Reader) ([]*validator.Asset, error) {
	const maxCapacity = 1024 * 1024
	scanner := bufio.NewScanner(reader)
	buf := make([]byte, maxCapacity)
	scanner.Buffer(buf, maxCapacity)
	var pbAssets []*validator.Asset
	for scanner.Scan() {
		pbAsset, err := getAssetFromJSON(scanner.Bytes())
		if err != nil {
			return nil, err
		}
		pbAssets = append(pbAssets, pbAsset)
	}
	return pbAssets, nil
}

func getDataFromBucket(config *ScoringConfig, bucketName string) ([]*validator.Asset, error) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		return nil, err
	}

	bucket := client.Bucket(bucketName)
	var pbAssets []*validator.Asset
	for _, objectName := range destinationObjectNames {
		reader, err := bucket.Object(objectName).NewReader(ctx)
		if err != nil {
			fmt.Println("WARNING: Unable to read inventory file :", objectName, err)
			continue
		}
		defer reader.Close()
		assets, err := getDataFromReader(config, reader)
		if err != nil {
			return nil, err
		}

		pbAssets = append(pbAssets, assets...)
	}
	if len(pbAssets) == 0 {
		return nil, fmt.Errorf("No inventory found")
	}
	return pbAssets, nil
}

func getDataFromFile(config *ScoringConfig, caiDirName string) ([]*validator.Asset, error) {
	var pbAssets []*validator.Asset
	for _, objectName := range destinationObjectNames {
		reader, err := os.Open(filepath.Join(caiDirName, objectName))
		if err != nil {
			fmt.Println("WARNING: Unable to read inventory file :", objectName, err)
			continue
		}
		defer reader.Close()
		assets, err := getDataFromReader(config, reader)
		if err != nil {
			return nil, err
		}

		pbAssets = append(pbAssets, assets...)
	}
	if len(pbAssets) == 0 {
		return nil, fmt.Errorf("No inventory found")
	}
	return pbAssets, nil
}

func getDataFromStdin(config *ScoringConfig) ([]*validator.Asset, error) {
	return getDataFromReader(config, os.Stdin)
}

// getViolations finds all Config Validator violations for a given Inventory
func getViolations(inventory *InventoryConfig, config *ScoringConfig) (*validator.AuditResponse, error) {
	var err error
	var pbAssets []*validator.Asset
	if inventory.bucketName != "" {
		pbAssets, err = getDataFromBucket(config, inventory.bucketName)
		if err != nil {
			return nil, errors.Wrap(err, "Fetching inventory from Bucket")
		}
	} else if inventory.dirPath != "" {
		pbAssets, err = getDataFromFile(config, inventory.dirPath)
		if err != nil {
			return nil, errors.Wrap(err, "Fetching inventory from local directory")
		}
	} else if inventory.readFromStdin {
		pbAssets, err = getDataFromStdin(config)
		if err != nil {
			return nil, errors.Wrap(err, "Reading from stdin")
		}
	}

	auditResult := &validator.AuditResponse{}
	for _, asset := range pbAssets {
		violations, err := config.validator.ReviewAsset(context.Background(), asset)

		if err != nil {
			return nil, errors.Wrapf(err, "reviewing asset %s", asset)
		}
		auditResult.Violations = append(auditResult.Violations, violations...)
	}
	return auditResult, nil
}

// converts raw JSON into Asset proto
func getAssetFromJSON(input []byte) (*validator.Asset, error) {
	pbAsset := &validator.Asset{}
	err := unMarshallAsset(input, pbAsset)
	if err != nil {
		return nil, errors.Wrapf(err, "converting asset %s to proto", string(input))
	}

	err = cvasset.SanitizeAncestryPath(pbAsset)
	if err != nil {
		return nil, errors.Wrapf(err, "fetching ancestry path for %s", pbAsset)
	}

	Log.Debug("Asset converted", "name", pbAsset.GetName(), "ancestry", pbAsset.GetAncestryPath())
	return pbAsset, nil
}
<commit_msg>fix(cli): Remove unused config parameter in violations.go (#829)<commit_after>// Copyright 2019 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 scorecard

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"cloud.google.com/go/storage"
	"github.com/forseti-security/config-validator/pkg/api/validator"
	cvasset "github.com/forseti-security/config-validator/pkg/asset"
	"github.com/pkg/errors"
)

func getDataFromReader(reader io.Reader) ([]*validator.Asset, error) {
	const maxCapacity = 1024 * 1024
	scanner := bufio.NewScanner(reader)
	buf := make([]byte, maxCapacity)
	scanner.Buffer(buf, maxCapacity)
	var pbAssets []*validator.Asset
	for scanner.Scan() {
		pbAsset, err := getAssetFromJSON(scanner.Bytes())
		if err != nil {
			return nil, err
		}
		pbAssets = append(pbAssets, pbAsset)
	}
	return pbAssets, nil
}

func getDataFromBucket(bucketName string) ([]*validator.Asset, error) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		return nil, err
	}

	bucket := client.Bucket(bucketName)
	var pbAssets []*validator.Asset
	for _, objectName := range destinationObjectNames {
		reader, err := bucket.Object(objectName).NewReader(ctx)
		if err != nil {
			fmt.Println("WARNING: Unable to read inventory file :", objectName, err)
			continue
		}
		defer reader.Close()
		assets, err := getDataFromReader(reader)
		if err != nil {
			return nil, err
		}

		pbAssets = append(pbAssets, assets...)
	}
	if len(pbAssets) == 0 {
		return nil, fmt.Errorf("No inventory found")
	}
	return pbAssets, nil
}

func getDataFromFile(caiDirName string) ([]*validator.Asset, error) {
	var pbAssets []*validator.Asset
	for _, objectName := range destinationObjectNames {
		reader, err := os.Open(filepath.Join(caiDirName, objectName))
		if err != nil {
			fmt.Println("WARNING: Unable to read inventory file :", objectName, err)
			continue
		}
		defer reader.Close()
		assets, err := getDataFromReader(reader)
		if err != nil {
			return nil, err
		}

		pbAssets = append(pbAssets, assets...)
	}
	if len(pbAssets) == 0 {
		return nil, fmt.Errorf("No inventory found")
	}
	return pbAssets, nil
}

func getDataFromStdin() ([]*validator.Asset, error) {
	return getDataFromReader(os.Stdin)
}

// getViolations finds all Config Validator violations for a given Inventory
func getViolations(inventory *InventoryConfig, config *ScoringConfig) (*validator.AuditResponse, error) {
	var err error
	var pbAssets []*validator.Asset
	if inventory.bucketName != "" {
		pbAssets, err = getDataFromBucket(inventory.bucketName)
		if err != nil {
			return nil, errors.Wrap(err, "Fetching inventory from Bucket")
		}
	} else if inventory.dirPath != "" {
		pbAssets, err = getDataFromFile(inventory.dirPath)
		if err != nil {
			return nil, errors.Wrap(err, "Fetching inventory from local directory")
		}
	} else if inventory.readFromStdin {
		pbAssets, err = getDataFromStdin()
		if err != nil {
			return nil, errors.Wrap(err, "Reading from stdin")
		}
	}

	auditResult := &validator.AuditResponse{}
	for _, asset := range pbAssets {
		violations, err := config.validator.ReviewAsset(context.Background(), asset)

		if err != nil {
			return nil, errors.Wrapf(err, "reviewing asset %s", asset)
		}
		auditResult.Violations = append(auditResult.Violations, violations...)
	}
	return auditResult, nil
}

// converts raw JSON into Asset proto
func getAssetFromJSON(input []byte) (*validator.Asset, error) {
	pbAsset := &validator.Asset{}
	err := unMarshallAsset(input, pbAsset)
	if err != nil {
		return nil, errors.Wrapf(err, "converting asset %s to proto", string(input))
	}

	err = cvasset.SanitizeAncestryPath(pbAsset)
	if err != nil {
		return nil, errors.Wrapf(err, "fetching ancestry path for %s", pbAsset)
	}

	Log.Debug("Asset converted", "name", pbAsset.GetName(), "ancestry", pbAsset.GetAncestryPath())
	return pbAsset, nil
}
<|endoftext|>


<commit_before>package helper

import (
	"encoding/json"
	"testing"
)

// Asserts that the request did not return an error.
// Optionally perform some checks only if the request did not fail
func AssertRequestOk(t *testing.T, response interface{}, err error, check_fn func()) {
	if err != nil {
		response_json, _ := json.MarshalIndent(response, "", "  ")
		errorPayload, _ := json.MarshalIndent(err, "", " ")
		t.Fatalf("Failed to perform request, because %s. Response:\n%s", errorPayload, response_json)
	} else {
		if check_fn != nil {
			check_fn()
		}
	}
}

// Asserts that the request _did_ return an error.
// Optionally perform some checks only if the request failed
func AssertRequestFail(t *testing.T, response interface{}, err error, check_fn func()) {
	if err == nil {
		response_json, _ := json.MarshalIndent(response, "", "  ")
		t.Fatalf("Request succeeded unexpectedly. Response:\n%s", response_json)
	} else {
		if check_fn != nil {
			check_fn()
		}
	}
}
<commit_msg>Improve acceptance test assertion by printing out error type<commit_after>package helper

import (
	"encoding/json"
	"reflect"
	"testing"
)

// Asserts that the request did not return an error.
// Optionally perform some checks only if the request did not fail
func AssertRequestOk(t *testing.T, response interface{}, err error, check_fn func()) {
	if err != nil {
		response_json, _ := json.MarshalIndent(response, "", "  ")
		errorPayload, _ := json.MarshalIndent(err, "", " ")
		t.Fatalf("Failed to perform request! Error: %s %s. Response: %s", getType(err), errorPayload, response_json)
	} else {
		if check_fn != nil {
			check_fn()
		}
	}
}

// Asserts that the request _did_ return an error.
// Optionally perform some checks only if the request failed
func AssertRequestFail(t *testing.T, response interface{}, err error, check_fn func()) {
	if err == nil {
		response_json, _ := json.MarshalIndent(response, "", "  ")
		t.Fatalf("Request succeeded unexpectedly. Response:\n%s", response_json)
	} else {
		if check_fn != nil {
			check_fn()
		}
	}
}

// Get type name of some value, according to https://stackoverflow.com/questions/35790935/using-reflection-in-go-to-get-the-name-of-a-struct
func getType(myvar interface{}) string {
	if t := reflect.TypeOf(myvar); t.Kind() == reflect.Ptr {
		return "*" + t.Elem().Name()
	} else {
		return t.Name()
	}
}
<|endoftext|>


<commit_before>package network

type Networks struct {
	Networks []Network `json:"networks,omitempty"`
}

// TODO add subnetworks
type Network struct {
	Status       string `json:"status,omitempty"`
	Name         string `json:"name,omitempty"`
	AdminStateUp bool   `json:"admin_state_up,omitempty"`
	TenantId     string `json:"tenant_id,omitempty"`
	Id           string `json:"id,omitempty"`
}

type NewNetwork struct {
	Name         string `json:"name,omitempty"`
	AdminStateUp bool   `json:"admin_state_up,omitempty"`
	Shared       bool   `json:"shared,omitempty"`
	TenantId     string `json:"tenant_id,omitempty"`
}

type Subnet struct {
	NetworkId  string `json:"network_id,omitempty"`
	Name       string `json:"name,omitempty"`
	TenantId   string `json:"tenant_id,omitempty"`
	Cidr       string `json:"cidr,omitempty"`
	IPVersion  int    `json:"ip_version,omitempty"`
	EnableDhcp bool   `json:"enable_dhcp,omitempty"`
	Id         string `json:"id,omitempty"`
}

type SecurityGroup struct {
	Id          string `json:"id,omitempty"`
	Name        string `json:"name,omitempty"`
	TenantId    string `json:"tenant_id,omitempty"`
	Description string `json:"description,omitempty"`
}

type NewSecurityGroup struct {
	Id          string `json:"id,omitempty"`
	Name        string `json:"name,omitempty"`
	TenantId    string `json:"tenant_id,omitempty"`
	Description string `json:"description,omitempty"`
}

type SecurityGroupRule struct {
	Id              string `json:"id,omitempty"`
	TenantId        string `json:"tenant_id,omitempty"`
	Direction       string `json:"direction,omitempty"`
	PortRangeMin    int    `json:"port_range_min,omitempty"`
	PortRangeMax    int    `json:"port_range_max,omitempty"`
	Protocol        string `json:"protocol,omitempty"`
	RemoteIpPrefix  string `json:"remote_ip_prefix,omitempty"`
	SecurityGroupId string `json:"security_group_id,omitempty"`
}
<commit_msg>Add subnetworks and rules on network and securityGroup struct<commit_after>package network

type Networks struct {
	Networks []Network `json:"networks,omitempty"`
}

type Network struct {
	Status       string   `json:"status,omitempty"`
	Name         string   `json:"name,omitempty"`
	AdminStateUp bool     `json:"admin_state_up,omitempty"`
	TenantId     string   `json:"tenant_id,omitempty"`
	Subnets      []string `json:"subnets,omitempty"`
	Id           string   `json:"id,omitempty"`
}

type NewNetwork struct {
	Name         string `json:"name,omitempty"`
	AdminStateUp bool   `json:"admin_state_up,omitempty"`
	Shared       bool   `json:"shared,omitempty"`
	TenantId     string `json:"tenant_id,omitempty"`
}

type Subnet struct {
	NetworkId  string `json:"network_id,omitempty"`
	Name       string `json:"name,omitempty"`
	TenantId   string `json:"tenant_id,omitempty"`
	Cidr       string `json:"cidr,omitempty"`
	IPVersion  int    `json:"ip_version,omitempty"`
	EnableDhcp bool   `json:"enable_dhcp,omitempty"`
	Id         string `json:"id,omitempty"`
}

type SecurityGroup struct {
	Id          string              `json:"id,omitempty"`
	Name        string              `json:"name,omitempty"`
	TenantId    string              `json:"tenant_id,omitempty"`
	Description string              `json:"description,omitempty"`
	Rules       []SecurityGroupRule `json:"security_group_rules,omitempty"`
}

type NewSecurityGroup struct {
	Id          string `json:"id,omitempty"`
	Name        string `json:"name,omitempty"`
	TenantId    string `json:"tenant_id,omitempty"`
	Description string `json:"description,omitempty"`
}

type SecurityGroupRule struct {
	Id              string `json:"id,omitempty"`
	TenantId        string `json:"tenant_id,omitempty"`
	Direction       string `json:"direction,omitempty"`
	PortRangeMin    int    `json:"port_range_min,omitempty"`
	PortRangeMax    int    `json:"port_range_max,omitempty"`
	Protocol        string `json:"protocol,omitempty"`
	RemoteIpPrefix  string `json:"remote_ip_prefix,omitempty"`
	SecurityGroupId string `json:"security_group_id,omitempty"`
}
<|endoftext|>