<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package discovery

import (
	"fmt"
	"path/filepath"
	"testing"

	"github.com/kubernetes-incubator/external-storage/local-volume/provisioner/pkg/cache"
	"github.com/kubernetes-incubator/external-storage/local-volume/provisioner/pkg/common"
	"github.com/kubernetes-incubator/external-storage/local-volume/provisioner/pkg/util"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/pkg/api/v1"
	v1helper "k8s.io/client-go/pkg/api/v1/helper"
)

const (
	testHostDir         = "/mnt/disks"
	testMountDir        = "/discoveryPath"
	testNodeName        = "test-node"
	testProvisionerName = "test-provisioner"
)

var testNode = &v1.Node{
	ObjectMeta: metav1.ObjectMeta{
		Name: testNodeName,
		Labels: map[string]string{
			common.NodeLabelKey: testNodeName,
		},
	},
}

var scMapping = map[string]string{
	"sc1": "dir1",
	"sc2": "dir2",
}

type testConfig struct {
	// The directory layout for the test
	// Key = directory, Value = list of volumes under that directory
	dirLayout map[string][]*util.FakeFile
	// The volumes that are expected to be created as PVs
	// Key = directory, Value = list of volumes under that directory
	expectedVolumes map[string][]*util.FakeFile
	// True if testing api failure
	apiShouldFail bool
	// The rest are set during setup
	volUtil *util.FakeVolumeUtil
	apiUtil *util.FakeAPIUtil
	cache   *cache.VolumeCache
}

func TestDiscoverVolumes_Basic(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount1", Hash: 0xaaaafef5, Capacity: 100 * 1024},
			{Name: "mount2", Hash: 0x79412c38, Capacity: 100 * 1024 * 1024},
		},
		"dir2": {
			{Name: "mount1", Hash: 0xa7aafa3c},
			{Name: "mount2", Hash: 0x7c4130f1},
		},
	}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: vols,
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()
	verifyCreatedPVs(t, test)
}

func TestDiscoverVolumes_BasicTwice(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount1", Hash: 0xaaaafef5},
			{Name: "mount2", Hash: 0x79412c38},
		},
		"dir2": {
			{Name: "mount1", Hash: 0xa7aafa3c},
			{Name: "mount2", Hash: 0x7c4130f1},
		},
	}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: vols,
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()
	verifyCreatedPVs(t, test)

	// Second time should not create any new volumes
	test.expectedVolumes = map[string][]*util.FakeFile{}
	d.DiscoverLocalVolumes()
	verifyCreatedPVs(t, test)
}

func TestDiscoverVolumes_NoDir(t *testing.T) {
	vols := map[string][]*util.FakeFile{}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: vols,
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()
	verifyCreatedPVs(t, test)
}

func TestDiscoverVolumes_EmptyDir(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {},
	}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: vols,
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()
	verifyCreatedPVs(t, test)
}

func TestDiscoverVolumes_NewVolumesLater(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount1", Hash: 0xaaaafef5},
			{Name: "mount2", Hash: 0x79412c38},
		},
		"dir2": {
			{Name: "mount1", Hash: 0xa7aafa3c},
			{Name: "mount2", Hash: 0x7c4130f1},
		},
	}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: vols,
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()

	verifyCreatedPVs(t, test)

	// Some new mount points show up
	newVols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount3", Hash: 0xf34b8003},
			{Name: "mount4", Hash: 0x144e29de},
		},
	}
	test.volUtil.AddNewFiles(testMountDir, newVols)
	test.expectedVolumes = newVols

	d.DiscoverLocalVolumes()

	verifyCreatedPVs(t, test)
}

func TestDiscoverVolumes_CreatePVFails(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount1", Hash: 0xaaaafef5},
			{Name: "mount2", Hash: 0x79412c38},
		},
		"dir2": {
			{Name: "mount1", Hash: 0xa7aafa3c},
			{Name: "mount2", Hash: 0x7c4130f1},
		},
	}
	test := &testConfig{
		apiShouldFail:   true,
		dirLayout:       vols,
		expectedVolumes: map[string][]*util.FakeFile{},
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()

	verifyCreatedPVs(t, test)
	verifyPVsNotInCache(t, test)
}

func TestDiscoverVolumes_BadVolume(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount1", IsNotDir: true},
		},
	}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: map[string][]*util.FakeFile{},
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()

	verifyCreatedPVs(t, test)
	verifyPVsNotInCache(t, test)
}

func testSetup(t *testing.T, test *testConfig) *Discoverer {
	test.cache = cache.NewVolumeCache()
	test.volUtil = util.NewFakeVolumeUtil(false)
	test.volUtil.AddNewFiles(testMountDir, test.dirLayout)
	test.apiUtil = util.NewFakeAPIUtil(test.apiShouldFail, test.cache)

	userConfig := &common.UserConfig{
		Node:         testNode,
		MountDir:     testMountDir,
		HostDir:      testHostDir,
		DiscoveryMap: scMapping,
	}
	runConfig := &common.RuntimeConfig{
		UserConfig: userConfig,
		Cache:      test.cache,
		VolUtil:    test.volUtil,
		APIUtil:    test.apiUtil,
		Name:       testProvisionerName,
	}
	d, err := NewDiscoverer(runConfig)
	if err != nil {
		t.Fatalf("Error setting up test discoverer: %v", err)
	}
	return d
}

func findSCName(t *testing.T, targetDir string, test *testConfig) string {
	for sc, dir := range scMapping {
		if dir == targetDir {
			return sc
		}
	}
	t.Fatalf("Failed to find SC Name for directory %v", targetDir)
	return ""
}

func verifyNodeAffinity(t *testing.T, pv *v1.PersistentVolume) {
	affinity, err := v1helper.GetStorageNodeAffinityFromAnnotation(pv.Annotations)
	if err != nil {
		t.Errorf("Could not get node affinity from annotation: %v", err)
		return
	}
	if affinity == nil {
		t.Errorf("No node affinity found")
		return
	}

	selector := affinity.RequiredDuringSchedulingIgnoredDuringExecution
	if selector == nil {
		t.Errorf("NodeAffinity node selector is nil")
		return
	}
	terms := selector.NodeSelectorTerms
	if len(terms) != 1 {
		t.Errorf("Node selector term count is %v, expected 1", len(terms))
		return
	}
	reqs := terms[0].MatchExpressions
	if len(reqs) != 1 {
		t.Errorf("Node selector term requirements count is %v, expected 1", len(reqs))
		return
	}

	req := reqs[0]
	if req.Key != common.NodeLabelKey {
		t.Errorf("Node selector requirement key is %v, expected %v", req.Key, common.NodeLabelKey)
	}
	if req.Operator != v1.NodeSelectorOpIn {
		t.Errorf("Node selector requirement operator is %v, expected %v", req.Operator, v1.NodeSelectorOpIn)
	}
	if len(req.Values) != 1 {
		t.Errorf("Node selector requirement value count is %v, expected 1", len(req.Values))
		return
	}
	if req.Values[0] != testNodeName {
		t.Errorf("Node selector requirement value is %v, expected %v", req.Values[0], testNodeName)
	}
}

func verifyProvisionerName(t *testing.T, pv *v1.PersistentVolume) {
	if len(pv.Annotations) == 0 {
		t.Errorf("Annotations not set")
		return
	}
	name, found := pv.Annotations[common.AnnProvisionedBy]
	if !found {
		t.Errorf("Provisioned by annotations not set")
		return
	}
	if name != testProvisionerName {
		t.Errorf("Provisioned name is %q, expected %q", name, testProvisionerName)
	}
}

// testPVInfo contains all the fields we are intested in validating.
type testPVInfo struct {
	pvName   string
	path     string
	capacity uint64
}

func verifyCreatedPVs(t *testing.T, test *testConfig) {
	expectedPVs := map[string]*testPVInfo{}
	for dir, files := range test.expectedVolumes {
		for _, file := range files {
			pvName := fmt.Sprintf("local-pv-%x", file.Hash)
			path := filepath.Join(testHostDir, dir, file.Name)
			expectedPVs[pvName] = &testPVInfo{pvName: pvName, path: path, capacity: file.Capacity}
		}
	}

	createdPVs := test.apiUtil.GetAndResetCreatedPVs()
	expectedLen := len(expectedPVs)
	actualLen := len(createdPVs)
	if expectedLen != actualLen {
		t.Errorf("Expected %v created PVs, got %v", expectedLen, actualLen)
	}

	for pvName, createdPV := range createdPVs {
		expectedPV, found := expectedPVs[pvName]
		if !found {
			t.Errorf("Did not expect created PVs %v", pvName)
		}
		if createdPV.Spec.PersistentVolumeSource.Local.Path != expectedPV.path {
			t.Errorf("Expected path %q, got %q", expectedPV.path, createdPV.Spec.PersistentVolumeSource.Local.Path)
		}
		_, exists := test.cache.GetPV(pvName)
		if !exists {
			t.Errorf("PV %q not in cache", pvName)
		}
		capacity, ok := createdPV.Spec.Capacity[v1.ResourceStorage]
		if !ok {
			t.Errorf("Unexpected empty resource storage")
		}
		capacityInt, ok := capacity.AsInt64()
		if !ok {
			t.Errorf("Unable to convert resource storage into int64")
		}
		if uint64(capacityInt) != expectedPV.capacity {
			t.Errorf("Expected capacity %d, got %d", expectedPV.capacity, capacityInt)
		}
		// TODO: verify storage class
		verifyProvisionerName(t, createdPV)
		verifyNodeAffinity(t, createdPV)
	}
}

func verifyPVsNotInCache(t *testing.T, test *testConfig) {
	for _, files := range test.dirLayout {
		for _, file := range files {
			pvName := fmt.Sprintf("local-pv-%x", file.Hash)
			_, exists := test.cache.GetPV(pvName)
			if exists {
				t.Errorf("Expected PV %q to not be in cache", pvName)
			}
		}
	}
}
<commit_msg>Refactor capacity verification test<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package discovery

import (
	"fmt"
	"path/filepath"
	"testing"

	"github.com/kubernetes-incubator/external-storage/local-volume/provisioner/pkg/cache"
	"github.com/kubernetes-incubator/external-storage/local-volume/provisioner/pkg/common"
	"github.com/kubernetes-incubator/external-storage/local-volume/provisioner/pkg/util"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/pkg/api/v1"
	v1helper "k8s.io/client-go/pkg/api/v1/helper"
)

const (
	testHostDir         = "/mnt/disks"
	testMountDir        = "/discoveryPath"
	testNodeName        = "test-node"
	testProvisionerName = "test-provisioner"
)

var testNode = &v1.Node{
	ObjectMeta: metav1.ObjectMeta{
		Name: testNodeName,
		Labels: map[string]string{
			common.NodeLabelKey: testNodeName,
		},
	},
}

var scMapping = map[string]string{
	"sc1": "dir1",
	"sc2": "dir2",
}

type testConfig struct {
	// The directory layout for the test
	// Key = directory, Value = list of volumes under that directory
	dirLayout map[string][]*util.FakeFile
	// The volumes that are expected to be created as PVs
	// Key = directory, Value = list of volumes under that directory
	expectedVolumes map[string][]*util.FakeFile
	// True if testing api failure
	apiShouldFail bool
	// The rest are set during setup
	volUtil *util.FakeVolumeUtil
	apiUtil *util.FakeAPIUtil
	cache   *cache.VolumeCache
}

func TestDiscoverVolumes_Basic(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount1", Hash: 0xaaaafef5, Capacity: 100 * 1024},
			{Name: "mount2", Hash: 0x79412c38, Capacity: 100 * 1024 * 1024},
		},
		"dir2": {
			{Name: "mount1", Hash: 0xa7aafa3c},
			{Name: "mount2", Hash: 0x7c4130f1},
		},
	}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: vols,
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()
	verifyCreatedPVs(t, test)
}

func TestDiscoverVolumes_BasicTwice(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount1", Hash: 0xaaaafef5},
			{Name: "mount2", Hash: 0x79412c38},
		},
		"dir2": {
			{Name: "mount1", Hash: 0xa7aafa3c},
			{Name: "mount2", Hash: 0x7c4130f1},
		},
	}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: vols,
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()
	verifyCreatedPVs(t, test)

	// Second time should not create any new volumes
	test.expectedVolumes = map[string][]*util.FakeFile{}
	d.DiscoverLocalVolumes()
	verifyCreatedPVs(t, test)
}

func TestDiscoverVolumes_NoDir(t *testing.T) {
	vols := map[string][]*util.FakeFile{}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: vols,
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()
	verifyCreatedPVs(t, test)
}

func TestDiscoverVolumes_EmptyDir(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {},
	}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: vols,
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()
	verifyCreatedPVs(t, test)
}

func TestDiscoverVolumes_NewVolumesLater(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount1", Hash: 0xaaaafef5},
			{Name: "mount2", Hash: 0x79412c38},
		},
		"dir2": {
			{Name: "mount1", Hash: 0xa7aafa3c},
			{Name: "mount2", Hash: 0x7c4130f1},
		},
	}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: vols,
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()

	verifyCreatedPVs(t, test)

	// Some new mount points show up
	newVols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount3", Hash: 0xf34b8003},
			{Name: "mount4", Hash: 0x144e29de},
		},
	}
	test.volUtil.AddNewFiles(testMountDir, newVols)
	test.expectedVolumes = newVols

	d.DiscoverLocalVolumes()

	verifyCreatedPVs(t, test)
}

func TestDiscoverVolumes_CreatePVFails(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount1", Hash: 0xaaaafef5},
			{Name: "mount2", Hash: 0x79412c38},
		},
		"dir2": {
			{Name: "mount1", Hash: 0xa7aafa3c},
			{Name: "mount2", Hash: 0x7c4130f1},
		},
	}
	test := &testConfig{
		apiShouldFail:   true,
		dirLayout:       vols,
		expectedVolumes: map[string][]*util.FakeFile{},
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()

	verifyCreatedPVs(t, test)
	verifyPVsNotInCache(t, test)
}

func TestDiscoverVolumes_BadVolume(t *testing.T) {
	vols := map[string][]*util.FakeFile{
		"dir1": {
			{Name: "mount1", IsNotDir: true},
		},
	}
	test := &testConfig{
		dirLayout:       vols,
		expectedVolumes: map[string][]*util.FakeFile{},
	}
	d := testSetup(t, test)

	d.DiscoverLocalVolumes()

	verifyCreatedPVs(t, test)
	verifyPVsNotInCache(t, test)
}

func testSetup(t *testing.T, test *testConfig) *Discoverer {
	test.cache = cache.NewVolumeCache()
	test.volUtil = util.NewFakeVolumeUtil(false)
	test.volUtil.AddNewFiles(testMountDir, test.dirLayout)
	test.apiUtil = util.NewFakeAPIUtil(test.apiShouldFail, test.cache)

	userConfig := &common.UserConfig{
		Node:         testNode,
		MountDir:     testMountDir,
		HostDir:      testHostDir,
		DiscoveryMap: scMapping,
	}
	runConfig := &common.RuntimeConfig{
		UserConfig: userConfig,
		Cache:      test.cache,
		VolUtil:    test.volUtil,
		APIUtil:    test.apiUtil,
		Name:       testProvisionerName,
	}
	d, err := NewDiscoverer(runConfig)
	if err != nil {
		t.Fatalf("Error setting up test discoverer: %v", err)
	}
	return d
}

func findSCName(t *testing.T, targetDir string, test *testConfig) string {
	for sc, dir := range scMapping {
		if dir == targetDir {
			return sc
		}
	}
	t.Fatalf("Failed to find SC Name for directory %v", targetDir)
	return ""
}

func verifyNodeAffinity(t *testing.T, pv *v1.PersistentVolume) {
	affinity, err := v1helper.GetStorageNodeAffinityFromAnnotation(pv.Annotations)
	if err != nil {
		t.Errorf("Could not get node affinity from annotation: %v", err)
		return
	}
	if affinity == nil {
		t.Errorf("No node affinity found")
		return
	}

	selector := affinity.RequiredDuringSchedulingIgnoredDuringExecution
	if selector == nil {
		t.Errorf("NodeAffinity node selector is nil")
		return
	}
	terms := selector.NodeSelectorTerms
	if len(terms) != 1 {
		t.Errorf("Node selector term count is %v, expected 1", len(terms))
		return
	}
	reqs := terms[0].MatchExpressions
	if len(reqs) != 1 {
		t.Errorf("Node selector term requirements count is %v, expected 1", len(reqs))
		return
	}

	req := reqs[0]
	if req.Key != common.NodeLabelKey {
		t.Errorf("Node selector requirement key is %v, expected %v", req.Key, common.NodeLabelKey)
	}
	if req.Operator != v1.NodeSelectorOpIn {
		t.Errorf("Node selector requirement operator is %v, expected %v", req.Operator, v1.NodeSelectorOpIn)
	}
	if len(req.Values) != 1 {
		t.Errorf("Node selector requirement value count is %v, expected 1", len(req.Values))
		return
	}
	if req.Values[0] != testNodeName {
		t.Errorf("Node selector requirement value is %v, expected %v", req.Values[0], testNodeName)
	}
}

func verifyProvisionerName(t *testing.T, pv *v1.PersistentVolume) {
	if len(pv.Annotations) == 0 {
		t.Errorf("Annotations not set")
		return
	}
	name, found := pv.Annotations[common.AnnProvisionedBy]
	if !found {
		t.Errorf("Provisioned by annotations not set")
		return
	}
	if name != testProvisionerName {
		t.Errorf("Provisioned name is %q, expected %q", name, testProvisionerName)
	}
}

func verifyCapacity(t *testing.T, createdPV *v1.PersistentVolume, expectedPV *testPVInfo) {
	capacity, ok := createdPV.Spec.Capacity[v1.ResourceStorage]
	if !ok {
		t.Errorf("Unexpected empty resource storage")
	}
	capacityInt, ok := capacity.AsInt64()
	if !ok {
		t.Errorf("Unable to convert resource storage into int64")
	}
	if uint64(capacityInt) != expectedPV.capacity {
		t.Errorf("Expected capacity %d, got %d", expectedPV.capacity, capacityInt)
	}
}

// testPVInfo contains all the fields we are intested in validating.
type testPVInfo struct {
	pvName   string
	path     string
	capacity uint64
}

func verifyCreatedPVs(t *testing.T, test *testConfig) {
	expectedPVs := map[string]*testPVInfo{}
	for dir, files := range test.expectedVolumes {
		for _, file := range files {
			pvName := fmt.Sprintf("local-pv-%x", file.Hash)
			path := filepath.Join(testHostDir, dir, file.Name)
			expectedPVs[pvName] = &testPVInfo{
				pvName:   pvName,
				path:     path,
				capacity: file.Capacity,
			}
		}
	}

	createdPVs := test.apiUtil.GetAndResetCreatedPVs()
	expectedLen := len(expectedPVs)
	actualLen := len(createdPVs)
	if expectedLen != actualLen {
		t.Errorf("Expected %v created PVs, got %v", expectedLen, actualLen)
	}

	for pvName, createdPV := range createdPVs {
		expectedPV, found := expectedPVs[pvName]
		if !found {
			t.Errorf("Did not expect created PVs %v", pvName)
		}
		if createdPV.Spec.PersistentVolumeSource.Local.Path != expectedPV.path {
			t.Errorf("Expected path %q, got %q", expectedPV.path, createdPV.Spec.PersistentVolumeSource.Local.Path)
		}
		_, exists := test.cache.GetPV(pvName)
		if !exists {
			t.Errorf("PV %q not in cache", pvName)
		}

		// TODO: verify storage class
		verifyProvisionerName(t, createdPV)
		verifyNodeAffinity(t, createdPV)
		verifyCapacity(t, createdPV, expectedPV)
	}
}

func verifyPVsNotInCache(t *testing.T, test *testConfig) {
	for _, files := range test.dirLayout {
		for _, file := range files {
			pvName := fmt.Sprintf("local-pv-%x", file.Hash)
			_, exists := test.cache.GetPV(pvName)
			if exists {
				t.Errorf("Expected PV %q to not be in cache", pvName)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"

	"github.com/spf13/cobra"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxc/config"
	"github.com/lxc/lxd/lxc/utils"
	"github.com/lxc/lxd/shared"
	cli "github.com/lxc/lxd/shared/cmd"
	"github.com/lxc/lxd/shared/i18n"
)

type cmdCopy struct {
	global *cmdGlobal

	flagNoProfiles    bool
	flagProfile       []string
	flagConfig        []string
	flagDevice        []string
	flagEphemeral     bool
	flagContainerOnly bool
	flagMode          string
	flagStateless     bool
	flagStorage       string
	flagTarget        string
}

func (c *cmdCopy) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = i18n.G("copy [<remote>:]<source>[/<snapshot>] [[<remote>:]<destination>]")
	cmd.Aliases = []string{"cp"}
	cmd.Short = i18n.G("Copy containers within or in between LXD instances")
	cmd.Long = cli.FormatSection(i18n.G("Description"), i18n.G(
		`Copy containers within or in between LXD instances`))

	cmd.RunE = c.Run
	cmd.Flags().StringArrayVarP(&c.flagConfig, "config", "c", nil, i18n.G("Config key/value to apply to the new container")+"``")
	cmd.Flags().StringArrayVarP(&c.flagDevice, "device", "d", nil, i18n.G("New key/value to apply to a specific device")+"``")
	cmd.Flags().StringArrayVarP(&c.flagProfile, "profile", "p", nil, i18n.G("Profile to apply to the new container")+"``")
	cmd.Flags().BoolVarP(&c.flagEphemeral, "ephemeral", "e", false, i18n.G("Ephemeral container"))
	cmd.Flags().StringVar(&c.flagMode, "mode", "pull", i18n.G("Transfer mode. One of pull (default), push or relay")+"``")
	cmd.Flags().BoolVar(&c.flagContainerOnly, "container-only", false, i18n.G("Copy the container without its snapshots"))
	cmd.Flags().BoolVar(&c.flagStateless, "stateless", false, i18n.G("Copy a stateful container stateless"))
	cmd.Flags().StringVarP(&c.flagStorage, "storage", "s", "", i18n.G("Storage pool name")+"``")
	cmd.Flags().StringVar(&c.flagTarget, "target", "", i18n.G("Cluster member name")+"``")
	cmd.Flags().BoolVar(&c.flagNoProfiles, "no-profiles", false, i18n.G("Create the container with no profiles applied"))

	return cmd
}

func (c *cmdCopy) copyContainer(conf *config.Config, sourceResource string,
	destResource string, keepVolatile bool, ephemeral int, stateful bool,
	containerOnly bool, mode string, pool string) error {
	// Parse the source
	sourceRemote, sourceName, err := conf.ParseRemote(sourceResource)
	if err != nil {
		return err
	}

	// Parse the destination
	destRemote, destName, err := conf.ParseRemote(destResource)
	if err != nil {
		return err
	}

	// Make sure we have a container or snapshot name
	if sourceName == "" {
		return fmt.Errorf(i18n.G("You must specify a source container name"))
	}

	// Check that a destination container was specified, if --target is passed.
	if destName == "" && c.flagTarget != "" {
		return fmt.Errorf(i18n.G("You must specify a destination container name when using --target"))
	}

	// If no destination name was provided, use the same as the source
	if destName == "" && destResource != "" {
		destName = sourceName
	}

	// Connect to the source host
	source, err := conf.GetContainerServer(sourceRemote)
	if err != nil {
		return err
	}

	// Connect to the destination host
	var dest lxd.ContainerServer
	if sourceRemote == destRemote {
		// Source and destination are the same
		dest = source
	} else {
		// Destination is different, connect to it
		dest, err = conf.GetContainerServer(destRemote)
		if err != nil {
			return err
		}
	}
	// Confirm that --target is only used with a cluster
	if c.flagTarget != "" && !dest.IsClustered() {
		return fmt.Errorf(i18n.G("To use --target, the destination remote must be a cluster"))
	}

	// Parse the config overrides
	configMap := map[string]string{}
	for _, entry := range c.flagConfig {
		if !strings.Contains(entry, "=") {
			return fmt.Errorf(i18n.G("Bad key=value pair: %s"), entry)
		}

		fields := strings.SplitN(entry, "=", 2)
		configMap[fields[0]] = fields[1]
	}

	// Parse the device overrides
	deviceMap := map[string]map[string]string{}
	for _, entry := range c.flagDevice {
		if !strings.Contains(entry, "=") || !strings.Contains(entry, ",") {
			return fmt.Errorf(i18n.G("Bad syntax, expecting <device>,<key>=<value>: %s"), entry)
		}

		deviceFields := strings.SplitN(entry, ",", 2)
		keyFields := strings.SplitN(deviceFields[1], "=", 2)

		if deviceMap[deviceFields[0]] == nil {
			deviceMap[deviceFields[0]] = map[string]string{}
		}

		deviceMap[deviceFields[0]][keyFields[0]] = keyFields[1]
	}

	var op lxd.RemoteOperation
	if shared.IsSnapshot(sourceName) {
		// Prepare the container creation request
		args := lxd.ContainerSnapshotCopyArgs{
			Name: destName,
			Mode: mode,
			Live: stateful,
		}

		// Copy of a snapshot into a new container
		srcFields := strings.SplitN(sourceName, shared.SnapshotDelimiter, 2)
		entry, _, err := source.GetContainerSnapshot(srcFields[0], srcFields[1])
		if err != nil {
			return err
		}

		// Allow adding additional profiles
		if c.flagProfile != nil {
			entry.Profiles = append(entry.Profiles, c.flagProfile...)
		} else if c.flagNoProfiles {
			entry.Profiles = []string{}
		}

		// Allow setting additional config keys
		if configMap != nil {
			for key, value := range configMap {
				entry.Config[key] = value
			}
		}

		// Allow setting device overrides
		if deviceMap != nil {
			for k, m := range deviceMap {
				if entry.Devices[k] == nil {
					entry.Devices[k] = m
					continue
				}

				for key, value := range m {
					entry.Devices[k][key] = value
				}
			}
		}

		// Allow overriding the ephemeral status
		if ephemeral == 1 {
			entry.Ephemeral = true
		} else if ephemeral == 0 {
			entry.Ephemeral = false
		}

		rootDiskDeviceKey, _, _ := shared.GetRootDiskDevice(entry.Devices)
		if err != nil {
			return err
		}

		if rootDiskDeviceKey != "" && pool != "" {
			entry.Devices[rootDiskDeviceKey]["pool"] = pool
		} else if pool != "" {
			entry.Devices["root"] = map[string]string{
				"type": "disk",
				"path": "/",
				"pool": pool,
			}
		}

		// Strip the volatile keys if requested
		if !keepVolatile {
			for k := range entry.Config {
				if k == "volatile.base_image" {
					continue
				}

				if strings.HasPrefix(k, "volatile") {
					delete(entry.Config, k)
				}
			}
		}

		// Do the actual copy
		if c.flagTarget != "" {
			dest = dest.UseTarget(c.flagTarget)
		}

		op, err = dest.CopyContainerSnapshot(source, srcFields[0], *entry, &args)
		if err != nil {
			return err
		}
	} else {
		// Prepare the container creation request
		args := lxd.ContainerCopyArgs{
			Name:          destName,
			Live:          stateful,
			ContainerOnly: containerOnly,
			Mode:          mode,
		}

		// Copy of a container into a new container
		entry, _, err := source.GetContainer(sourceName)
		if err != nil {
			return err
		}

		// Allow adding additional profiles
		if c.flagProfile != nil {
			entry.Profiles = append(entry.Profiles, c.flagProfile...)
		} else if c.flagNoProfiles {
			entry.Profiles = []string{}
		}

		// Allow setting additional config keys
		if configMap != nil {
			for key, value := range configMap {
				entry.Config[key] = value
			}
		}

		// Allow setting device overrides
		if deviceMap != nil {
			for k, m := range deviceMap {
				if entry.Devices[k] == nil {
					entry.Devices[k] = m
					continue
				}

				for key, value := range m {
					entry.Devices[k][key] = value
				}
			}
		}

		// Allow overriding the ephemeral status
		if ephemeral == 1 {
			entry.Ephemeral = true
		} else if ephemeral == 0 {
			entry.Ephemeral = false
		}

		rootDiskDeviceKey, _, _ := shared.GetRootDiskDevice(entry.Devices)
		if err != nil {
			return err
		}

		if rootDiskDeviceKey != "" && pool != "" {
			entry.Devices[rootDiskDeviceKey]["pool"] = pool
		} else if pool != "" {
			entry.Devices["root"] = map[string]string{
				"type": "disk",
				"path": "/",
				"pool": pool,
			}
		}

		// Strip the volatile keys if requested
		if !keepVolatile {
			for k := range entry.Config {
				if k == "volatile.base_image" {
					continue
				}

				if strings.HasPrefix(k, "volatile") {
					delete(entry.Config, k)
				}
			}
		}

		// Do the actual copy
		if c.flagTarget != "" {
			dest = dest.UseTarget(c.flagTarget)
		}

		op, err = dest.CopyContainer(source, *entry, &args)
		if err != nil {
			return err
		}
	}

	// Watch the background operation
	progress := utils.ProgressRenderer{
		Format: i18n.G("Transferring container: %s"),
		Quiet:  c.global.flagQuiet,
	}

	_, err = op.AddHandler(progress.UpdateOp)
	if err != nil {
		progress.Done("")
		return err
	}

	// Wait for the copy to complete
	err = utils.CancelableWait(op, &progress)
	if err != nil {
		progress.Done("")
		return err
	}
	progress.Done("")

	// If choosing a random name, show it to the user
	if destResource == "" {
		// Get the successful operation data
		opInfo, err := op.GetTarget()
		if err != nil {
			return err
		}

		// Extract the list of affected containers
		containers, ok := opInfo.Resources["containers"]
		if !ok || len(containers) != 1 {
			return fmt.Errorf(i18n.G("Failed to get the new container name"))
		}

		// Extract the name of the container
		fields := strings.Split(containers[0], "/")
		fmt.Printf(i18n.G("Container name is: %s")+"\n", fields[len(fields)-1])
	}

	return nil
}

func (c *cmdCopy) Run(cmd *cobra.Command, args []string) error {
	conf := c.global.conf

	// Sanity checks
	exit, err := c.global.CheckArgs(cmd, args, 1, 2)
	if exit {
		return err
	}

	// For copies, default to non-ephemeral and allow override (move uses -1)
	ephem := 0
	if c.flagEphemeral {
		ephem = 1
	}

	// Parse the mode
	mode := "pull"
	if c.flagMode != "" {
		mode = c.flagMode
	}

	stateful := !c.flagStateless

	// If not target name is specified, one will be chosed by the server
	if len(args) < 2 {
		return c.copyContainer(conf, args[0], "", false, ephem,
			stateful, c.flagContainerOnly, mode, c.flagStorage)
	}

	// Normal copy with a pre-determined name
	return c.copyContainer(conf, args[0], args[1], false, ephem,
		stateful, c.flagContainerOnly, mode, c.flagStorage)
}
<commit_msg>lxc/copy: --container-only is meaningless for snapshots<commit_after>package main

import (
	"fmt"
	"strings"

	"github.com/spf13/cobra"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxc/config"
	"github.com/lxc/lxd/lxc/utils"
	"github.com/lxc/lxd/shared"
	cli "github.com/lxc/lxd/shared/cmd"
	"github.com/lxc/lxd/shared/i18n"
)

type cmdCopy struct {
	global *cmdGlobal

	flagNoProfiles    bool
	flagProfile       []string
	flagConfig        []string
	flagDevice        []string
	flagEphemeral     bool
	flagContainerOnly bool
	flagMode          string
	flagStateless     bool
	flagStorage       string
	flagTarget        string
}

func (c *cmdCopy) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = i18n.G("copy [<remote>:]<source>[/<snapshot>] [[<remote>:]<destination>]")
	cmd.Aliases = []string{"cp"}
	cmd.Short = i18n.G("Copy containers within or in between LXD instances")
	cmd.Long = cli.FormatSection(i18n.G("Description"), i18n.G(
		`Copy containers within or in between LXD instances`))

	cmd.RunE = c.Run
	cmd.Flags().StringArrayVarP(&c.flagConfig, "config", "c", nil, i18n.G("Config key/value to apply to the new container")+"``")
	cmd.Flags().StringArrayVarP(&c.flagDevice, "device", "d", nil, i18n.G("New key/value to apply to a specific device")+"``")
	cmd.Flags().StringArrayVarP(&c.flagProfile, "profile", "p", nil, i18n.G("Profile to apply to the new container")+"``")
	cmd.Flags().BoolVarP(&c.flagEphemeral, "ephemeral", "e", false, i18n.G("Ephemeral container"))
	cmd.Flags().StringVar(&c.flagMode, "mode", "pull", i18n.G("Transfer mode. One of pull (default), push or relay")+"``")
	cmd.Flags().BoolVar(&c.flagContainerOnly, "container-only", false, i18n.G("Copy the container without its snapshots"))
	cmd.Flags().BoolVar(&c.flagStateless, "stateless", false, i18n.G("Copy a stateful container stateless"))
	cmd.Flags().StringVarP(&c.flagStorage, "storage", "s", "", i18n.G("Storage pool name")+"``")
	cmd.Flags().StringVar(&c.flagTarget, "target", "", i18n.G("Cluster member name")+"``")
	cmd.Flags().BoolVar(&c.flagNoProfiles, "no-profiles", false, i18n.G("Create the container with no profiles applied"))

	return cmd
}

func (c *cmdCopy) copyContainer(conf *config.Config, sourceResource string,
	destResource string, keepVolatile bool, ephemeral int, stateful bool,
	containerOnly bool, mode string, pool string) error {
	// Parse the source
	sourceRemote, sourceName, err := conf.ParseRemote(sourceResource)
	if err != nil {
		return err
	}

	// Parse the destination
	destRemote, destName, err := conf.ParseRemote(destResource)
	if err != nil {
		return err
	}

	// Make sure we have a container or snapshot name
	if sourceName == "" {
		return fmt.Errorf(i18n.G("You must specify a source container name"))
	}

	// Check that a destination container was specified, if --target is passed.
	if destName == "" && c.flagTarget != "" {
		return fmt.Errorf(i18n.G("You must specify a destination container name when using --target"))
	}

	// If no destination name was provided, use the same as the source
	if destName == "" && destResource != "" {
		destName = sourceName
	}

	// Connect to the source host
	source, err := conf.GetContainerServer(sourceRemote)
	if err != nil {
		return err
	}

	// Connect to the destination host
	var dest lxd.ContainerServer
	if sourceRemote == destRemote {
		// Source and destination are the same
		dest = source
	} else {
		// Destination is different, connect to it
		dest, err = conf.GetContainerServer(destRemote)
		if err != nil {
			return err
		}
	}
	// Confirm that --target is only used with a cluster
	if c.flagTarget != "" && !dest.IsClustered() {
		return fmt.Errorf(i18n.G("To use --target, the destination remote must be a cluster"))
	}

	// Parse the config overrides
	configMap := map[string]string{}
	for _, entry := range c.flagConfig {
		if !strings.Contains(entry, "=") {
			return fmt.Errorf(i18n.G("Bad key=value pair: %s"), entry)
		}

		fields := strings.SplitN(entry, "=", 2)
		configMap[fields[0]] = fields[1]
	}

	// Parse the device overrides
	deviceMap := map[string]map[string]string{}
	for _, entry := range c.flagDevice {
		if !strings.Contains(entry, "=") || !strings.Contains(entry, ",") {
			return fmt.Errorf(i18n.G("Bad syntax, expecting <device>,<key>=<value>: %s"), entry)
		}

		deviceFields := strings.SplitN(entry, ",", 2)
		keyFields := strings.SplitN(deviceFields[1], "=", 2)

		if deviceMap[deviceFields[0]] == nil {
			deviceMap[deviceFields[0]] = map[string]string{}
		}

		deviceMap[deviceFields[0]][keyFields[0]] = keyFields[1]
	}

	var op lxd.RemoteOperation
	if shared.IsSnapshot(sourceName) {
		if containerOnly {
			return fmt.Errorf(i18n.G("--container-only can't be passed when the source is a snapshot"))
		}

		// Prepare the container creation request
		args := lxd.ContainerSnapshotCopyArgs{
			Name: destName,
			Mode: mode,
			Live: stateful,
		}

		// Copy of a snapshot into a new container
		srcFields := strings.SplitN(sourceName, shared.SnapshotDelimiter, 2)
		entry, _, err := source.GetContainerSnapshot(srcFields[0], srcFields[1])
		if err != nil {
			return err
		}

		// Allow adding additional profiles
		if c.flagProfile != nil {
			entry.Profiles = append(entry.Profiles, c.flagProfile...)
		} else if c.flagNoProfiles {
			entry.Profiles = []string{}
		}

		// Allow setting additional config keys
		if configMap != nil {
			for key, value := range configMap {
				entry.Config[key] = value
			}
		}

		// Allow setting device overrides
		if deviceMap != nil {
			for k, m := range deviceMap {
				if entry.Devices[k] == nil {
					entry.Devices[k] = m
					continue
				}

				for key, value := range m {
					entry.Devices[k][key] = value
				}
			}
		}

		// Allow overriding the ephemeral status
		if ephemeral == 1 {
			entry.Ephemeral = true
		} else if ephemeral == 0 {
			entry.Ephemeral = false
		}

		rootDiskDeviceKey, _, _ := shared.GetRootDiskDevice(entry.Devices)
		if err != nil {
			return err
		}

		if rootDiskDeviceKey != "" && pool != "" {
			entry.Devices[rootDiskDeviceKey]["pool"] = pool
		} else if pool != "" {
			entry.Devices["root"] = map[string]string{
				"type": "disk",
				"path": "/",
				"pool": pool,
			}
		}

		// Strip the volatile keys if requested
		if !keepVolatile {
			for k := range entry.Config {
				if k == "volatile.base_image" {
					continue
				}

				if strings.HasPrefix(k, "volatile") {
					delete(entry.Config, k)
				}
			}
		}

		// Do the actual copy
		if c.flagTarget != "" {
			dest = dest.UseTarget(c.flagTarget)
		}

		op, err = dest.CopyContainerSnapshot(source, srcFields[0], *entry, &args)
		if err != nil {
			return err
		}
	} else {
		// Prepare the container creation request
		args := lxd.ContainerCopyArgs{
			Name:          destName,
			Live:          stateful,
			ContainerOnly: containerOnly,
			Mode:          mode,
		}

		// Copy of a container into a new container
		entry, _, err := source.GetContainer(sourceName)
		if err != nil {
			return err
		}

		// Allow adding additional profiles
		if c.flagProfile != nil {
			entry.Profiles = append(entry.Profiles, c.flagProfile...)
		} else if c.flagNoProfiles {
			entry.Profiles = []string{}
		}

		// Allow setting additional config keys
		if configMap != nil {
			for key, value := range configMap {
				entry.Config[key] = value
			}
		}

		// Allow setting device overrides
		if deviceMap != nil {
			for k, m := range deviceMap {
				if entry.Devices[k] == nil {
					entry.Devices[k] = m
					continue
				}

				for key, value := range m {
					entry.Devices[k][key] = value
				}
			}
		}

		// Allow overriding the ephemeral status
		if ephemeral == 1 {
			entry.Ephemeral = true
		} else if ephemeral == 0 {
			entry.Ephemeral = false
		}

		rootDiskDeviceKey, _, _ := shared.GetRootDiskDevice(entry.Devices)
		if err != nil {
			return err
		}

		if rootDiskDeviceKey != "" && pool != "" {
			entry.Devices[rootDiskDeviceKey]["pool"] = pool
		} else if pool != "" {
			entry.Devices["root"] = map[string]string{
				"type": "disk",
				"path": "/",
				"pool": pool,
			}
		}

		// Strip the volatile keys if requested
		if !keepVolatile {
			for k := range entry.Config {
				if k == "volatile.base_image" {
					continue
				}

				if strings.HasPrefix(k, "volatile") {
					delete(entry.Config, k)
				}
			}
		}

		// Do the actual copy
		if c.flagTarget != "" {
			dest = dest.UseTarget(c.flagTarget)
		}

		op, err = dest.CopyContainer(source, *entry, &args)
		if err != nil {
			return err
		}
	}

	// Watch the background operation
	progress := utils.ProgressRenderer{
		Format: i18n.G("Transferring container: %s"),
		Quiet:  c.global.flagQuiet,
	}

	_, err = op.AddHandler(progress.UpdateOp)
	if err != nil {
		progress.Done("")
		return err
	}

	// Wait for the copy to complete
	err = utils.CancelableWait(op, &progress)
	if err != nil {
		progress.Done("")
		return err
	}
	progress.Done("")

	// If choosing a random name, show it to the user
	if destResource == "" {
		// Get the successful operation data
		opInfo, err := op.GetTarget()
		if err != nil {
			return err
		}

		// Extract the list of affected containers
		containers, ok := opInfo.Resources["containers"]
		if !ok || len(containers) != 1 {
			return fmt.Errorf(i18n.G("Failed to get the new container name"))
		}

		// Extract the name of the container
		fields := strings.Split(containers[0], "/")
		fmt.Printf(i18n.G("Container name is: %s")+"\n", fields[len(fields)-1])
	}

	return nil
}

func (c *cmdCopy) Run(cmd *cobra.Command, args []string) error {
	conf := c.global.conf

	// Sanity checks
	exit, err := c.global.CheckArgs(cmd, args, 1, 2)
	if exit {
		return err
	}

	// For copies, default to non-ephemeral and allow override (move uses -1)
	ephem := 0
	if c.flagEphemeral {
		ephem = 1
	}

	// Parse the mode
	mode := "pull"
	if c.flagMode != "" {
		mode = c.flagMode
	}

	stateful := !c.flagStateless

	// If not target name is specified, one will be chosed by the server
	if len(args) < 2 {
		return c.copyContainer(conf, args[0], "", false, ephem,
			stateful, c.flagContainerOnly, mode, c.flagStorage)
	}

	// Normal copy with a pre-determined name
	return c.copyContainer(conf, args[0], args[1], false, ephem,
		stateful, c.flagContainerOnly, mode, c.flagStorage)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/i18n"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/gnuflag"
)

type copyCmd struct {
	ephem bool
}

func (c *copyCmd) showByDefault() bool {
	return true
}

func (c *copyCmd) usage() string {
	return i18n.G(
		`Copy containers within or in between lxd instances.

lxc copy [remote:]<source container> [remote:]<destination container> [--ephemeral|e]`)
}

func (c *copyCmd) flags() {
	gnuflag.BoolVar(&c.ephem, "ephemeral", false, i18n.G("Ephemeral container"))
	gnuflag.BoolVar(&c.ephem, "e", false, i18n.G("Ephemeral container"))
}

func copyContainer(config *lxd.Config, sourceResource string, destResource string, keepVolatile bool, ephemeral int) error {
	sourceRemote, sourceName := config.ParseRemoteAndContainer(sourceResource)
	destRemote, destName := config.ParseRemoteAndContainer(destResource)

	if sourceName == "" {
		return fmt.Errorf(i18n.G("you must specify a source container name"))
	}

	if destName == "" {
		destName = sourceName
	}

	source, err := lxd.NewClient(config, sourceRemote)
	if err != nil {
		return err
	}

	status := &shared.ContainerState{}

	// TODO: presumably we want to do this for copying snapshots too? We
	// need to think a bit more about how we track the baseImage in the
	// face of LVM and snapshots in general; this will probably make more
	// sense once that work is done.
	baseImage := ""

	if !shared.IsSnapshot(sourceName) {
		status, err = source.ContainerStatus(sourceName)
		if err != nil {
			return err
		}

		baseImage = status.Config["volatile.base_image"]

		if !keepVolatile {
			for k := range status.Config {
				if strings.HasPrefix(k, "volatile") {
					delete(status.Config, k)
				}
			}
		}
	}

	// Do a local copy if the remotes are the same, otherwise do a migration
	if sourceRemote == destRemote {
		if sourceName == destName {
			return fmt.Errorf(i18n.G("can't copy to the same container name"))
		}

		cp, err := source.LocalCopy(sourceName, destName, status.Config, status.Profiles, ephemeral == 1)
		if err != nil {
			return err
		}

		return source.WaitForSuccess(cp.Operation)
	} else {
		dest, err := lxd.NewClient(config, destRemote)
		if err != nil {
			return err
		}

		sourceProfs := shared.NewStringSet(status.Profiles)
		destProfs, err := dest.ListProfiles()
		if err != nil {
			return err
		}

		if !sourceProfs.IsSubset(shared.NewStringSet(destProfs)) {
			return fmt.Errorf(i18n.G("not all the profiles from the source exist on the target"))
		}

		if ephemeral == -1 {
			ct, err := source.ContainerStatus(sourceName)
			if err != nil {
				return err
			}

			if ct.Ephemeral {
				ephemeral = 1
			} else {
				ephemeral = 0
			}
		}

		sourceWSResponse, err := source.GetMigrationSourceWS(sourceName)
		if err != nil {
			return err
		}

		secrets := map[string]string{}

		op, err := sourceWSResponse.MetadataAsOperation()
		if err == nil && op.Metadata != nil {
			for k, v := range *op.Metadata {
				secrets[k] = v.(string)
			}
		} else {
			// FIXME: This is a backward compatibility codepath
			if err := json.Unmarshal(sourceWSResponse.Metadata, &secrets); err != nil {
				return err
			}
		}

		addresses, err := source.Addresses()
		if err != nil {
			return err
		}

		for _, addr := range addresses {
			sourceWSUrl := "https://" + addr + sourceWSResponse.Operation

			var migration *lxd.Response
			migration, err = dest.MigrateFrom(destName, sourceWSUrl, secrets, status.Architecture, status.Config, status.Devices, status.Profiles, baseImage, ephemeral == 1)
			if err != nil {
				continue
			}

			if err = dest.WaitForSuccess(migration.Operation); err != nil {
				continue
			}

			return nil
		}

		return err
	}
}

func (c *copyCmd) run(config *lxd.Config, args []string) error {
	if len(args) != 2 {
		return errArgs
	}

	ephem := 0
	if c.ephem {
		ephem = 1
	}

	return copyContainer(config, args[0], args[1], false, ephem)
}
<commit_msg>Add backward compatibility<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/i18n"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/gnuflag"
)

type copyCmd struct {
	ephem bool
}

func (c *copyCmd) showByDefault() bool {
	return true
}

func (c *copyCmd) usage() string {
	return i18n.G(
		`Copy containers within or in between lxd instances.

lxc copy [remote:]<source container> [remote:]<destination container> [--ephemeral|e]`)
}

func (c *copyCmd) flags() {
	gnuflag.BoolVar(&c.ephem, "ephemeral", false, i18n.G("Ephemeral container"))
	gnuflag.BoolVar(&c.ephem, "e", false, i18n.G("Ephemeral container"))
}

func copyContainer(config *lxd.Config, sourceResource string, destResource string, keepVolatile bool, ephemeral int) error {
	sourceRemote, sourceName := config.ParseRemoteAndContainer(sourceResource)
	destRemote, destName := config.ParseRemoteAndContainer(destResource)

	if sourceName == "" {
		return fmt.Errorf(i18n.G("you must specify a source container name"))
	}

	if destName == "" {
		destName = sourceName
	}

	source, err := lxd.NewClient(config, sourceRemote)
	if err != nil {
		return err
	}

	status := &shared.ContainerState{}

	// TODO: presumably we want to do this for copying snapshots too? We
	// need to think a bit more about how we track the baseImage in the
	// face of LVM and snapshots in general; this will probably make more
	// sense once that work is done.
	baseImage := ""

	if !shared.IsSnapshot(sourceName) {
		status, err = source.ContainerStatus(sourceName)
		if err != nil {
			return err
		}

		baseImage = status.Config["volatile.base_image"]

		if !keepVolatile {
			for k := range status.Config {
				if strings.HasPrefix(k, "volatile") {
					delete(status.Config, k)
				}
			}
		}
	}

	// Do a local copy if the remotes are the same, otherwise do a migration
	if sourceRemote == destRemote {
		if sourceName == destName {
			return fmt.Errorf(i18n.G("can't copy to the same container name"))
		}

		cp, err := source.LocalCopy(sourceName, destName, status.Config, status.Profiles, ephemeral == 1)
		if err != nil {
			return err
		}

		return source.WaitForSuccess(cp.Operation)
	} else {
		dest, err := lxd.NewClient(config, destRemote)
		if err != nil {
			return err
		}

		sourceProfs := shared.NewStringSet(status.Profiles)
		destProfs, err := dest.ListProfiles()
		if err != nil {
			return err
		}

		if !sourceProfs.IsSubset(shared.NewStringSet(destProfs)) {
			return fmt.Errorf(i18n.G("not all the profiles from the source exist on the target"))
		}

		if ephemeral == -1 {
			ct, err := source.ContainerStatus(sourceName)
			if err != nil {
				return err
			}

			if ct.Ephemeral {
				ephemeral = 1
			} else {
				ephemeral = 0
			}
		}

		sourceWSResponse, err := source.GetMigrationSourceWS(sourceName)
		if err != nil {
			return err
		}

		secrets := map[string]string{}

		op, err := sourceWSResponse.MetadataAsOperation()
		if err == nil && op.Metadata != nil {
			for k, v := range *op.Metadata {
				secrets[k] = v.(string)
			}
		} else {
			// FIXME: This is a backward compatibility codepath
			if err := json.Unmarshal(sourceWSResponse.Metadata, &secrets); err != nil {
				return err
			}
		}

		addresses, err := source.Addresses()
		if err != nil {
			return err
		}

		for _, addr := range addresses {
			var migration *lxd.Response

			sourceWSUrl := "https://" + addr + sourceWSResponse.Operation
			migration, err = dest.MigrateFrom(destName, sourceWSUrl, secrets, status.Architecture, status.Config, status.Devices, status.Profiles, baseImage, ephemeral == 1)
			if err != nil {
				continue
			}

			if err = dest.WaitForSuccess(migration.Operation); err != nil {
				// FIXME: This is a backward compatibility codepath
				sourceWSUrl := "wss://" + addr + sourceWSResponse.Operation + "/websocket"

				migration, err = dest.MigrateFrom(destName, sourceWSUrl, secrets, status.Architecture, status.Config, status.Devices, status.Profiles, baseImage, ephemeral == 1)
				if err != nil {
					continue
				}

				if err = dest.WaitForSuccess(migration.Operation); err != nil {
					continue
				}
			}

			return nil
		}

		return err
	}
}

func (c *copyCmd) run(config *lxd.Config, args []string) error {
	if len(args) != 2 {
		return errArgs
	}

	ephem := 0
	if c.ephem {
		ephem = 1
	}

	return copyContainer(config, args[0], args[1], false, ephem)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/gnuflag"
	"github.com/lxc/lxd/shared/i18n"
	"github.com/lxc/lxd/shared/termios"
)

type fileCmd struct {
	uid  int
	gid  int
	mode string

	recursive bool

	mkdirs bool
}

func (c *fileCmd) showByDefault() bool {
	return true
}

func (c *fileCmd) usage() string {
	return i18n.G(
		`Manage files on a container.

lxc file pull [-r|--recursive] <source> [<source>...] <target>
lxc file push [-r|--recursive] [-p|create-dirs] [--uid=UID] [--gid=GID] [--mode=MODE] <source> [<source>...] <target>
lxc file edit <file>

<source> in the case of pull, <target> in the case of push and <file> in the case of edit are <container name>/<path>

Examples:

To push /etc/hosts into the container foo:
  lxc file push /etc/hosts foo/etc/hosts

To pull /etc/hosts from the container:
  lxc file pull foo/etc/hosts .
`)
}

func (c *fileCmd) flags() {
	gnuflag.IntVar(&c.uid, "uid", -1, i18n.G("Set the file's uid on push"))
	gnuflag.IntVar(&c.gid, "gid", -1, i18n.G("Set the file's gid on push"))
	gnuflag.StringVar(&c.mode, "mode", "", i18n.G("Set the file's perms on push"))
	gnuflag.BoolVar(&c.recursive, "recusrive", false, i18n.G("Recursively push or pull files"))
	gnuflag.BoolVar(&c.recursive, "r", false, i18n.G("Recursively push or pull files"))
	gnuflag.BoolVar(&c.mkdirs, "create-dirs", false, i18n.G("Create any directories necessary"))
	gnuflag.BoolVar(&c.mkdirs, "p", false, i18n.G("Create any directories necessary"))
}

func (c *fileCmd) push(config *lxd.Config, send_file_perms bool, args []string) error {
	if len(args) < 2 {
		return errArgs
	}

	target := args[len(args)-1]
	pathSpec := strings.SplitN(target, "/", 2)

	if len(pathSpec) != 2 {
		return fmt.Errorf(i18n.G("Invalid target %s"), target)
	}

	targetPath := pathSpec[1]
	remote, container := config.ParseRemoteAndContainer(pathSpec[0])

	d, err := lxd.NewClient(config, remote)
	if err != nil {
		return err
	}

	var sourcefilenames []string
	for _, fname := range args[:len(args)-1] {
		if !strings.HasPrefix(fname, "--") {
			sourcefilenames = append(sourcefilenames, fname)
		}
	}

	mode := os.FileMode(0755)
	if c.mode != "" {
		if len(c.mode) == 3 {
			c.mode = "0" + c.mode
		}

		m, err := strconv.ParseInt(c.mode, 0, 0)
		if err != nil {
			return err
		}
		mode = os.FileMode(m)
	}

	if c.recursive {
		if c.uid != -1 || c.gid != -1 || c.mode != "" {
			return fmt.Errorf(i18n.G("can't supply uid/gid/mode in recursive mode"))
		}

		for _, fname := range sourcefilenames {
			if c.mkdirs {
				if err := d.MkdirP(container, fname, mode); err != nil {
					return err
				}
			}

			if err := d.RecursivePushFile(container, fname, pathSpec[1]); err != nil {
				return err
			}
		}

		return nil
	}

	uid := 0
	if c.uid >= 0 {
		uid = c.uid
	}

	gid := 0
	if c.gid >= 0 {
		gid = c.gid
	}

	_, targetfilename := filepath.Split(targetPath)

	if (targetfilename != "") && (len(sourcefilenames) > 1) {
		return errArgs
	}

	/* Make sure all of the files are accessible by us before trying to
	 * push any of them. */
	var files []*os.File
	for _, f := range sourcefilenames {
		var file *os.File
		if f == "-" {
			file = os.Stdin
		} else {
			file, err = os.Open(f)
			if err != nil {
				return err
			}
		}

		defer file.Close()
		files = append(files, file)
	}

	for _, f := range files {
		fpath := targetPath
		if targetfilename == "" {
			fpath = path.Join(fpath, path.Base(f.Name()))
		}

		if c.mkdirs {
			if err := d.MkdirP(container, filepath.Dir(fpath), mode); err != nil {
				return err
			}
		}

		if send_file_perms {
			if c.mode == "" || c.uid == -1 || c.gid == -1 {
				fMode, fUid, fGid, err := c.getOwner(f)
				if err != nil {
					return err
				}

				if c.mode == "" {
					mode = fMode
				}

				if c.uid == -1 {
					uid = fUid
				}

				if c.gid == -1 {
					gid = fGid
				}
			}

			err = d.PushFile(container, fpath, gid, uid, fmt.Sprintf("%04o", mode.Perm()), f)
		} else {
			err = d.PushFile(container, fpath, -1, -1, "", f)
		}

		if err != nil {
			return err
		}
	}

	return nil
}

func (c *fileCmd) pull(config *lxd.Config, args []string) error {
	if len(args) < 2 {
		return errArgs
	}

	target := args[len(args)-1]
	targetIsDir := false
	sb, err := os.Stat(target)
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	/*
	 * If the path exists, just use it. If it doesn't exist, it might be a
	 * directory in one of two cases:
	 *   1. Someone explicitly put "/" at the end
	 *   2. Someone provided more than one source. In this case the target
	 *      should be a directory so we can save all the files into it.
	 */
	if err == nil {
		targetIsDir = sb.IsDir()
		if !targetIsDir && len(args)-1 > 1 {
			return fmt.Errorf(i18n.G("More than one file to download, but target is not a directory"))
		}
	} else if strings.HasSuffix(target, string(os.PathSeparator)) || len(args)-1 > 1 {
		if err := os.MkdirAll(target, 0755); err != nil {
			return err
		}
		targetIsDir = true
	}

	for _, f := range args[:len(args)-1] {
		pathSpec := strings.SplitN(f, "/", 2)
		if len(pathSpec) != 2 {
			return fmt.Errorf(i18n.G("Invalid source %s"), f)
		}

		remote, container := config.ParseRemoteAndContainer(pathSpec[0])
		d, err := lxd.NewClient(config, remote)
		if err != nil {
			return err
		}

		if c.recursive {
			if err := d.RecursivePullFile(container, pathSpec[1], target); err != nil {
				return err
			}

			continue
		}

		_, _, mode, type_, buf, _, err := d.PullFile(container, pathSpec[1])
		if err != nil {
			return err
		}

		if type_ == "directory" {
			return fmt.Errorf(i18n.G("can't pull a directory without --recursive"))
		}

		var targetPath string
		if targetIsDir {
			targetPath = path.Join(target, path.Base(pathSpec[1]))
		} else {
			targetPath = target
		}

		var f *os.File
		if targetPath == "-" {
			f = os.Stdout
		} else {
			f, err = os.Create(targetPath)
			if err != nil {
				return err
			}
			defer f.Close()

			err = f.Chmod(os.FileMode(mode))
			if err != nil {
				return err
			}
		}

		_, err = io.Copy(f, buf)
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *fileCmd) edit(config *lxd.Config, args []string) error {
	if len(args) != 1 {
		return errArgs
	}

	if c.recursive {
		return fmt.Errorf(i18n.G("recursive edit doesn't make sense :("))
	}

	// If stdin isn't a terminal, read text from it
	if !termios.IsTerminal(int(syscall.Stdin)) {
		return c.push(config, false, append([]string{os.Stdin.Name()}, args[0]))
	}

	// Create temp file
	f, err := ioutil.TempFile("", "lxd_file_edit_")
	fname := f.Name()
	f.Close()
	os.Remove(fname)
	defer os.Remove(fname)

	// Extract current value
	err = c.pull(config, append([]string{args[0]}, fname))
	if err != nil {
		return err
	}

	_, err = shared.TextEditor(fname, []byte{})
	if err != nil {
		return err
	}

	err = c.push(config, false, append([]string{fname}, args[0]))
	if err != nil {
		return err
	}

	return nil
}

func (c *fileCmd) run(config *lxd.Config, args []string) error {
	if len(args) < 1 {
		return errArgs
	}

	switch args[0] {
	case "push":
		return c.push(config, true, args[1:])
	case "pull":
		return c.pull(config, args[1:])
	case "edit":
		return c.edit(config, args[1:])
	default:
		return errArgs
	}
}
<commit_msg>fix typo<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/gnuflag"
	"github.com/lxc/lxd/shared/i18n"
	"github.com/lxc/lxd/shared/termios"
)

type fileCmd struct {
	uid  int
	gid  int
	mode string

	recursive bool

	mkdirs bool
}

func (c *fileCmd) showByDefault() bool {
	return true
}

func (c *fileCmd) usage() string {
	return i18n.G(
		`Manage files on a container.

lxc file pull [-r|--recursive] <source> [<source>...] <target>
lxc file push [-r|--recursive] [-p|create-dirs] [--uid=UID] [--gid=GID] [--mode=MODE] <source> [<source>...] <target>
lxc file edit <file>

<source> in the case of pull, <target> in the case of push and <file> in the case of edit are <container name>/<path>

Examples:

To push /etc/hosts into the container foo:
  lxc file push /etc/hosts foo/etc/hosts

To pull /etc/hosts from the container:
  lxc file pull foo/etc/hosts .
`)
}

func (c *fileCmd) flags() {
	gnuflag.IntVar(&c.uid, "uid", -1, i18n.G("Set the file's uid on push"))
	gnuflag.IntVar(&c.gid, "gid", -1, i18n.G("Set the file's gid on push"))
	gnuflag.StringVar(&c.mode, "mode", "", i18n.G("Set the file's perms on push"))
	gnuflag.BoolVar(&c.recursive, "recursive", false, i18n.G("Recursively push or pull files"))
	gnuflag.BoolVar(&c.recursive, "r", false, i18n.G("Recursively push or pull files"))
	gnuflag.BoolVar(&c.mkdirs, "create-dirs", false, i18n.G("Create any directories necessary"))
	gnuflag.BoolVar(&c.mkdirs, "p", false, i18n.G("Create any directories necessary"))
}

func (c *fileCmd) push(config *lxd.Config, send_file_perms bool, args []string) error {
	if len(args) < 2 {
		return errArgs
	}

	target := args[len(args)-1]
	pathSpec := strings.SplitN(target, "/", 2)

	if len(pathSpec) != 2 {
		return fmt.Errorf(i18n.G("Invalid target %s"), target)
	}

	targetPath := pathSpec[1]
	remote, container := config.ParseRemoteAndContainer(pathSpec[0])

	d, err := lxd.NewClient(config, remote)
	if err != nil {
		return err
	}

	var sourcefilenames []string
	for _, fname := range args[:len(args)-1] {
		if !strings.HasPrefix(fname, "--") {
			sourcefilenames = append(sourcefilenames, fname)
		}
	}

	mode := os.FileMode(0755)
	if c.mode != "" {
		if len(c.mode) == 3 {
			c.mode = "0" + c.mode
		}

		m, err := strconv.ParseInt(c.mode, 0, 0)
		if err != nil {
			return err
		}
		mode = os.FileMode(m)
	}

	if c.recursive {
		if c.uid != -1 || c.gid != -1 || c.mode != "" {
			return fmt.Errorf(i18n.G("can't supply uid/gid/mode in recursive mode"))
		}

		for _, fname := range sourcefilenames {
			if c.mkdirs {
				if err := d.MkdirP(container, fname, mode); err != nil {
					return err
				}
			}

			if err := d.RecursivePushFile(container, fname, pathSpec[1]); err != nil {
				return err
			}
		}

		return nil
	}

	uid := 0
	if c.uid >= 0 {
		uid = c.uid
	}

	gid := 0
	if c.gid >= 0 {
		gid = c.gid
	}

	_, targetfilename := filepath.Split(targetPath)

	if (targetfilename != "") && (len(sourcefilenames) > 1) {
		return errArgs
	}

	/* Make sure all of the files are accessible by us before trying to
	 * push any of them. */
	var files []*os.File
	for _, f := range sourcefilenames {
		var file *os.File
		if f == "-" {
			file = os.Stdin
		} else {
			file, err = os.Open(f)
			if err != nil {
				return err
			}
		}

		defer file.Close()
		files = append(files, file)
	}

	for _, f := range files {
		fpath := targetPath
		if targetfilename == "" {
			fpath = path.Join(fpath, path.Base(f.Name()))
		}

		if c.mkdirs {
			if err := d.MkdirP(container, filepath.Dir(fpath), mode); err != nil {
				return err
			}
		}

		if send_file_perms {
			if c.mode == "" || c.uid == -1 || c.gid == -1 {
				fMode, fUid, fGid, err := c.getOwner(f)
				if err != nil {
					return err
				}

				if c.mode == "" {
					mode = fMode
				}

				if c.uid == -1 {
					uid = fUid
				}

				if c.gid == -1 {
					gid = fGid
				}
			}

			err = d.PushFile(container, fpath, gid, uid, fmt.Sprintf("%04o", mode.Perm()), f)
		} else {
			err = d.PushFile(container, fpath, -1, -1, "", f)
		}

		if err != nil {
			return err
		}
	}

	return nil
}

func (c *fileCmd) pull(config *lxd.Config, args []string) error {
	if len(args) < 2 {
		return errArgs
	}

	target := args[len(args)-1]
	targetIsDir := false
	sb, err := os.Stat(target)
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	/*
	 * If the path exists, just use it. If it doesn't exist, it might be a
	 * directory in one of two cases:
	 *   1. Someone explicitly put "/" at the end
	 *   2. Someone provided more than one source. In this case the target
	 *      should be a directory so we can save all the files into it.
	 */
	if err == nil {
		targetIsDir = sb.IsDir()
		if !targetIsDir && len(args)-1 > 1 {
			return fmt.Errorf(i18n.G("More than one file to download, but target is not a directory"))
		}
	} else if strings.HasSuffix(target, string(os.PathSeparator)) || len(args)-1 > 1 {
		if err := os.MkdirAll(target, 0755); err != nil {
			return err
		}
		targetIsDir = true
	}

	for _, f := range args[:len(args)-1] {
		pathSpec := strings.SplitN(f, "/", 2)
		if len(pathSpec) != 2 {
			return fmt.Errorf(i18n.G("Invalid source %s"), f)
		}

		remote, container := config.ParseRemoteAndContainer(pathSpec[0])
		d, err := lxd.NewClient(config, remote)
		if err != nil {
			return err
		}

		if c.recursive {
			if err := d.RecursivePullFile(container, pathSpec[1], target); err != nil {
				return err
			}

			continue
		}

		_, _, mode, type_, buf, _, err := d.PullFile(container, pathSpec[1])
		if err != nil {
			return err
		}

		if type_ == "directory" {
			return fmt.Errorf(i18n.G("can't pull a directory without --recursive"))
		}

		var targetPath string
		if targetIsDir {
			targetPath = path.Join(target, path.Base(pathSpec[1]))
		} else {
			targetPath = target
		}

		var f *os.File
		if targetPath == "-" {
			f = os.Stdout
		} else {
			f, err = os.Create(targetPath)
			if err != nil {
				return err
			}
			defer f.Close()

			err = f.Chmod(os.FileMode(mode))
			if err != nil {
				return err
			}
		}

		_, err = io.Copy(f, buf)
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *fileCmd) edit(config *lxd.Config, args []string) error {
	if len(args) != 1 {
		return errArgs
	}

	if c.recursive {
		return fmt.Errorf(i18n.G("recursive edit doesn't make sense :("))
	}

	// If stdin isn't a terminal, read text from it
	if !termios.IsTerminal(int(syscall.Stdin)) {
		return c.push(config, false, append([]string{os.Stdin.Name()}, args[0]))
	}

	// Create temp file
	f, err := ioutil.TempFile("", "lxd_file_edit_")
	fname := f.Name()
	f.Close()
	os.Remove(fname)
	defer os.Remove(fname)

	// Extract current value
	err = c.pull(config, append([]string{args[0]}, fname))
	if err != nil {
		return err
	}

	_, err = shared.TextEditor(fname, []byte{})
	if err != nil {
		return err
	}

	err = c.push(config, false, append([]string{fname}, args[0]))
	if err != nil {
		return err
	}

	return nil
}

func (c *fileCmd) run(config *lxd.Config, args []string) error {
	if len(args) < 1 {
		return errArgs
	}

	switch args[0] {
	case "push":
		return c.push(config, true, args[1:])
	case "pull":
		return c.pull(config, args[1:])
	case "edit":
		return c.edit(config, args[1:])
	default:
		return errArgs
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/project"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

type initDataNode struct {
	api.ServerPut `yaml:",inline"`
	Networks      []internalClusterPostNetwork `json:"networks" yaml:"networks"`
	StoragePools  []api.StoragePoolsPost       `json:"storage_pools" yaml:"storage_pools"`
	Profiles      []api.ProfilesPost           `json:"profiles" yaml:"profiles"`
	Projects      []api.ProjectsPost           `json:"projects" yaml:"projects"`
}

type initDataCluster struct {
	api.ClusterPut `yaml:",inline"`

	// The path to the cluster certificate
	// Example: /tmp/cluster.crt
	ClusterCertificatePath string `json:"cluster_certificate_path" yaml:"cluster_certificate_path"`

	// A cluster join token
	// Example: BASE64-TOKEN
	ClusterToken string `json:"cluster_token" yaml:"cluster_token"`
}

// Helper to initialize node-specific entities on a LXD instance using the
// definitions from the given initDataNode object.
//
// It's used both by the 'lxd init' command and by the PUT /1.0/cluster API.
//
// In case of error, the returned function can be used to revert the changes.
func initDataNodeApply(d lxd.InstanceServer, config initDataNode) (func(), error) {
	revert := revert.New()
	defer revert.Fail()

	// Apply server configuration.
	if config.Config != nil && len(config.Config) > 0 {
		// Get current config.
		currentServer, etag, err := d.GetServer()
		if err != nil {
			return nil, fmt.Errorf("Failed to retrieve current server configuration: %w", err)
		}

		// Setup reverter.
		revert.Add(func() { d.UpdateServer(currentServer.Writable(), "") })

		// Prepare the update.
		newServer := api.ServerPut{}
		err = shared.DeepCopy(currentServer.Writable(), &newServer)
		if err != nil {
			return nil, fmt.Errorf("Failed to copy server configuration: %w", err)
		}

		for k, v := range config.Config {
			newServer.Config[k] = fmt.Sprintf("%v", v)
		}

		// Apply it.
		err = d.UpdateServer(newServer, etag)
		if err != nil {
			return nil, fmt.Errorf("Failed to update server configuration: %w", err)
		}
	}

	// Apply storage configuration.
	if config.StoragePools != nil && len(config.StoragePools) > 0 {
		// Get the list of storagePools.
		storagePoolNames, err := d.GetStoragePoolNames()
		if err != nil {
			return nil, fmt.Errorf("Failed to retrieve list of storage pools: %w", err)
		}

		// StoragePool creator
		createStoragePool := func(storagePool api.StoragePoolsPost) error {
			// Create the storagePool if doesn't exist.
			err := d.CreateStoragePool(storagePool)
			if err != nil {
				return fmt.Errorf("Failed to create storage pool '%s': %w", storagePool.Name, err)
			}

			// Setup reverter.
			revert.Add(func() { d.DeleteStoragePool(storagePool.Name) })
			return nil
		}

		// StoragePool updater.
		updateStoragePool := func(storagePool api.StoragePoolsPost) error {
			// Get the current storagePool.
			currentStoragePool, etag, err := d.GetStoragePool(storagePool.Name)
			if err != nil {
				return fmt.Errorf("Failed to retrieve current storage pool '%s': %w", storagePool.Name, err)
			}

			// Quick check.
			if currentStoragePool.Driver != storagePool.Driver {
				return fmt.Errorf("Storage pool '%s' is of type '%s' instead of '%s'", currentStoragePool.Name, currentStoragePool.Driver, storagePool.Driver)
			}

			// Setup reverter.
			revert.Add(func() { d.UpdateStoragePool(currentStoragePool.Name, currentStoragePool.Writable(), "") })

			// Prepare the update.
			newStoragePool := api.StoragePoolPut{}
			err = shared.DeepCopy(currentStoragePool.Writable(), &newStoragePool)
			if err != nil {
				return fmt.Errorf("Failed to copy configuration of storage pool '%s': %w", storagePool.Name, err)
			}

			// Description override.
			if storagePool.Description != "" {
				newStoragePool.Description = storagePool.Description
			}

			// Config overrides.
			for k, v := range storagePool.Config {
				newStoragePool.Config[k] = fmt.Sprintf("%v", v)
			}

			// Apply it.
			err = d.UpdateStoragePool(currentStoragePool.Name, newStoragePool, etag)
			if err != nil {
				return fmt.Errorf("Failed to update storage pool '%s': %w", storagePool.Name, err)
			}

			return nil
		}

		for _, storagePool := range config.StoragePools {
			// New storagePool.
			if !shared.StringInSlice(storagePool.Name, storagePoolNames) {
				err := createStoragePool(storagePool)
				if err != nil {
					return nil, err
				}

				continue
			}

			// Existing storagePool.
			err := updateStoragePool(storagePool)
			if err != nil {
				return nil, err
			}
		}
	}

	// Apply project configuration.
	if config.Projects != nil && len(config.Projects) > 0 {
		// Get the list of projects.
		projectNames, err := d.GetProjectNames()
		if err != nil {
			return nil, fmt.Errorf("Failed to retrieve list of projects: %w", err)
		}

		// Project creator.
		createProject := func(project api.ProjectsPost) error {
			// Create the project if doesn't exist.
			err := d.CreateProject(project)
			if err != nil {
				return fmt.Errorf("Failed to create local member project '%s': %w", project.Name, err)
			}

			// Setup reverter.
			revert.Add(func() { d.DeleteProject(project.Name) })
			return nil
		}

		// Project updater.
		updateProject := func(project api.ProjectsPost) error {
			// Get the current project.
			currentProject, etag, err := d.GetProject(project.Name)
			if err != nil {
				return fmt.Errorf("Failed to retrieve current project '%s': %w", project.Name, err)
			}

			// Setup reverter.
			revert.Add(func() { d.UpdateProject(currentProject.Name, currentProject.Writable(), "") })

			// Prepare the update.
			newProject := api.ProjectPut{}
			err = shared.DeepCopy(currentProject.Writable(), &newProject)
			if err != nil {
				return fmt.Errorf("Failed to copy configuration of project '%s': %w", project.Name, err)
			}

			// Description override.
			if project.Description != "" {
				newProject.Description = project.Description
			}

			// Config overrides.
			for k, v := range project.Config {
				newProject.Config[k] = fmt.Sprintf("%v", v)
			}

			// Apply it.
			err = d.UpdateProject(currentProject.Name, newProject, etag)
			if err != nil {
				return fmt.Errorf("Failed to update local member project '%s': %w", project.Name, err)
			}

			return nil
		}

		for _, project := range config.Projects {
			// New project.
			if !shared.StringInSlice(project.Name, projectNames) {
				err := createProject(project)
				if err != nil {
					return nil, err
				}

				continue
			}

			// Existing project.
			err := updateProject(project)
			if err != nil {
				return nil, err
			}
		}
	}

	// Apply network configuration.
	if config.Networks != nil && len(config.Networks) > 0 {
		// Network creator.
		createNetwork := func(network internalClusterPostNetwork) error {
			// Create the network if doesn't exist.
			err := d.UseProject(network.Project).CreateNetwork(network.NetworksPost)
			if err != nil {
				return fmt.Errorf("Failed to create local member network %q in project %q: %w", network.Name, network.Project, err)
			}

			// Setup reverter.
			revert.Add(func() { d.UseProject(network.Project).DeleteNetwork(network.Name) })
			return nil
		}

		// Network updater.
		updateNetwork := func(network internalClusterPostNetwork) error {
			// Get the current network.
			currentNetwork, etag, err := d.UseProject(network.Project).GetNetwork(network.Name)
			if err != nil {
				return fmt.Errorf("Failed to retrieve current network %q in project %q: %w", network.Name, network.Project, err)
			}

			// Prepare the update.
			newNetwork := api.NetworkPut{}
			err = shared.DeepCopy(currentNetwork.Writable(), &newNetwork)
			if err != nil {
				return fmt.Errorf("Failed to copy configuration of network %q in project %q: %w", network.Name, network.Project, err)
			}

			// Description override.
			if network.Description != "" {
				newNetwork.Description = network.Description
			}

			// Config overrides.
			for k, v := range network.Config {
				newNetwork.Config[k] = fmt.Sprintf("%v", v)
			}

			// Apply it.
			err = d.UseProject(network.Project).UpdateNetwork(currentNetwork.Name, newNetwork, etag)
			if err != nil {
				return fmt.Errorf("Failed to update local member network %q in project %q: %w", network.Name, network.Project, err)
			}

			// Setup reverter.
			revert.Add(func() {
				d.UseProject(network.Project).UpdateNetwork(currentNetwork.Name, currentNetwork.Writable(), "")
			})

			return nil
		}

		for _, network := range config.Networks {
			// Populate default project if not specified for backwards compatbility with earlier
			// preseed dump files.
			if network.Project == "" {
				network.Project = project.Default
			}

			_, _, err := d.UseProject(network.Project).GetNetwork(network.Name)
			if err != nil {
				// New network.
				err = createNetwork(network)
				if err != nil {
					return nil, err
				}
			} else {
				// Existing network.
				err = updateNetwork(network)
				if err != nil {
					return nil, err
				}
			}
		}
	}

	// Apply profile configuration.
	if config.Profiles != nil && len(config.Profiles) > 0 {
		// Get the list of profiles.
		profileNames, err := d.GetProfileNames()
		if err != nil {
			return nil, fmt.Errorf("Failed to retrieve list of profiles: %w", err)
		}

		// Profile creator.
		createProfile := func(profile api.ProfilesPost) error {
			// Create the profile if doesn't exist.
			err := d.CreateProfile(profile)
			if err != nil {
				return fmt.Errorf("Failed to create profile '%s': %w", profile.Name, err)
			}

			// Setup reverter.
			revert.Add(func() { d.DeleteProfile(profile.Name) })
			return nil
		}

		// Profile updater.
		updateProfile := func(profile api.ProfilesPost) error {
			// Get the current profile.
			currentProfile, etag, err := d.GetProfile(profile.Name)
			if err != nil {
				return fmt.Errorf("Failed to retrieve current profile '%s': %w", profile.Name, err)
			}

			// Setup reverter.
			revert.Add(func() { d.UpdateProfile(currentProfile.Name, currentProfile.Writable(), "") })

			// Prepare the update.
			newProfile := api.ProfilePut{}
			err = shared.DeepCopy(currentProfile.Writable(), &newProfile)
			if err != nil {
				return fmt.Errorf("Failed to copy configuration of profile '%s': %w", profile.Name, err)
			}

			// Description override.
			if profile.Description != "" {
				newProfile.Description = profile.Description
			}

			// Config overrides.
			for k, v := range profile.Config {
				newProfile.Config[k] = fmt.Sprintf("%v", v)
			}

			// Device overrides.
			for k, v := range profile.Devices {
				// New device.
				_, ok := newProfile.Devices[k]
				if !ok {
					newProfile.Devices[k] = v
					continue
				}

				// Existing device.
				for configKey, configValue := range v {
					newProfile.Devices[k][configKey] = fmt.Sprintf("%v", configValue)
				}
			}

			// Apply it.
			err = d.UpdateProfile(currentProfile.Name, newProfile, etag)
			if err != nil {
				return fmt.Errorf("Failed to update profile '%s': %w", profile.Name, err)
			}

			return nil
		}

		for _, profile := range config.Profiles {
			// New profile.
			if !shared.StringInSlice(profile.Name, profileNames) {
				err := createProfile(profile)
				if err != nil {
					return nil, err
				}

				continue
			}

			// Existing profile.
			err := updateProfile(profile)
			if err != nil {
				return nil, err
			}
		}
	}

	revertExternal := revert.Clone() // Clone before calling revert.Success() so we can return the Fail func.
	revert.Success()
	return revertExternal.Fail, nil
}

// Helper to initialize LXD clustering.
//
// Used by the 'lxd init' command.
func initDataClusterApply(d lxd.InstanceServer, config *initDataCluster) error {
	if config == nil || !config.Enabled {
		return nil
	}

	// Get the current cluster configuration
	currentCluster, etag, err := d.GetCluster()
	if err != nil {
		return fmt.Errorf("Failed to retrieve current cluster config: %w", err)
	}

	// Check if already enabled
	if !currentCluster.Enabled {
		// Configure the cluster
		op, err := d.UpdateCluster(config.ClusterPut, etag)
		if err != nil {
			return fmt.Errorf("Failed to configure cluster: %w", err)
		}

		err = op.Wait()
		if err != nil {
			return fmt.Errorf("Failed to configure cluster: %w", err)
		}
	}

	return nil
}
<commit_msg>lxd/init: Error quoting in initDataNodeApply<commit_after>package main

import (
	"fmt"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/project"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

type initDataNode struct {
	api.ServerPut `yaml:",inline"`
	Networks      []internalClusterPostNetwork `json:"networks" yaml:"networks"`
	StoragePools  []api.StoragePoolsPost       `json:"storage_pools" yaml:"storage_pools"`
	Profiles      []api.ProfilesPost           `json:"profiles" yaml:"profiles"`
	Projects      []api.ProjectsPost           `json:"projects" yaml:"projects"`
}

type initDataCluster struct {
	api.ClusterPut `yaml:",inline"`

	// The path to the cluster certificate
	// Example: /tmp/cluster.crt
	ClusterCertificatePath string `json:"cluster_certificate_path" yaml:"cluster_certificate_path"`

	// A cluster join token
	// Example: BASE64-TOKEN
	ClusterToken string `json:"cluster_token" yaml:"cluster_token"`
}

// Helper to initialize node-specific entities on a LXD instance using the
// definitions from the given initDataNode object.
//
// It's used both by the 'lxd init' command and by the PUT /1.0/cluster API.
//
// In case of error, the returned function can be used to revert the changes.
func initDataNodeApply(d lxd.InstanceServer, config initDataNode) (func(), error) {
	revert := revert.New()
	defer revert.Fail()

	// Apply server configuration.
	if config.Config != nil && len(config.Config) > 0 {
		// Get current config.
		currentServer, etag, err := d.GetServer()
		if err != nil {
			return nil, fmt.Errorf("Failed to retrieve current server configuration: %w", err)
		}

		// Setup reverter.
		revert.Add(func() { d.UpdateServer(currentServer.Writable(), "") })

		// Prepare the update.
		newServer := api.ServerPut{}
		err = shared.DeepCopy(currentServer.Writable(), &newServer)
		if err != nil {
			return nil, fmt.Errorf("Failed to copy server configuration: %w", err)
		}

		for k, v := range config.Config {
			newServer.Config[k] = fmt.Sprintf("%v", v)
		}

		// Apply it.
		err = d.UpdateServer(newServer, etag)
		if err != nil {
			return nil, fmt.Errorf("Failed to update server configuration: %w", err)
		}
	}

	// Apply storage configuration.
	if config.StoragePools != nil && len(config.StoragePools) > 0 {
		// Get the list of storagePools.
		storagePoolNames, err := d.GetStoragePoolNames()
		if err != nil {
			return nil, fmt.Errorf("Failed to retrieve list of storage pools: %w", err)
		}

		// StoragePool creator
		createStoragePool := func(storagePool api.StoragePoolsPost) error {
			// Create the storagePool if doesn't exist.
			err := d.CreateStoragePool(storagePool)
			if err != nil {
				return fmt.Errorf("Failed to create storage pool %q: %w", storagePool.Name, err)
			}

			// Setup reverter.
			revert.Add(func() { d.DeleteStoragePool(storagePool.Name) })
			return nil
		}

		// StoragePool updater.
		updateStoragePool := func(storagePool api.StoragePoolsPost) error {
			// Get the current storagePool.
			currentStoragePool, etag, err := d.GetStoragePool(storagePool.Name)
			if err != nil {
				return fmt.Errorf("Failed to retrieve current storage pool %q: %w", storagePool.Name, err)
			}

			// Quick check.
			if currentStoragePool.Driver != storagePool.Driver {
				return fmt.Errorf("Storage pool %q is of type %q instead of %q", currentStoragePool.Name, currentStoragePool.Driver, storagePool.Driver)
			}

			// Setup reverter.
			revert.Add(func() { d.UpdateStoragePool(currentStoragePool.Name, currentStoragePool.Writable(), "") })

			// Prepare the update.
			newStoragePool := api.StoragePoolPut{}
			err = shared.DeepCopy(currentStoragePool.Writable(), &newStoragePool)
			if err != nil {
				return fmt.Errorf("Failed to copy configuration of storage pool %q: %w", storagePool.Name, err)
			}

			// Description override.
			if storagePool.Description != "" {
				newStoragePool.Description = storagePool.Description
			}

			// Config overrides.
			for k, v := range storagePool.Config {
				newStoragePool.Config[k] = fmt.Sprintf("%v", v)
			}

			// Apply it.
			err = d.UpdateStoragePool(currentStoragePool.Name, newStoragePool, etag)
			if err != nil {
				return fmt.Errorf("Failed to update storage pool %q: %w", storagePool.Name, err)
			}

			return nil
		}

		for _, storagePool := range config.StoragePools {
			// New storagePool.
			if !shared.StringInSlice(storagePool.Name, storagePoolNames) {
				err := createStoragePool(storagePool)
				if err != nil {
					return nil, err
				}

				continue
			}

			// Existing storagePool.
			err := updateStoragePool(storagePool)
			if err != nil {
				return nil, err
			}
		}
	}

	// Apply project configuration.
	if config.Projects != nil && len(config.Projects) > 0 {
		// Get the list of projects.
		projectNames, err := d.GetProjectNames()
		if err != nil {
			return nil, fmt.Errorf("Failed to retrieve list of projects: %w", err)
		}

		// Project creator.
		createProject := func(project api.ProjectsPost) error {
			// Create the project if doesn't exist.
			err := d.CreateProject(project)
			if err != nil {
				return fmt.Errorf("Failed to create local member project %q: %w", project.Name, err)
			}

			// Setup reverter.
			revert.Add(func() { d.DeleteProject(project.Name) })
			return nil
		}

		// Project updater.
		updateProject := func(project api.ProjectsPost) error {
			// Get the current project.
			currentProject, etag, err := d.GetProject(project.Name)
			if err != nil {
				return fmt.Errorf("Failed to retrieve current project %q: %w", project.Name, err)
			}

			// Setup reverter.
			revert.Add(func() { d.UpdateProject(currentProject.Name, currentProject.Writable(), "") })

			// Prepare the update.
			newProject := api.ProjectPut{}
			err = shared.DeepCopy(currentProject.Writable(), &newProject)
			if err != nil {
				return fmt.Errorf("Failed to copy configuration of project %q: %w", project.Name, err)
			}

			// Description override.
			if project.Description != "" {
				newProject.Description = project.Description
			}

			// Config overrides.
			for k, v := range project.Config {
				newProject.Config[k] = fmt.Sprintf("%v", v)
			}

			// Apply it.
			err = d.UpdateProject(currentProject.Name, newProject, etag)
			if err != nil {
				return fmt.Errorf("Failed to update local member project %q: %w", project.Name, err)
			}

			return nil
		}

		for _, project := range config.Projects {
			// New project.
			if !shared.StringInSlice(project.Name, projectNames) {
				err := createProject(project)
				if err != nil {
					return nil, err
				}

				continue
			}

			// Existing project.
			err := updateProject(project)
			if err != nil {
				return nil, err
			}
		}
	}

	// Apply network configuration.
	if config.Networks != nil && len(config.Networks) > 0 {
		// Network creator.
		createNetwork := func(network internalClusterPostNetwork) error {
			// Create the network if doesn't exist.
			err := d.UseProject(network.Project).CreateNetwork(network.NetworksPost)
			if err != nil {
				return fmt.Errorf("Failed to create local member network %q in project %q: %w", network.Name, network.Project, err)
			}

			// Setup reverter.
			revert.Add(func() { d.UseProject(network.Project).DeleteNetwork(network.Name) })
			return nil
		}

		// Network updater.
		updateNetwork := func(network internalClusterPostNetwork) error {
			// Get the current network.
			currentNetwork, etag, err := d.UseProject(network.Project).GetNetwork(network.Name)
			if err != nil {
				return fmt.Errorf("Failed to retrieve current network %q in project %q: %w", network.Name, network.Project, err)
			}

			// Prepare the update.
			newNetwork := api.NetworkPut{}
			err = shared.DeepCopy(currentNetwork.Writable(), &newNetwork)
			if err != nil {
				return fmt.Errorf("Failed to copy configuration of network %q in project %q: %w", network.Name, network.Project, err)
			}

			// Description override.
			if network.Description != "" {
				newNetwork.Description = network.Description
			}

			// Config overrides.
			for k, v := range network.Config {
				newNetwork.Config[k] = fmt.Sprintf("%v", v)
			}

			// Apply it.
			err = d.UseProject(network.Project).UpdateNetwork(currentNetwork.Name, newNetwork, etag)
			if err != nil {
				return fmt.Errorf("Failed to update local member network %q in project %q: %w", network.Name, network.Project, err)
			}

			// Setup reverter.
			revert.Add(func() {
				d.UseProject(network.Project).UpdateNetwork(currentNetwork.Name, currentNetwork.Writable(), "")
			})

			return nil
		}

		for _, network := range config.Networks {
			// Populate default project if not specified for backwards compatbility with earlier
			// preseed dump files.
			if network.Project == "" {
				network.Project = project.Default
			}

			_, _, err := d.UseProject(network.Project).GetNetwork(network.Name)
			if err != nil {
				// New network.
				err = createNetwork(network)
				if err != nil {
					return nil, err
				}
			} else {
				// Existing network.
				err = updateNetwork(network)
				if err != nil {
					return nil, err
				}
			}
		}
	}

	// Apply profile configuration.
	if config.Profiles != nil && len(config.Profiles) > 0 {
		// Get the list of profiles.
		profileNames, err := d.GetProfileNames()
		if err != nil {
			return nil, fmt.Errorf("Failed to retrieve list of profiles: %w", err)
		}

		// Profile creator.
		createProfile := func(profile api.ProfilesPost) error {
			// Create the profile if doesn't exist.
			err := d.CreateProfile(profile)
			if err != nil {
				return fmt.Errorf("Failed to create profile %q: %w", profile.Name, err)
			}

			// Setup reverter.
			revert.Add(func() { d.DeleteProfile(profile.Name) })
			return nil
		}

		// Profile updater.
		updateProfile := func(profile api.ProfilesPost) error {
			// Get the current profile.
			currentProfile, etag, err := d.GetProfile(profile.Name)
			if err != nil {
				return fmt.Errorf("Failed to retrieve current profile %q: %w", profile.Name, err)
			}

			// Setup reverter.
			revert.Add(func() { d.UpdateProfile(currentProfile.Name, currentProfile.Writable(), "") })

			// Prepare the update.
			newProfile := api.ProfilePut{}
			err = shared.DeepCopy(currentProfile.Writable(), &newProfile)
			if err != nil {
				return fmt.Errorf("Failed to copy configuration of profile %q: %w", profile.Name, err)
			}

			// Description override.
			if profile.Description != "" {
				newProfile.Description = profile.Description
			}

			// Config overrides.
			for k, v := range profile.Config {
				newProfile.Config[k] = fmt.Sprintf("%v", v)
			}

			// Device overrides.
			for k, v := range profile.Devices {
				// New device.
				_, ok := newProfile.Devices[k]
				if !ok {
					newProfile.Devices[k] = v
					continue
				}

				// Existing device.
				for configKey, configValue := range v {
					newProfile.Devices[k][configKey] = fmt.Sprintf("%v", configValue)
				}
			}

			// Apply it.
			err = d.UpdateProfile(currentProfile.Name, newProfile, etag)
			if err != nil {
				return fmt.Errorf("Failed to update profile %q: %w", profile.Name, err)
			}

			return nil
		}

		for _, profile := range config.Profiles {
			// New profile.
			if !shared.StringInSlice(profile.Name, profileNames) {
				err := createProfile(profile)
				if err != nil {
					return nil, err
				}

				continue
			}

			// Existing profile.
			err := updateProfile(profile)
			if err != nil {
				return nil, err
			}
		}
	}

	revertExternal := revert.Clone() // Clone before calling revert.Success() so we can return the Fail func.
	revert.Success()
	return revertExternal.Fail, nil
}

// Helper to initialize LXD clustering.
//
// Used by the 'lxd init' command.
func initDataClusterApply(d lxd.InstanceServer, config *initDataCluster) error {
	if config == nil || !config.Enabled {
		return nil
	}

	// Get the current cluster configuration
	currentCluster, etag, err := d.GetCluster()
	if err != nil {
		return fmt.Errorf("Failed to retrieve current cluster config: %w", err)
	}

	// Check if already enabled
	if !currentCluster.Enabled {
		// Configure the cluster
		op, err := d.UpdateCluster(config.ClusterPut, etag)
		if err != nil {
			return fmt.Errorf("Failed to configure cluster: %w", err)
		}

		err = op.Wait()
		if err != nil {
			return fmt.Errorf("Failed to configure cluster: %w", err)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package scraperService

import (
	"fmt"
	"net"
	"net/url"
	"time"

	"github.com/NyaaPantsu/nyaa/config"
	"github.com/NyaaPantsu/nyaa/db"
	"github.com/NyaaPantsu/nyaa/model"
	"github.com/NyaaPantsu/nyaa/util/log"
)

// MTU yes this is the ipv6 mtu
const MTU = 1500

// max number of scrapes per packet
const ScrapesPerPacket = 74

// bittorrent scraper
type Scraper struct {
	done             chan int
	sendQueue        chan *SendEvent
	recvQueue        chan *RecvEvent
	errQueue         chan error
	trackers         map[string]*Bucket
	ticker           *time.Ticker
	cleanup          *time.Ticker
	interval         time.Duration
	PacketsPerSecond uint
}

func New(conf *config.ScraperConfig) (sc *Scraper, err error) {
	sc = &Scraper{
		done:      make(chan int),
		sendQueue: make(chan *SendEvent, 1024),
		recvQueue: make(chan *RecvEvent, 1024),
		errQueue:  make(chan error),
		trackers:  make(map[string]*Bucket),
		ticker:    time.NewTicker(time.Second * 10),
		interval:  time.Second * time.Duration(conf.IntervalSeconds),
		cleanup:   time.NewTicker(time.Minute),
	}

	if sc.PacketsPerSecond == 0 {
		sc.PacketsPerSecond = 10
	}

	for idx := range conf.Trackers {
		err = sc.AddTracker(&conf.Trackers[idx])
		if err != nil {
			break
		}
	}
	return
}

func (sc *Scraper) AddTracker(conf *config.ScrapeConfig) (err error) {
	var u *url.URL
	u, err = url.Parse(conf.URL)
	if err == nil {
		var ips []net.IP
		ips, err = net.LookupIP(u.Hostname())
		if err == nil {
			// TODO: use more than 1 ip ?
			addr := &net.UDPAddr{
				IP: ips[0],
			}
			addr.Port, err = net.LookupPort("udp", u.Port())
			if err == nil {
				sc.trackers[addr.String()] = NewBucket(addr)
			}
		}
	}
	return
}

func (sc *Scraper) Close() (err error) {
	close(sc.sendQueue)
	close(sc.recvQueue)
	close(sc.errQueue)
	sc.ticker.Stop()
	sc.done <- 1
	return
}

func (sc *Scraper) runRecv(pc net.PacketConn) {
	for {
		var buff [MTU]byte
		n, from, err := pc.ReadFrom(buff[:])

		if err == nil {

			log.Debugf("got %d from %s", n, from)
			sc.recvQueue <- &RecvEvent{
				From: from,
				Data: buff[:n],
			}
		} else {
			sc.errQueue <- err
		}
	}
}

func (sc *Scraper) runSend(pc net.PacketConn) {
	for {
		ev, ok := <-sc.sendQueue
		if !ok {
			return
		}
		log.Debugf("write %d to %s", len(ev.Data), ev.To)
		pc.WriteTo(ev.Data, ev.To)
	}
}

func (sc *Scraper) RunWorker(pc net.PacketConn) (err error) {

	go sc.runRecv(pc)
	go sc.runSend(pc)
	for {
		var bucket *Bucket
		ev, ok := <-sc.recvQueue
		if !ok {
			break
		}
		tid, err := ev.TID()
		if err != nil {
			log.Warnf("failed: %s", err)
			break
		}
		action, err := ev.Action()
		if err != nil {
			log.Warnf("failed: %s", err)
			break
		}
		log.Debugf("transaction = %d action = %d", tid, action)
		bucket, ok = sc.trackers[ev.From.String()]
		if !ok || bucket == nil {
			log.Warnf("bucket not found for %s", ev.From)
			break
		}

		bucket.VisitTransaction(tid, func(t *Transaction) {
			if t == nil {
				log.Warnf("no transaction %d", tid)
				return
			}
			if t.GotData(ev.Data) {
				err := t.Sync()
				if err != nil {
					log.Warnf("failed to sync swarm: %s", err)
				}
				t.Done()
				log.Debugf("transaction %d done", tid)
			} else {
				sc.sendQueue <- t.SendEvent(ev.From)
			}
		})
	}
	return
}

func (sc *Scraper) Run() {
	for {
		select {
		case <-sc.ticker.C:
			sc.Scrape(sc.PacketsPerSecond)
			break
		case <-sc.cleanup.C:
			sc.removeStale()
			break
		}
	}
}

func (sc *Scraper) removeStale() {

	for k := range sc.trackers {
		sc.trackers[k].ForEachTransaction(func(tid uint32, t *Transaction) {
			if t == nil || t.IsTimedOut() {
				sc.trackers[k].Forget(tid)
			}
		})
	}
}

func (sc *Scraper) Scrape(packets uint) {
	now := time.Now().Add(0 - sc.interval)
	// only scrape torretns uploaded within 90 days
	oldest := now.Add(0 - (time.Hour * 24 * 90))

	query := fmt.Sprintf(
		"SELECT * FROM ("+

			// previously scraped torrents that will be scraped again:
			"SELECT %[1]s.torrent_id, torrent_hash FROM %[1]s, %[2]s WHERE "+
			"date > ? AND "+
			"%[1]s.torrent_id = %[2]s.torrent_id AND "+
			"scrape.last_scrape < ?"+

			// torrents that weren't scraped before:
			" UNION "+
			"SELECT torrent_id, torrent_hash FROM %[1]s WHERE "+
			"date > ? AND "+
			"torrent_id NOT IN (SELECT torrent_id FROM %[2]s)"+

			") AS x ORDER BY torrent_id DESC LIMIT ?",
		config.Conf.Models.TorrentsTableName, config.Conf.Models.ScrapeTableName)
	rows, err := db.ORM.Raw(query, oldest, now, oldest, packets*ScrapesPerPacket).Rows()

	if err == nil {
		counter := 0
		var scrape [ScrapesPerPacket]model.Torrent
		for rows.Next() {
			idx := counter % ScrapesPerPacket
			rows.Scan(&scrape[idx].ID, &scrape[idx].Hash)
			counter++
			if counter%ScrapesPerPacket == 0 {
				for _, b := range sc.trackers {
					t := b.NewTransaction(scrape[:])
					sc.sendQueue <- t.SendEvent(b.Addr)
				}
			}
		}
		idx := counter % ScrapesPerPacket
		if idx > 0 {
			for _, b := range sc.trackers {
				t := b.NewTransaction(scrape[:idx])
				sc.sendQueue <- t.SendEvent(b.Addr)
			}
		}
		log.Infof("scrape %d", counter)
		rows.Close()
	} else {
		log.Warnf("failed to select torrents for scrape: %s", err)
	}
}

func (sc *Scraper) Wait() {
	<-sc.done
}
<commit_msg>Actually fix scraper<commit_after>package scraperService

import (
	"fmt"
	"net"
	"net/url"
	"time"

	"github.com/NyaaPantsu/nyaa/config"
	"github.com/NyaaPantsu/nyaa/db"
	"github.com/NyaaPantsu/nyaa/model"
	"github.com/NyaaPantsu/nyaa/util/log"
)

// MTU yes this is the ipv6 mtu
const MTU = 1500

// max number of scrapes per packet
const ScrapesPerPacket = 74

// bittorrent scraper
type Scraper struct {
	done             chan int
	sendQueue        chan *SendEvent
	recvQueue        chan *RecvEvent
	errQueue         chan error
	trackers         map[string]*Bucket
	ticker           *time.Ticker
	cleanup          *time.Ticker
	interval         time.Duration
	PacketsPerSecond uint
}

func New(conf *config.ScraperConfig) (sc *Scraper, err error) {
	sc = &Scraper{
		done:      make(chan int),
		sendQueue: make(chan *SendEvent, 1024),
		recvQueue: make(chan *RecvEvent, 1024),
		errQueue:  make(chan error),
		trackers:  make(map[string]*Bucket),
		ticker:    time.NewTicker(time.Second * 10),
		interval:  time.Second * time.Duration(conf.IntervalSeconds),
		cleanup:   time.NewTicker(time.Minute),
	}

	if sc.PacketsPerSecond == 0 {
		sc.PacketsPerSecond = 10
	}

	for idx := range conf.Trackers {
		err = sc.AddTracker(&conf.Trackers[idx])
		if err != nil {
			break
		}
	}
	return
}

func (sc *Scraper) AddTracker(conf *config.ScrapeConfig) (err error) {
	var u *url.URL
	u, err = url.Parse(conf.URL)
	if err == nil {
		var ips []net.IP
		ips, err = net.LookupIP(u.Hostname())
		if err == nil {
			// TODO: use more than 1 ip ?
			addr := &net.UDPAddr{
				IP: ips[0],
			}
			addr.Port, err = net.LookupPort("udp", u.Port())
			if err == nil {
				sc.trackers[addr.String()] = NewBucket(addr)
			}
		}
	}
	return
}

func (sc *Scraper) Close() (err error) {
	close(sc.sendQueue)
	close(sc.recvQueue)
	close(sc.errQueue)
	sc.ticker.Stop()
	sc.done <- 1
	return
}

func (sc *Scraper) runRecv(pc net.PacketConn) {
	for {
		var buff [MTU]byte
		n, from, err := pc.ReadFrom(buff[:])

		if err == nil {

			log.Debugf("got %d from %s", n, from)
			sc.recvQueue <- &RecvEvent{
				From: from,
				Data: buff[:n],
			}
		} else {
			sc.errQueue <- err
		}
	}
}

func (sc *Scraper) runSend(pc net.PacketConn) {
	for {
		ev, ok := <-sc.sendQueue
		if !ok {
			return
		}
		log.Debugf("write %d to %s", len(ev.Data), ev.To)
		pc.WriteTo(ev.Data, ev.To)
	}
}

func (sc *Scraper) RunWorker(pc net.PacketConn) (err error) {

	go sc.runRecv(pc)
	go sc.runSend(pc)
	for {
		var bucket *Bucket
		ev, ok := <-sc.recvQueue
		if !ok {
			break
		}
		tid, err := ev.TID()
		if err != nil {
			log.Warnf("failed: %s", err)
			break
		}
		action, err := ev.Action()
		if err != nil {
			log.Warnf("failed: %s", err)
			break
		}
		log.Debugf("transaction = %d action = %d", tid, action)
		bucket, ok = sc.trackers[ev.From.String()]
		if !ok || bucket == nil {
			log.Warnf("bucket not found for %s", ev.From)
			break
		}

		bucket.VisitTransaction(tid, func(t *Transaction) {
			if t == nil {
				log.Warnf("no transaction %d", tid)
				return
			}
			if t.GotData(ev.Data) {
				err := t.Sync()
				if err != nil {
					log.Warnf("failed to sync swarm: %s", err)
				}
				t.Done()
				log.Debugf("transaction %d done", tid)
			} else {
				sc.sendQueue <- t.SendEvent(ev.From)
			}
		})
	}
	return
}

func (sc *Scraper) Run() {
	for {
		select {
		case <-sc.ticker.C:
			sc.Scrape(sc.PacketsPerSecond)
			break
		case <-sc.cleanup.C:
			sc.removeStale()
			break
		}
	}
}

func (sc *Scraper) removeStale() {

	for k := range sc.trackers {
		sc.trackers[k].ForEachTransaction(func(tid uint32, t *Transaction) {
			if t == nil || t.IsTimedOut() {
				sc.trackers[k].Forget(tid)
			}
		})
	}
}

func (sc *Scraper) Scrape(packets uint) {
	now := time.Now().Add(0 - sc.interval)
	// only scrape torretns uploaded within 90 days
	oldest := now.Add(0 - (time.Hour * 24 * 90))

	query := fmt.Sprintf(
		"SELECT * FROM ("+

			// previously scraped torrents that will be scraped again:
			"SELECT %[1]s.torrent_id, torrent_hash FROM %[1]s, %[2]s WHERE "+
			"date > ? AND "+
			"%[1]s.torrent_id = %[2]s.torrent_id AND "+
			"$[2]s.last_scrape < ?"+

			// torrents that weren't scraped before:
			" UNION "+
			"SELECT torrent_id, torrent_hash FROM %[1]s WHERE "+
			"date > ? AND "+
			"torrent_id NOT IN (SELECT torrent_id FROM %[2]s)"+

			") AS x ORDER BY torrent_id DESC LIMIT ?",
		config.Conf.Models.TorrentsTableName, config.Conf.Models.ScrapeTableName)
	rows, err := db.ORM.Raw(query, oldest, now, oldest, packets*ScrapesPerPacket).Rows()

	if err == nil {
		counter := 0
		var scrape [ScrapesPerPacket]model.Torrent
		for rows.Next() {
			idx := counter % ScrapesPerPacket
			rows.Scan(&scrape[idx].ID, &scrape[idx].Hash)
			counter++
			if counter%ScrapesPerPacket == 0 {
				for _, b := range sc.trackers {
					t := b.NewTransaction(scrape[:])
					sc.sendQueue <- t.SendEvent(b.Addr)
				}
			}
		}
		idx := counter % ScrapesPerPacket
		if idx > 0 {
			for _, b := range sc.trackers {
				t := b.NewTransaction(scrape[:idx])
				sc.sendQueue <- t.SendEvent(b.Addr)
			}
		}
		log.Infof("scrape %d", counter)
		rows.Close()
	} else {
		log.Warnf("failed to select torrents for scrape: %s", err)
	}
}

func (sc *Scraper) Wait() {
	<-sc.done
}
<|endoftext|>