<commit_before>/*
Copyright 2020 The cert-manager Authors.

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

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

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

package framework

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"

	admissionregistrationv1beta1 "k8s.io/api/admissionregistration/v1beta1"
	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions"
	apiextensionsinstall "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/install"
	v1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	jsonserializer "k8s.io/apimachinery/pkg/runtime/serializer/json"
	"k8s.io/apimachinery/pkg/runtime/serializer/versioning"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/client-go/rest"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/envtest"

	webhooktesting "github.com/jetstack/cert-manager/cmd/webhook/app/testing"
	"github.com/jetstack/cert-manager/pkg/api"
	apitesting "github.com/jetstack/cert-manager/pkg/api/testing"
)

func init() {
	// Set environment variables for controller-runtime's envtest package.
	// This is done once as we cannot scope environment variables to a single
	// invocation of RunControlPlane due to envtest's design.
	setUpEnvTestEnv()
}

type StopFunc func()

func RunControlPlane(t *testing.T) (*rest.Config, StopFunc) {
	webhookOpts, stopWebhook := webhooktesting.StartWebhookServer(t, []string{})
	crdsDir := apitesting.CRDDirectory(t)
	crds := readCustomResourcesAtPath(t, crdsDir)
	for _, crd := range crds {
		t.Logf("Found CRD with name %q", crd.Name)
	}
	patchCRDConversion(crds, webhookOpts.URL, webhookOpts.CAPEM)

	env := &envtest.Environment{
		AttachControlPlaneOutput: false,
		CRDs:                     crdsToRuntimeObjects(crds),
	}

	config, err := env.Start()
	if err != nil {
		t.Fatalf("failed to start control plane: %v", err)
	}

	cl, err := client.New(config, client.Options{Scheme: api.Scheme})
	if err != nil {
		t.Fatal(err)
	}

	// installing the validating webhooks, not using WebhookInstallOptions as it patches the CA to be it's own
	err = cl.Create(context.Background(), getValidatingWebhookConfig(webhookOpts.URL, webhookOpts.CAPEM))
	if err != nil {
		t.Fatal(err)
	}

	// installing the mutating webhooks, not using WebhookInstallOptions as it patches the CA to be it's own
	err = cl.Create(context.Background(), getMutatingWebhookConfig(webhookOpts.URL, webhookOpts.CAPEM))
	if err != nil {
		t.Fatal(err)
	}

	return config, func() {
		defer stopWebhook()
		if err := env.Stop(); err != nil {
			t.Logf("failed to shut down control plane, not failing test: %v", err)
		}
	}
}

var (
	internalScheme = runtime.NewScheme()
)

func init() {
	utilruntime.Must(metav1.AddMetaToScheme(internalScheme))
	apiextensionsinstall.Install(internalScheme)
}

func patchCRDConversion(crds []*v1.CustomResourceDefinition, url string, caPEM []byte) {
	for _, crd := range crds {
		if crd.Spec.Conversion == nil {
			continue
		}
		if crd.Spec.Conversion.Webhook == nil {
			continue
		}
		if crd.Spec.Conversion.Webhook.ClientConfig == nil {
			continue
		}
		if crd.Spec.Conversion.Webhook.ClientConfig.Service == nil {
			continue
		}
		path := ""
		if crd.Spec.Conversion.Webhook.ClientConfig.Service.Path != nil {
			path = *crd.Spec.Conversion.Webhook.ClientConfig.Service.Path
		}
		url := fmt.Sprintf("%s%s", url, path)
		crd.Spec.Conversion.Webhook.ClientConfig.URL = &url
		crd.Spec.Conversion.Webhook.ClientConfig.CABundle = caPEM
		crd.Spec.Conversion.Webhook.ClientConfig.Service = nil
	}
}

func readCustomResourcesAtPath(t *testing.T, path string) []*v1.CustomResourceDefinition {
	serializer := jsonserializer.NewSerializerWithOptions(jsonserializer.DefaultMetaFactory, internalScheme, internalScheme, jsonserializer.SerializerOptions{
		Yaml: true,
	})
	converter := runtime.UnsafeObjectConvertor(internalScheme)
	codec := versioning.NewCodec(serializer, serializer, converter, internalScheme, internalScheme, internalScheme, runtime.InternalGroupVersioner, runtime.InternalGroupVersioner, internalScheme.Name())

	var crds []*v1.CustomResourceDefinition
	if err := filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if filepath.Ext(path) != ".yaml" {
			return nil
		}
		crd, err := readCRDsAtPath(codec, converter, path)
		if err != nil {
			return err
		}
		crds = append(crds, crd...)
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	return crds
}

func readCRDsAtPath(codec runtime.Codec, converter runtime.ObjectConvertor, path string) ([]*v1.CustomResourceDefinition, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	var crds []*v1.CustomResourceDefinition
	for _, d := range strings.Split(string(data), "\n---\n") {
		// skip empty YAML documents
		if strings.TrimSpace(d) == "" {
			continue
		}

		internalCRD := &apiextensions.CustomResourceDefinition{}
		if _, _, err := codec.Decode([]byte(d), nil, internalCRD); err != nil {
			return nil, err
		}

		out := &v1.CustomResourceDefinition{}
		if err := converter.Convert(internalCRD, out, nil); err != nil {
			return nil, err
		}

		crds = append(crds, out)
	}

	return crds, nil
}

func crdsToRuntimeObjects(in []*v1.CustomResourceDefinition) []runtime.Object {
	out := make([]runtime.Object, len(in))

	for i, crd := range in {
		out[i] = runtime.Object(crd)
	}

	return out
}

func getValidatingWebhookConfig(url string, caPEM []byte) runtime.Object {
	failurePolicy := admissionregistrationv1beta1.Fail
	sideEffects := admissionregistrationv1beta1.SideEffectClassNone
	validateURL := fmt.Sprintf("%s/validate", url)
	webhook := admissionregistrationv1beta1.ValidatingWebhookConfiguration{
		ObjectMeta: metav1.ObjectMeta{
			Name: "cert-manager-webhook",
		},
		Webhooks: []admissionregistrationv1beta1.ValidatingWebhook{
			{
				Name: "webhook.cert-manager.io",
				ClientConfig: admissionregistrationv1beta1.WebhookClientConfig{
					URL:      &validateURL,
					CABundle: caPEM,
				},
				Rules: []admissionregistrationv1beta1.RuleWithOperations{
					{
						Operations: []admissionregistrationv1beta1.OperationType{
							admissionregistrationv1beta1.Create,
							admissionregistrationv1beta1.Update,
						},
						Rule: admissionregistrationv1beta1.Rule{
							APIGroups:   []string{"cert-manager.io", "acme.cert-manager.io"},
							APIVersions: []string{"*"},
							Resources:   []string{"*/*"},
						},
					},
				},
				FailurePolicy: &failurePolicy,
				SideEffects:   &sideEffects,
			},
		},
	}

	return &webhook
}

func getMutatingWebhookConfig(url string, caPEM []byte) runtime.Object {
	failurePolicy := admissionregistrationv1beta1.Fail
	sideEffects := admissionregistrationv1beta1.SideEffectClassNone
	validateURL := fmt.Sprintf("%s/mutate", url)
	webhook := admissionregistrationv1beta1.MutatingWebhookConfiguration{
		ObjectMeta: metav1.ObjectMeta{
			Name: "cert-manager-webhook",
		},
		Webhooks: []admissionregistrationv1beta1.MutatingWebhook{
			{
				Name: "webhook.cert-manager.io",
				ClientConfig: admissionregistrationv1beta1.WebhookClientConfig{
					URL:      &validateURL,
					CABundle: caPEM,
				},
				Rules: []admissionregistrationv1beta1.RuleWithOperations{
					{
						Operations: []admissionregistrationv1beta1.OperationType{
							admissionregistrationv1beta1.Create,
							admissionregistrationv1beta1.Update,
						},
						Rule: admissionregistrationv1beta1.Rule{
							APIGroups:   []string{"cert-manager.io", "acme.cert-manager.io"},
							APIVersions: []string{"*"},
							Resources:   []string{"*/*"},
						},
					},
				},
				FailurePolicy: &failurePolicy,
				SideEffects:   &sideEffects,
			},
		},
	}

	return &webhook
}
<commit_msg>Update integration test framework to restart the API to share the address with the webhook<commit_after>/*
Copyright 2020 The cert-manager Authors.

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

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

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

package framework

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"

	admissionregistrationv1beta1 "k8s.io/api/admissionregistration/v1beta1"
	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions"
	apiextensionsinstall "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/install"
	v1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	jsonserializer "k8s.io/apimachinery/pkg/runtime/serializer/json"
	"k8s.io/apimachinery/pkg/runtime/serializer/versioning"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/client-go/rest"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/envtest"

	webhooktesting "github.com/jetstack/cert-manager/cmd/webhook/app/testing"
	"github.com/jetstack/cert-manager/pkg/api"
	apitesting "github.com/jetstack/cert-manager/pkg/api/testing"
)

func init() {
	// Set environment variables for controller-runtime's envtest package.
	// This is done once as we cannot scope environment variables to a single
	// invocation of RunControlPlane due to envtest's design.
	setUpEnvTestEnv()
}

type StopFunc func()

func RunControlPlane(t *testing.T) (*rest.Config, StopFunc) {
	// Here we start the API server so its address can be given to the webhook on
	// start. We then restart the API with the CRDs in the webhook.
	env := &envtest.Environment{
		AttachControlPlaneOutput: false,
	}

	config, err := env.Start()
	if err != nil {
		t.Fatalf("failed to start control plane: %v", err)
	}

	if err := env.Stop(); err != nil {
		t.Fatal(err)
	}

	webhookOpts, stopWebhook := webhooktesting.StartWebhookServer(t, []string{"--master=" + config.Host})
	crdsDir := apitesting.CRDDirectory(t)
	crds := readCustomResourcesAtPath(t, crdsDir)
	for _, crd := range crds {
		t.Logf("Found CRD with name %q", crd.Name)
	}
	patchCRDConversion(crds, webhookOpts.URL, webhookOpts.CAPEM)

	env.CRDs = crdsToRuntimeObjects(crds)
	env.Config = config

	config, err = env.Start()
	if err != nil {
		t.Fatalf("failed to start control plane: %v", err)
	}

	cl, err := client.New(config, client.Options{Scheme: api.Scheme})
	if err != nil {
		t.Fatal(err)
	}

	// installing the validating webhooks, not using WebhookInstallOptions as it patches the CA to be it's own
	err = cl.Create(context.Background(), getValidatingWebhookConfig(webhookOpts.URL, webhookOpts.CAPEM))
	if err != nil {
		t.Fatal(err)
	}

	// installing the mutating webhooks, not using WebhookInstallOptions as it patches the CA to be it's own
	err = cl.Create(context.Background(), getMutatingWebhookConfig(webhookOpts.URL, webhookOpts.CAPEM))
	if err != nil {
		t.Fatal(err)
	}

	return config, func() {
		defer stopWebhook()
		if err := env.Stop(); err != nil {
			t.Logf("failed to shut down control plane, not failing test: %v", err)
		}
	}
}

var (
	internalScheme = runtime.NewScheme()
)

func init() {
	utilruntime.Must(metav1.AddMetaToScheme(internalScheme))
	apiextensionsinstall.Install(internalScheme)
}

func patchCRDConversion(crds []*v1.CustomResourceDefinition, url string, caPEM []byte) {
	for _, crd := range crds {
		if crd.Spec.Conversion == nil {
			continue
		}
		if crd.Spec.Conversion.Webhook == nil {
			continue
		}
		if crd.Spec.Conversion.Webhook.ClientConfig == nil {
			continue
		}
		if crd.Spec.Conversion.Webhook.ClientConfig.Service == nil {
			continue
		}
		path := ""
		if crd.Spec.Conversion.Webhook.ClientConfig.Service.Path != nil {
			path = *crd.Spec.Conversion.Webhook.ClientConfig.Service.Path
		}
		url := fmt.Sprintf("%s%s", url, path)
		crd.Spec.Conversion.Webhook.ClientConfig.URL = &url
		crd.Spec.Conversion.Webhook.ClientConfig.CABundle = caPEM
		crd.Spec.Conversion.Webhook.ClientConfig.Service = nil
	}
}

func readCustomResourcesAtPath(t *testing.T, path string) []*v1.CustomResourceDefinition {
	serializer := jsonserializer.NewSerializerWithOptions(jsonserializer.DefaultMetaFactory, internalScheme, internalScheme, jsonserializer.SerializerOptions{
		Yaml: true,
	})
	converter := runtime.UnsafeObjectConvertor(internalScheme)
	codec := versioning.NewCodec(serializer, serializer, converter, internalScheme, internalScheme, internalScheme, runtime.InternalGroupVersioner, runtime.InternalGroupVersioner, internalScheme.Name())

	var crds []*v1.CustomResourceDefinition
	if err := filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if filepath.Ext(path) != ".yaml" {
			return nil
		}
		crd, err := readCRDsAtPath(codec, converter, path)
		if err != nil {
			return err
		}
		crds = append(crds, crd...)
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	return crds
}

func readCRDsAtPath(codec runtime.Codec, converter runtime.ObjectConvertor, path string) ([]*v1.CustomResourceDefinition, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	var crds []*v1.CustomResourceDefinition
	for _, d := range strings.Split(string(data), "\n---\n") {
		// skip empty YAML documents
		if strings.TrimSpace(d) == "" {
			continue
		}

		internalCRD := &apiextensions.CustomResourceDefinition{}
		if _, _, err := codec.Decode([]byte(d), nil, internalCRD); err != nil {
			return nil, err
		}

		out := &v1.CustomResourceDefinition{}
		if err := converter.Convert(internalCRD, out, nil); err != nil {
			return nil, err
		}

		crds = append(crds, out)
	}

	return crds, nil
}

func crdsToRuntimeObjects(in []*v1.CustomResourceDefinition) []runtime.Object {
	out := make([]runtime.Object, len(in))

	for i, crd := range in {
		out[i] = runtime.Object(crd)
	}

	return out
}

func getValidatingWebhookConfig(url string, caPEM []byte) runtime.Object {
	failurePolicy := admissionregistrationv1beta1.Fail
	sideEffects := admissionregistrationv1beta1.SideEffectClassNone
	validateURL := fmt.Sprintf("%s/validate", url)
	webhook := admissionregistrationv1beta1.ValidatingWebhookConfiguration{
		ObjectMeta: metav1.ObjectMeta{
			Name: "cert-manager-webhook",
		},
		Webhooks: []admissionregistrationv1beta1.ValidatingWebhook{
			{
				Name: "webhook.cert-manager.io",
				ClientConfig: admissionregistrationv1beta1.WebhookClientConfig{
					URL:      &validateURL,
					CABundle: caPEM,
				},
				Rules: []admissionregistrationv1beta1.RuleWithOperations{
					{
						Operations: []admissionregistrationv1beta1.OperationType{
							admissionregistrationv1beta1.Create,
							admissionregistrationv1beta1.Update,
						},
						Rule: admissionregistrationv1beta1.Rule{
							APIGroups:   []string{"cert-manager.io", "acme.cert-manager.io"},
							APIVersions: []string{"*"},
							Resources:   []string{"*/*"},
						},
					},
				},
				FailurePolicy: &failurePolicy,
				SideEffects:   &sideEffects,
			},
		},
	}

	return &webhook
}

func getMutatingWebhookConfig(url string, caPEM []byte) runtime.Object {
	failurePolicy := admissionregistrationv1beta1.Fail
	sideEffects := admissionregistrationv1beta1.SideEffectClassNone
	validateURL := fmt.Sprintf("%s/mutate", url)
	webhook := admissionregistrationv1beta1.MutatingWebhookConfiguration{
		ObjectMeta: metav1.ObjectMeta{
			Name: "cert-manager-webhook",
		},
		Webhooks: []admissionregistrationv1beta1.MutatingWebhook{
			{
				Name: "webhook.cert-manager.io",
				ClientConfig: admissionregistrationv1beta1.WebhookClientConfig{
					URL:      &validateURL,
					CABundle: caPEM,
				},
				Rules: []admissionregistrationv1beta1.RuleWithOperations{
					{
						Operations: []admissionregistrationv1beta1.OperationType{
							admissionregistrationv1beta1.Create,
							admissionregistrationv1beta1.Update,
						},
						Rule: admissionregistrationv1beta1.Rule{
							APIGroups:   []string{"cert-manager.io", "acme.cert-manager.io"},
							APIVersions: []string{"*"},
							Resources:   []string{"*/*"},
						},
					},
				},
				FailurePolicy: &failurePolicy,
				SideEffects:   &sideEffects,
			},
		},
	}

	return &webhook
}
<|endoftext|>


<commit_before>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package context

import (
	"bytes"
	"io"
	"os"
	"path/filepath"

	"github.com/juju/errors"
	"github.com/juju/juju/resource"
	"github.com/juju/utils"
	"github.com/juju/utils/hash"
	charmresource "gopkg.in/juju/charm.v6-unstable/resource"
)

// HookContextFacade is the name of the API facade for resources in the uniter.
const HookContextFacade = resource.ComponentName + "-hook-context"

// APIClient exposes the uniter API functionality needed for resources.
type APIClient interface {
	// GetResource returns the resource info and content for the given
	// name (and unit-implied service).
	GetResource(resourceName string) (resource.Resource, io.ReadCloser, error)
}

// HookContext exposes the functionality exposed by the resources context.
type HookContext interface {
	// DownloadResource downloads the named resource and returns
	// the path to which it was downloaded.
	DownloadResource(name string) (filePath string, _ error)
}

// NewContextAPI returns a new Content for the given API client and data dir.
func NewContextAPI(apiClient APIClient, dataDir string) *Context {
	return &Context{
		apiClient: apiClient,
		dataDir:   dataDir,

		createResourceFile: createResourceFile,
		writeResource:      writeResource,
	}
}

// Content is the resources portion of a uniter hook context.
type Context struct {
	apiClient APIClient
	dataDir   string

	createResourceFile func(path string) (io.WriteCloser, error)
	writeResource      func(io.Writer, io.Reader) (int64, charmresource.Fingerprint, error)
}

// DownloadResource downloads the named resource and returns the path
// to which it was downloaded. If the resource does not exist or has
// not been uploaded yet then errors.NotFound is returned.
//
// Note that the downloaded file is checked for correctness.
func (c *Context) DownloadResource(name string) (string, error) {
	// TODO(katco): Potential race-condition: two commands running at
	// once.
	// TODO(katco): Check to see if we have latest version

	resourcePath, _, err := c.downloadResource(name)
	if err != nil {
		return "", errors.Trace(err)
	}

	return resourcePath, nil
}

// Flush implements jujuc.Context.
func (c *Context) Flush() error {
	return nil
}

// downloadResource downloads the named resource to the provided path.
func (c *Context) downloadResource(name string) (string, resource.Resource, error) {
	// TODO(ericsnow) This needs to be atomic?
	// (e.g. write to separate dir and move the dir into place)

	info, resourceReader, err := c.apiClient.GetResource(name)
	if err != nil {
		return "", resource.Resource{}, errors.Trace(err)
	}
	defer resourceReader.Close()

	resourcePath := resolveResourcePath(c.dataDir, info)

	target, err := c.createResourceFile(resourcePath)
	if err != nil {
		return "", resource.Resource{}, errors.Trace(err)
	}
	defer target.Close()

	content := resourceContent{
		data:        resourceReader,
		size:        info.Size,
		fingerprint: info.Fingerprint,
	}
	size, fp, err := c.writeResource(target, content.data)
	if err != nil {
		return "", resource.Resource{}, errors.Trace(err)
	}

	if err := content.verify(size, fp); err != nil {
		return "", resource.Resource{}, errors.Trace(err)
	}

	return resourcePath, info, nil
}

// resourceContent holds a reader for the content of a resource along
// with details about that content.
type resourceContent struct {
	data        io.Reader
	size        int64
	fingerprint charmresource.Fingerprint
}

// verify ensures that the actual resource content details match
// the expected ones.
func (c resourceContent) verify(size int64, fp charmresource.Fingerprint) error {
	if size != c.size {
		return errors.Errorf("resource size does not match expected (%d != %d)", size, c.size)
	}
	if !bytes.Equal(fp.Bytes(), c.fingerprint.Bytes()) {
		return errors.Errorf("resource fingerprint does not match expected (%q != %q)", fp, c.fingerprint)
	}
	return nil
}

// resolveResourcePath returns the full path to the resource.
func resolveResourcePath(unitPath string, resourceInfo resource.Resource) string {
	return filepath.Join(unitPath, resourceInfo.Name, resourceInfo.Path)
}

// createResourceFile creates the file into which a resource's content
// should be written.
func createResourceFile(path string) (io.WriteCloser, error) {
	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return nil, errors.Annotate(err, "could not create resource dir")
	}
	target, err := os.Create(path)
	if err != nil {
		return nil, errors.Annotate(err, "could not create new file for resource")
	}
	// TODO(ericsnow) chmod 0644?
	return target, nil
}

// writeResource writes the provided (source) resource content
// to the target. The size and fingerprint are returned.
func writeResource(target io.Writer, source io.Reader) (size int64, fp charmresource.Fingerprint, err error) {
	checksumWriter := charmresource.NewFingerprintHash()
	hashingReader := hash.NewHashingReader(source, checksumWriter)
	sizingReader := utils.NewSizingReader(hashingReader)
	source = sizingReader

	if _, err := io.Copy(target, source); err != nil {
		return size, fp, errors.Annotate(err, "could not write resource to file")
	}

	size = sizingReader.Size()
	fp = checksumWriter.Fingerprint()
	return size, fp, nil
}
<commit_msg>Download the resource into a temp dir.<commit_after>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package context

import (
	"bytes"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/juju/errors"
	"github.com/juju/juju/resource"
	"github.com/juju/utils"
	"github.com/juju/utils/hash"
	charmresource "gopkg.in/juju/charm.v6-unstable/resource"
)

// HookContextFacade is the name of the API facade for resources in the uniter.
const HookContextFacade = resource.ComponentName + "-hook-context"

// APIClient exposes the uniter API functionality needed for resources.
type APIClient interface {
	// GetResource returns the resource info and content for the given
	// name (and unit-implied service).
	GetResource(resourceName string) (resource.Resource, io.ReadCloser, error)
}

// HookContext exposes the functionality exposed by the resources context.
type HookContext interface {
	// DownloadResource downloads the named resource and returns
	// the path to which it was downloaded.
	DownloadResource(name string) (filePath string, _ error)
}

// NewContextAPI returns a new Content for the given API client and data dir.
func NewContextAPI(apiClient APIClient, dataDir string) *Context {
	return &Context{
		apiClient: apiClient,
		dataDir:   dataDir,

		tempDir:            func() (string, error) { return ioutil.TempDir("", "juju-resource-") },
		removeDir:          os.RemoveAll,
		createResourceFile: createResourceFile,
		writeResource:      writeResource,
		replaceDirectory:   replaceDirectory,
	}
}

// Content is the resources portion of a uniter hook context.
type Context struct {
	apiClient APIClient
	dataDir   string

	tempDir            func() (string, error)
	removeDir          func(string) error
	createResourceFile func(path string) (io.WriteCloser, error)
	writeResource      func(io.Writer, io.Reader) (int64, charmresource.Fingerprint, error)
	replaceDirectory   func(string, string) error
}

// DownloadResource downloads the named resource and returns the path
// to which it was downloaded. If the resource does not exist or has
// not been uploaded yet then errors.NotFound is returned.
//
// Note that the downloaded file is checked for correctness.
func (c *Context) DownloadResource(name string) (string, error) {
	// TODO(katco): Potential race-condition: two commands running at
	// once.
	// TODO(katco): Check to see if we have latest version

	tempDir, err := c.tempDir()
	if err != nil {
		return "", errors.Trace(err)
	}
	defer c.removeDir(tempDir)

	info, err := c.downloadResource(name, tempDir)
	if err != nil {
		return "", errors.Trace(err)
	}

	resourcePath, err := c.replaceResourceDir(tempDir, info)
	if err != nil {
		return "", errors.Trace(err)
	}

	return resourcePath, nil
}

// Flush implements jujuc.Context.
func (c *Context) Flush() error {
	return nil
}

// downloadResource downloads the named resource to the provided path.
func (c *Context) downloadResource(name, tempDir string) (resource.Resource, error) {
	info, resourceReader, err := c.apiClient.GetResource(name)
	if err != nil {
		return resource.Resource{}, errors.Trace(err)
	}
	defer resourceReader.Close()

	resourcePath := resolveResourcePath(tempDir, info)

	target, err := c.createResourceFile(resourcePath)
	if err != nil {
		return resource.Resource{}, errors.Trace(err)
	}
	defer target.Close()

	content := resourceContent{
		data:        resourceReader,
		size:        info.Size,
		fingerprint: info.Fingerprint,
	}
	size, fp, err := c.writeResource(target, content.data)
	if err != nil {
		return resource.Resource{}, errors.Trace(err)
	}

	if err := content.verify(size, fp); err != nil {
		return resource.Resource{}, errors.Trace(err)
	}

	return info, nil
}

// replaceResourceDir replaces the original resource dir
// with the temporary resource dir.
func (c *Context) replaceResourceDir(tempDir string, info resource.Resource) (string, error) {
	oldDir := filepath.Dir(resolveResourcePath(tempDir, info))
	resourcePath := resolveResourcePath(c.dataDir, info)
	resDir := filepath.Dir(resourcePath)

	if err := c.replaceDirectory(resDir, oldDir); err != nil {
		return "", errors.Annotate(err, "could not replace existing resource directory")
	}

	return resourcePath, nil
}

// resourceContent holds a reader for the content of a resource along
// with details about that content.
type resourceContent struct {
	data        io.Reader
	size        int64
	fingerprint charmresource.Fingerprint
}

// verify ensures that the actual resource content details match
// the expected ones.
func (c resourceContent) verify(size int64, fp charmresource.Fingerprint) error {
	if size != c.size {
		return errors.Errorf("resource size does not match expected (%d != %d)", size, c.size)
	}
	if !bytes.Equal(fp.Bytes(), c.fingerprint.Bytes()) {
		return errors.Errorf("resource fingerprint does not match expected (%q != %q)", fp, c.fingerprint)
	}
	return nil
}

// resolveResourcePath returns the full path to the resource.
func resolveResourcePath(dataDir string, resourceInfo resource.Resource) string {
	return filepath.Join(dataDir, resourceInfo.Name, resourceInfo.Path)
}

// createResourceFile creates the file into which a resource's content
// should be written.
func createResourceFile(path string) (io.WriteCloser, error) {
	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return nil, errors.Annotate(err, "could not create resource dir")
	}
	target, err := os.Create(path)
	if err != nil {
		return nil, errors.Annotate(err, "could not create new file for resource")
	}
	// TODO(ericsnow) chmod 0644?
	return target, nil
}

// writeResource writes the provided (source) resource content
// to the target. The size and fingerprint are returned.
func writeResource(target io.Writer, source io.Reader) (size int64, fp charmresource.Fingerprint, err error) {
	checksumWriter := charmresource.NewFingerprintHash()
	hashingReader := hash.NewHashingReader(source, checksumWriter)
	sizingReader := utils.NewSizingReader(hashingReader)
	source = sizingReader

	if _, err := io.Copy(target, source); err != nil {
		return size, fp, errors.Annotate(err, "could not write resource to file")
	}

	size = sizingReader.Size()
	fp = checksumWriter.Fingerprint()
	return size, fp, nil
}

// replaceDirectory replaces the target directory with the source. This
// involves removing the target if it exists and then moving the source
// into place.
func replaceDirectory(targetDir, sourceDir string) error {
	if err := os.RemoveAll(targetDir); err != nil {
		return errors.Trace(err)
	}
	if err := os.Rename(targetDir, sourceDir); err != nil {
		return errors.Trace(err)
	}
	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"bytes"
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/goamz/ec2"
)

func resourceAwsRouteTable() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsRouteTableCreate,
		Read:   resourceAwsRouteTableRead,
		Update: resourceAwsRouteTableUpdate,
		Delete: resourceAwsRouteTableDelete,

		Schema: map[string]*schema.Schema{
			"vpc_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"tags": tagsSchema(),

			"route": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"cidr_block": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"gateway_id": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},

						"instance_id": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
				Set: resourceAwsRouteTableHash,
			},
		},
	}
}

func resourceAwsRouteTableCreate(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	// Create the routing table
	createOpts := &ec2.CreateRouteTable{
		VpcId: d.Get("vpc_id").(string),
	}
	log.Printf("[DEBUG] RouteTable create config: %#v", createOpts)

	resp, err := ec2conn.CreateRouteTable(createOpts)
	if err != nil {
		return fmt.Errorf("Error creating route table: %s", err)
	}

	// Get the ID and store it
	rt := &resp.RouteTable
	d.SetId(rt.RouteTableId)
	log.Printf("[INFO] Route Table ID: %s", d.Id())

	// Wait for the route table to become available
	log.Printf(
		"[DEBUG] Waiting for route table (%s) to become available",
		d.Id())
	stateConf := &resource.StateChangeConf{
		Pending: []string{"pending"},
		Target:  "ready",
		Refresh: resourceAwsRouteTableStateRefreshFunc(ec2conn, d.Id()),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for route table (%s) to become available: %s",
			d.Id(), err)
	}

	return resourceAwsRouteTableUpdate(d, meta)
}

func resourceAwsRouteTableRead(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	rtRaw, _, err := resourceAwsRouteTableStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if rtRaw == nil {
		return nil
	}

	rt := rtRaw.(*ec2.RouteTable)
	d.Set("vpc_id", rt.VpcId)

	// Create an empty schema.Set to hold all routes
	route := &schema.Set{F: resourceAwsRouteTableHash}

	// Loop through the routes and add them to the set
	for _, r := range rt.Routes {
		if r.GatewayId == "local" {
			continue
		}

		if r.Origin == "EnableVgwRoutePropagation" {
			continue
		}

		m := make(map[string]interface{})
		m["cidr_block"] = r.DestinationCidrBlock

		if r.GatewayId != "" {
			m["gateway_id"] = r.GatewayId
		}

		if r.InstanceId != "" {
			m["instance_id"] = r.InstanceId
		}

		route.Add(m)
	}
	d.Set("route", route)

	// Tags
	d.Set("tags", tagsToMap(rt.Tags))

	return nil
}

func resourceAwsRouteTableUpdate(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	// Check if the route set as a whole has changed
	if d.HasChange("route") {
		o, n := d.GetChange("route")
		ors := o.(*schema.Set).Difference(n.(*schema.Set))
		nrs := n.(*schema.Set).Difference(o.(*schema.Set))

		// Now first loop through all the old routes and delete any obsolete ones
		for _, route := range ors.List() {
			m := route.(map[string]interface{})

			// Delete the route as it no longer exists in the config
			_, err := ec2conn.DeleteRoute(
				d.Id(), m["cidr_block"].(string))
			if err != nil {
				return err
			}
		}

		// Make sure we save the state of the currently configured rules
		routes := o.(*schema.Set).Intersection(n.(*schema.Set))
		d.Set("route", routes)

		// Then loop through al the newly configured routes and create them
		for _, route := range nrs.List() {
			m := route.(map[string]interface{})

			opts := ec2.CreateRoute{
				RouteTableId:         d.Id(),
				DestinationCidrBlock: m["cidr_block"].(string),
				GatewayId:            m["gateway_id"].(string),
				InstanceId:           m["instance_id"].(string),
			}

			_, err := ec2conn.CreateRoute(&opts)
			if err != nil {
				return err
			}

			routes.Add(route)
			d.Set("route", routes)
		}
	}

	if err := setTags(ec2conn, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	return resourceAwsRouteTableRead(d, meta)
}

func resourceAwsRouteTableDelete(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	// First request the routing table since we'll have to disassociate
	// all the subnets first.
	rtRaw, _, err := resourceAwsRouteTableStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if rtRaw == nil {
		return nil
	}
	rt := rtRaw.(*ec2.RouteTable)

	// Do all the disassociations
	for _, a := range rt.Associations {
		log.Printf("[INFO] Disassociating association: %s", a.AssociationId)
		if _, err := ec2conn.DisassociateRouteTable(a.AssociationId); err != nil {
			return err
		}
	}

	// Delete the route table
	log.Printf("[INFO] Deleting Route Table: %s", d.Id())
	if _, err := ec2conn.DeleteRouteTable(d.Id()); err != nil {
		ec2err, ok := err.(*ec2.Error)
		if ok && ec2err.Code == "InvalidRouteTableID.NotFound" {
			return nil
		}

		return fmt.Errorf("Error deleting route table: %s", err)
	}

	// Wait for the route table to really destroy
	log.Printf(
		"[DEBUG] Waiting for route table (%s) to become destroyed",
		d.Id())

	stateConf := &resource.StateChangeConf{
		Pending: []string{"ready"},
		Target:  "",
		Refresh: resourceAwsRouteTableStateRefreshFunc(ec2conn, d.Id()),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for route table (%s) to become destroyed: %s",
			d.Id(), err)
	}

	return nil
}

func resourceAwsRouteTableHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%s-", m["cidr_block"].(string)))

	if v, ok := m["gateway_id"]; ok {
		buf.WriteString(fmt.Sprintf("%s-", v.(string)))
	}

	if v, ok := m["instance_id"]; ok {
		buf.WriteString(fmt.Sprintf("%s-", v.(string)))
	}

	return hashcode.String(buf.String())
}

// resourceAwsRouteTableStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// a RouteTable.
func resourceAwsRouteTableStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.DescribeRouteTables([]string{id}, ec2.NewFilter())
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok && ec2err.Code == "InvalidRouteTableID.NotFound" {
				resp = nil
			} else {
				log.Printf("Error on RouteTableStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			// Sometimes AWS just has consistency issues and doesn't see
			// our instance yet. Return an empty state.
			return nil, "", nil
		}

		rt := &resp.RouteTables[0]
		return rt, "ready", nil
	}
}
<commit_msg>This change belongs in a branch<commit_after>package aws

import (
	"bytes"
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/goamz/ec2"
)

func resourceAwsRouteTable() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsRouteTableCreate,
		Read:   resourceAwsRouteTableRead,
		Update: resourceAwsRouteTableUpdate,
		Delete: resourceAwsRouteTableDelete,

		Schema: map[string]*schema.Schema{
			"vpc_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"tags": tagsSchema(),

			"route": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"cidr_block": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"gateway_id": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},

						"instance_id": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
				Set: resourceAwsRouteTableHash,
			},
		},
	}
}

func resourceAwsRouteTableCreate(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	// Create the routing table
	createOpts := &ec2.CreateRouteTable{
		VpcId: d.Get("vpc_id").(string),
	}
	log.Printf("[DEBUG] RouteTable create config: %#v", createOpts)

	resp, err := ec2conn.CreateRouteTable(createOpts)
	if err != nil {
		return fmt.Errorf("Error creating route table: %s", err)
	}

	// Get the ID and store it
	rt := &resp.RouteTable
	d.SetId(rt.RouteTableId)
	log.Printf("[INFO] Route Table ID: %s", d.Id())

	// Wait for the route table to become available
	log.Printf(
		"[DEBUG] Waiting for route table (%s) to become available",
		d.Id())
	stateConf := &resource.StateChangeConf{
		Pending: []string{"pending"},
		Target:  "ready",
		Refresh: resourceAwsRouteTableStateRefreshFunc(ec2conn, d.Id()),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for route table (%s) to become available: %s",
			d.Id(), err)
	}

	return resourceAwsRouteTableUpdate(d, meta)
}

func resourceAwsRouteTableRead(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	rtRaw, _, err := resourceAwsRouteTableStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if rtRaw == nil {
		return nil
	}

	rt := rtRaw.(*ec2.RouteTable)
	d.Set("vpc_id", rt.VpcId)

	// Create an empty schema.Set to hold all routes
	route := &schema.Set{F: resourceAwsRouteTableHash}

	// Loop through the routes and add them to the set
	for _, r := range rt.Routes {
		if r.GatewayId == "local" {
			continue
		}

		m := make(map[string]interface{})
		m["cidr_block"] = r.DestinationCidrBlock

		if r.GatewayId != "" {
			m["gateway_id"] = r.GatewayId
		}

		if r.InstanceId != "" {
			m["instance_id"] = r.InstanceId
		}

		route.Add(m)
	}
	d.Set("route", route)

	// Tags
	d.Set("tags", tagsToMap(rt.Tags))

	return nil
}

func resourceAwsRouteTableUpdate(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	// Check if the route set as a whole has changed
	if d.HasChange("route") {
		o, n := d.GetChange("route")
		ors := o.(*schema.Set).Difference(n.(*schema.Set))
		nrs := n.(*schema.Set).Difference(o.(*schema.Set))

		// Now first loop through all the old routes and delete any obsolete ones
		for _, route := range ors.List() {
			m := route.(map[string]interface{})

			// Delete the route as it no longer exists in the config
			_, err := ec2conn.DeleteRoute(
				d.Id(), m["cidr_block"].(string))
			if err != nil {
				return err
			}
		}

		// Make sure we save the state of the currently configured rules
		routes := o.(*schema.Set).Intersection(n.(*schema.Set))
		d.Set("route", routes)

		// Then loop through al the newly configured routes and create them
		for _, route := range nrs.List() {
			m := route.(map[string]interface{})

			opts := ec2.CreateRoute{
				RouteTableId:         d.Id(),
				DestinationCidrBlock: m["cidr_block"].(string),
				GatewayId:            m["gateway_id"].(string),
				InstanceId:           m["instance_id"].(string),
			}

			_, err := ec2conn.CreateRoute(&opts)
			if err != nil {
				return err
			}

			routes.Add(route)
			d.Set("route", routes)
		}
	}

	if err := setTags(ec2conn, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	return resourceAwsRouteTableRead(d, meta)
}

func resourceAwsRouteTableDelete(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	// First request the routing table since we'll have to disassociate
	// all the subnets first.
	rtRaw, _, err := resourceAwsRouteTableStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if rtRaw == nil {
		return nil
	}
	rt := rtRaw.(*ec2.RouteTable)

	// Do all the disassociations
	for _, a := range rt.Associations {
		log.Printf("[INFO] Disassociating association: %s", a.AssociationId)
		if _, err := ec2conn.DisassociateRouteTable(a.AssociationId); err != nil {
			return err
		}
	}

	// Delete the route table
	log.Printf("[INFO] Deleting Route Table: %s", d.Id())
	if _, err := ec2conn.DeleteRouteTable(d.Id()); err != nil {
		ec2err, ok := err.(*ec2.Error)
		if ok && ec2err.Code == "InvalidRouteTableID.NotFound" {
			return nil
		}

		return fmt.Errorf("Error deleting route table: %s", err)
	}

	// Wait for the route table to really destroy
	log.Printf(
		"[DEBUG] Waiting for route table (%s) to become destroyed",
		d.Id())

	stateConf := &resource.StateChangeConf{
		Pending: []string{"ready"},
		Target:  "",
		Refresh: resourceAwsRouteTableStateRefreshFunc(ec2conn, d.Id()),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for route table (%s) to become destroyed: %s",
			d.Id(), err)
	}

	return nil
}

func resourceAwsRouteTableHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%s-", m["cidr_block"].(string)))

	if v, ok := m["gateway_id"]; ok {
		buf.WriteString(fmt.Sprintf("%s-", v.(string)))
	}

	if v, ok := m["instance_id"]; ok {
		buf.WriteString(fmt.Sprintf("%s-", v.(string)))
	}

	return hashcode.String(buf.String())
}

// resourceAwsRouteTableStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// a RouteTable.
func resourceAwsRouteTableStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.DescribeRouteTables([]string{id}, ec2.NewFilter())
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok && ec2err.Code == "InvalidRouteTableID.NotFound" {
				resp = nil
			} else {
				log.Printf("Error on RouteTableStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			// Sometimes AWS just has consistency issues and doesn't see
			// our instance yet. Return an empty state.
			return nil, "", nil
		}

		rt := &resp.RouteTables[0]
		return rt, "ready", nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/vmware/govmomi"
	"github.com/vmware/govmomi/vim25/mo"
	"github.com/vmware/govmomi/vim25/types"
)

func resourceVirtualMachine() *schema.Resource {
	return &schema.Resource{
		Create: resourceVirtualMachineCreate,
		Read:   resourceVirtualMachineRead,
		Delete: resourceVirtualMachineDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"source": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"datacenter": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"folder": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"host": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
                Computed: true,
			},
			"pool": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
                Computed: true,
            },
			"linked_clone": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
			},
			"cpus": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
                Computed: true,
			},
			"memory": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
                Computed: true,
			},
			"power_on": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
                Default:  true,
			},
			"ip_address": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceVirtualMachineCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*govmomi.Client)

	source_vm_ref, err := client.SearchIndex().FindByInventoryPath(fmt.Sprintf("%s/vm/%s", d.Get("datacenter").(string), d.Get("source").(string)))
	if err != nil {
		return fmt.Errorf("Error reading vm: %s", err)
	}
	source_vm := source_vm_ref.(*govmomi.VirtualMachine)

	folder_ref, err := client.SearchIndex().FindByInventoryPath(fmt.Sprintf("%v/vm/%v", d.Get("datacenter").(string), d.Get("folder").(string)))
	if err != nil {
		return fmt.Errorf("Error reading folder: %s", err)
	}
	folder := folder_ref.(*govmomi.Folder)

	var o mo.VirtualMachine
	err = client.Properties(source_vm.Reference(), []string{"snapshot"}, &o)
	if err != nil {
		return fmt.Errorf("Error reading snapshot")
	}
	if o.Snapshot == nil {
		return fmt.Errorf("Base VM has no snapshots")
	}
	snapshot := o.Snapshot.CurrentSnapshot

	var relocateSpec types.VirtualMachineRelocateSpec

    var pool_mor types.ManagedObjectReference
    if d.Get("host").(string) != "" && d.Get("pool").(string) != ""	{
        pool_ref, err := client.SearchIndex().FindByInventoryPath(fmt.Sprintf("%v/host/%v/Resources/%v", d.Get("datacenter").(string), d.Get("host").(string), d.Get("pool").(string)))
        if err != nil {
            return fmt.Errorf("Error reading resource pool: %s", err)
        }
        pool_mor = pool_ref.Reference()
        relocateSpec.Pool = &pool_mor
    }

	if d.Get("linked_clone").(bool) {
		relocateSpec.DiskMoveType = "createNewChildDiskBacking"
	}
    var confSpec types.VirtualMachineConfigSpec
    if d.Get("cpus") != nil {
        confSpec.NumCPUs = d.Get("cpus").(int)
    }
    if d.Get("memory") != nil {
        confSpec.MemoryMB = int64(d.Get("memory").(int))
    }

	cloneSpec := types.VirtualMachineCloneSpec{
		Location: relocateSpec,
        Config:   &confSpec,
        PowerOn:  d.Get("power_on").(bool),
	}
    if d.Get("linked_clone").(bool) {
        cloneSpec.Snapshot = snapshot
    }

	task, err := source_vm.Clone(folder, d.Get("name").(string), cloneSpec)
	if err != nil {
		return fmt.Errorf("Error clonning vm: %s", err)
	}
	info, err := task.WaitForResult(nil)
	if err != nil {
		return fmt.Errorf("Error clonning vm: %s", err)
	}

	vm_mor := info.Result.(types.ManagedObjectReference)
    d.SetId(vm_mor.Value)
    vm := govmomi.NewVirtualMachine(client, vm_mor)
    // workaround for https://github.com/vmware/govmomi/issues/218
    if d.Get("power_on").(bool) {
        ip, err := vm.WaitForIP()
        if err != nil {
            log.Printf("[ERROR] Cannot read ip address: %s", err)
        } else {
            d.Set("ip_address", ip)
        }
    }

    return nil
}

func resourceVirtualMachineRead(d *schema.ResourceData, meta interface{}) error {
    client := meta.(*govmomi.Client)
    vm_mor := types.ManagedObjectReference{Type: "VirtualMachine", Value: d.Id() }
    vm := govmomi.NewVirtualMachine(client, vm_mor)

    var o mo.VirtualMachine
    err := client.Properties(vm.Reference(), []string{"summary"}, &o)
    if err != nil {
        d.SetId("")
        return nil
    }
    d.Set("name", o.Summary.Config.Name)
    d.Set("cpus", o.Summary.Config.NumCpu)
    d.Set("memory", o.Summary.Config.MemorySizeMB)

    if o.Summary.Runtime.PowerState == "poweredOn" {
        d.Set("power_on", true)
    } else {
        d.Set("power_on", false)
    }

    if d.Get("power_on").(bool) {
        ip, err := vm.WaitForIP()
        if err != nil {
            log.Printf("[ERROR] Cannot read ip address: %s", err)
        } else {
            d.Set("ip_address", ip)
        }
    }

	return nil
}

func resourceVirtualMachineDelete(d *schema.ResourceData, meta interface{}) error {
    client := meta.(*govmomi.Client)
    vm_mor := types.ManagedObjectReference{Type: "VirtualMachine", Value: d.Id() }
    vm := govmomi.NewVirtualMachine(client, vm_mor)

    task, err := vm.PowerOff()
    if err != nil {
        return fmt.Errorf("Error powering vm off: %s", err)
    }
    task.WaitForResult(nil)

    task, err = vm.Destroy()
    if err != nil {
        return fmt.Errorf("Error deleting vm: %s", err)
    }
    _, err = task.WaitForResult(nil)
    if err != nil {
        return fmt.Errorf("Error deleting vm: %s", err)
    }

    return nil
}
<commit_msg>bugfix: search for current snapshot only when it is required<commit_after>package main

import (
	"fmt"
	"log"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/vmware/govmomi"
	"github.com/vmware/govmomi/vim25/mo"
	"github.com/vmware/govmomi/vim25/types"
)

func resourceVirtualMachine() *schema.Resource {
	return &schema.Resource{
		Create: resourceVirtualMachineCreate,
		Read:   resourceVirtualMachineRead,
		Delete: resourceVirtualMachineDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"source": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"datacenter": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"folder": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"host": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
                Computed: true,
			},
			"pool": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
                Computed: true,
            },
			"linked_clone": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
			},
			"cpus": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
                Computed: true,
			},
			"memory": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
                Computed: true,
			},
			"power_on": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
                Default:  true,
			},
			"ip_address": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceVirtualMachineCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*govmomi.Client)

	source_vm_ref, err := client.SearchIndex().FindByInventoryPath(fmt.Sprintf("%s/vm/%s", d.Get("datacenter").(string), d.Get("source").(string)))
	if err != nil {
		return fmt.Errorf("Error reading vm: %s", err)
	}
	source_vm := source_vm_ref.(*govmomi.VirtualMachine)

	folder_ref, err := client.SearchIndex().FindByInventoryPath(fmt.Sprintf("%v/vm/%v", d.Get("datacenter").(string), d.Get("folder").(string)))
	if err != nil {
		return fmt.Errorf("Error reading folder: %s", err)
	}
	folder := folder_ref.(*govmomi.Folder)

	var relocateSpec types.VirtualMachineRelocateSpec

    var pool_mor types.ManagedObjectReference
    if d.Get("host").(string) != "" && d.Get("pool").(string) != ""	{
        pool_ref, err := client.SearchIndex().FindByInventoryPath(fmt.Sprintf("%v/host/%v/Resources/%v", d.Get("datacenter").(string), d.Get("host").(string), d.Get("pool").(string)))
        if err != nil {
            return fmt.Errorf("Error reading resource pool: %s", err)
        }
        pool_mor = pool_ref.Reference()
        relocateSpec.Pool = &pool_mor
    }

	if d.Get("linked_clone").(bool) {
		relocateSpec.DiskMoveType = "createNewChildDiskBacking"
	}
    var confSpec types.VirtualMachineConfigSpec
    if d.Get("cpus") != nil {
        confSpec.NumCPUs = d.Get("cpus").(int)
    }
    if d.Get("memory") != nil {
        confSpec.MemoryMB = int64(d.Get("memory").(int))
    }

	cloneSpec := types.VirtualMachineCloneSpec{
		Location: relocateSpec,
        Config:   &confSpec,
        PowerOn:  d.Get("power_on").(bool),
	}
    if d.Get("linked_clone").(bool) {
        var o mo.VirtualMachine
        err = client.Properties(source_vm.Reference(), []string{"snapshot"}, &o)
        if err != nil {
            return fmt.Errorf("Error reading snapshot")
        }
        if o.Snapshot == nil {
            return fmt.Errorf("Base VM has no snapshots")
        }
        cloneSpec.Snapshot = o.Snapshot.CurrentSnapshot
    }

	task, err := source_vm.Clone(folder, d.Get("name").(string), cloneSpec)
	if err != nil {
		return fmt.Errorf("Error clonning vm: %s", err)
	}
	info, err := task.WaitForResult(nil)
	if err != nil {
		return fmt.Errorf("Error clonning vm: %s", err)
	}

	vm_mor := info.Result.(types.ManagedObjectReference)
    d.SetId(vm_mor.Value)
    vm := govmomi.NewVirtualMachine(client, vm_mor)
    // workaround for https://github.com/vmware/govmomi/issues/218
    if d.Get("power_on").(bool) {
        ip, err := vm.WaitForIP()
        if err != nil {
            log.Printf("[ERROR] Cannot read ip address: %s", err)
        } else {
            d.Set("ip_address", ip)
        }
    }

    return nil
}

func resourceVirtualMachineRead(d *schema.ResourceData, meta interface{}) error {
    client := meta.(*govmomi.Client)
    vm_mor := types.ManagedObjectReference{Type: "VirtualMachine", Value: d.Id() }
    vm := govmomi.NewVirtualMachine(client, vm_mor)

    var o mo.VirtualMachine
    err := client.Properties(vm.Reference(), []string{"summary"}, &o)
    if err != nil {
        d.SetId("")
        return nil
    }
    d.Set("name", o.Summary.Config.Name)
    d.Set("cpus", o.Summary.Config.NumCpu)
    d.Set("memory", o.Summary.Config.MemorySizeMB)

    if o.Summary.Runtime.PowerState == "poweredOn" {
        d.Set("power_on", true)
    } else {
        d.Set("power_on", false)
    }

    if d.Get("power_on").(bool) {
        ip, err := vm.WaitForIP()
        if err != nil {
            log.Printf("[ERROR] Cannot read ip address: %s", err)
        } else {
            d.Set("ip_address", ip)
        }
    }

	return nil
}

func resourceVirtualMachineDelete(d *schema.ResourceData, meta interface{}) error {
    client := meta.(*govmomi.Client)
    vm_mor := types.ManagedObjectReference{Type: "VirtualMachine", Value: d.Id() }
    vm := govmomi.NewVirtualMachine(client, vm_mor)

    task, err := vm.PowerOff()
    if err != nil {
        return fmt.Errorf("Error powering vm off: %s", err)
    }
    task.WaitForResult(nil)

    task, err = vm.Destroy()
    if err != nil {
        return fmt.Errorf("Error deleting vm: %s", err)
    }
    _, err = task.WaitForResult(nil)
    if err != nil {
        return fmt.Errorf("Error deleting vm: %s", err)
    }

    return nil
}
<|endoftext|>


<commit_before>package api

import (
	"context"
	"time"

	"github.com/go-gorp/gorp"
	"github.com/rockbears/log"

	"github.com/ovh/cds/engine/api/workflow"
	"github.com/ovh/cds/sdk"
)

func (api *API) cleanWorkflowRunSecrets(ctx context.Context) {
	// Load workflow run older than now - snapshot retention delay
	maxRetentionDate := time.Now().Add(-time.Hour * time.Duration(24*api.Config.Secrets.SnapshotRetentionDelay))

	db := api.mustDB()

	delay := 10 * time.Minute
	if api.Config.Secrets.SnapshotRetentionDelay > 0 {
		delay = time.Duration(api.Config.Secrets.SnapshotRetentionDelay) * time.Minute
	}

	limit := int64(100)
	if api.Config.Secrets.SnapshotCleanBatchSize > 0 {
		limit = api.Config.Secrets.SnapshotCleanBatchSize
	}

	log.Info(ctx, "Starting workflow run secrets clean routine")

	ticker := time.NewTicker(delay)

	for range ticker.C {
		runIDs, err := workflow.LoadRunsIDsCreatedBefore(ctx, db, maxRetentionDate, limit)
		if err != nil {
			log.ErrorWithStackTrace(ctx, err)
			continue
		}
		for _, id := range runIDs {
			if err := api.cleanWorkflowRunSecretsForRun(ctx, db, id); err != nil {
				log.ErrorWithStackTrace(ctx, err)
			}
		}
	}
}

func (api *API) cleanWorkflowRunSecretsForRun(ctx context.Context, db *gorp.DbMap, workflowRunID int64) error {
	tx, err := db.Begin()
	if err != nil {
		return sdk.WithStack(err)
	}
	defer tx.Rollback() // nolint
	if err := workflow.SetRunReadOnlyByID(ctx, tx, workflowRunID); err != nil {
		return sdk.WithStack(err)
	}
	if err := workflow.DeleteRunSecretsByWorkflowRunID(ctx, tx, workflowRunID); err != nil {
		return sdk.WithStack(err)
	}
	if err := tx.Commit(); err != nil {
		return sdk.WithStack(err)
	}
	return nil
}
<commit_msg>fix(api): workflow run secret clean interval (#6272)<commit_after>package api

import (
	"context"
	"time"

	"github.com/go-gorp/gorp"
	"github.com/rockbears/log"

	"github.com/ovh/cds/engine/api/workflow"
	"github.com/ovh/cds/sdk"
)

func (api *API) cleanWorkflowRunSecrets(ctx context.Context) {
	// Load workflow run older than now - snapshot retention delay
	maxRetentionDate := time.Now().Add(-time.Hour * time.Duration(24*api.Config.Secrets.SnapshotRetentionDelay))

	db := api.mustDB()

	delay := 10 * time.Minute
	if api.Config.Secrets.SnapshotCleanInterval > 0 {
		delay = time.Duration(api.Config.Secrets.SnapshotCleanInterval) * time.Minute
	}

	limit := int64(100)
	if api.Config.Secrets.SnapshotCleanBatchSize > 0 {
		limit = api.Config.Secrets.SnapshotCleanBatchSize
	}

	log.Info(ctx, "Starting workflow run secrets clean routine")

	ticker := time.NewTicker(delay)

	for range ticker.C {
		runIDs, err := workflow.LoadRunsIDsCreatedBefore(ctx, db, maxRetentionDate, limit)
		if err != nil {
			log.ErrorWithStackTrace(ctx, err)
			continue
		}
		for _, id := range runIDs {
			if err := api.cleanWorkflowRunSecretsForRun(ctx, db, id); err != nil {
				log.ErrorWithStackTrace(ctx, err)
			}
		}
	}
}

func (api *API) cleanWorkflowRunSecretsForRun(ctx context.Context, db *gorp.DbMap, workflowRunID int64) error {
	tx, err := db.Begin()
	if err != nil {
		return sdk.WithStack(err)
	}
	defer tx.Rollback() // nolint
	if err := workflow.SetRunReadOnlyByID(ctx, tx, workflowRunID); err != nil {
		return sdk.WithStack(err)
	}
	if err := workflow.DeleteRunSecretsByWorkflowRunID(ctx, tx, workflowRunID); err != nil {
		return sdk.WithStack(err)
	}
	if err := tx.Commit(); err != nil {
		return sdk.WithStack(err)
	}
	return nil
}
<|endoftext|>


<commit_before>//go:build cgo
// +build cgo

package sqliteStorage

import (
	"errors"
	"fmt"
	"runtime"
	"sync"
	"time"

	"crawshaw.io/sqlite"
	"crawshaw.io/sqlite/sqlitex"

	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/storage"
)

type NewDirectStorageOpts struct {
	NewConnOpts
	InitDbOpts
	InitConnOpts
	GcBlobs           bool
	NoCacheBlobs      bool
	BlobFlushInterval time.Duration
}

// A convenience function that creates a connection pool, resource provider, and a pieces storage
// ClientImpl and returns them all with a Close attached.
func NewDirectStorage(opts NewDirectStorageOpts) (_ storage.ClientImplCloser, err error) {
	conn, err := newConn(opts.NewConnOpts)
	if err != nil {
		return
	}
	if opts.PageSize == 0 {
		// The largest size sqlite supports. I think we want this to be the smallest piece size we
		// can expect, which is probably 1<<17.
		opts.PageSize = 1 << 16
	}
	err = initDatabase(conn, opts.InitDbOpts)
	if err != nil {
		conn.Close()
		return
	}
	err = initConn(conn, opts.InitConnOpts)
	if err != nil {
		conn.Close()
		return
	}
	if opts.BlobFlushInterval == 0 && !opts.GcBlobs {
		// This is influenced by typical busy timeouts, of 5-10s. We want to give other connections
		// a few chances at getting a transaction through.
		opts.BlobFlushInterval = time.Second
	}
	cl := &client{
		conn:  conn,
		blobs: make(map[string]*sqlite.Blob),
		opts:  opts,
	}
	// Avoid race with cl.blobFlusherFunc
	cl.l.Lock()
	defer cl.l.Unlock()
	if opts.BlobFlushInterval != 0 {
		cl.blobFlusher = time.AfterFunc(opts.BlobFlushInterval, cl.blobFlusherFunc)
	}
	cl.capacity = cl.getCapacity
	return cl, nil
}

func (cl *client) getCapacity() (ret *int64) {
	cl.l.Lock()
	defer cl.l.Unlock()
	err := sqlitex.Exec(cl.conn, "select value from setting where name='capacity'", func(stmt *sqlite.Stmt) error {
		ret = new(int64)
		*ret = stmt.ColumnInt64(0)
		return nil
	})
	if err != nil {
		panic(err)
	}
	return
}

type client struct {
	l           sync.Mutex
	conn        conn
	blobs       map[string]*sqlite.Blob
	blobFlusher *time.Timer
	opts        NewDirectStorageOpts
	closed      bool
	capacity    func() *int64
}

func (c *client) blobFlusherFunc() {
	c.l.Lock()
	defer c.l.Unlock()
	c.flushBlobs()
	if !c.closed {
		c.blobFlusher.Reset(c.opts.BlobFlushInterval)
	}
}

func (c *client) flushBlobs() {
	for key, b := range c.blobs {
		// Need the lock to prevent racing with the GC finalizers.
		b.Close()
		delete(c.blobs, key)
	}
}

func (c *client) OpenTorrent(info *metainfo.Info, infoHash metainfo.Hash) (storage.TorrentImpl, error) {
	t := torrent{c}
	return storage.TorrentImpl{Piece: t.Piece, Close: t.Close, Capacity: &c.capacity}, nil
}

func (c *client) Close() (err error) {
	c.l.Lock()
	defer c.l.Unlock()
	c.flushBlobs()
	if c.opts.BlobFlushInterval != 0 {
		c.blobFlusher.Stop()
	}
	if !c.closed {
		c.closed = true
		err = c.conn.Close()
		c.conn = nil
	}
	return
}

type torrent struct {
	c *client
}

func rowidForBlob(c conn, name string, length int64, create bool) (rowid int64, err error) {
	rowidOk := false
	err = sqlitex.Exec(c, "select rowid from blob where name=?", func(stmt *sqlite.Stmt) error {
		if rowidOk {
			panic("expected at most one row")
		}
		// TODO: How do we know if we got this wrong?
		rowid = stmt.ColumnInt64(0)
		rowidOk = true
		return nil
	}, name)
	if err != nil {
		return
	}
	if rowidOk {
		return
	}
	if !create {
		err = errors.New("no existing row")
		return
	}
	err = sqlitex.Exec(c, "insert into blob(name, data) values(?, zeroblob(?))", nil, name, length)
	if err != nil {
		return
	}
	rowid = c.LastInsertRowID()
	return
}

func (t torrent) Piece(p metainfo.Piece) storage.PieceImpl {
	t.c.l.Lock()
	defer t.c.l.Unlock()
	name := p.Hash().HexString()
	return piece{
		name,
		p.Length(),
		t.c,
	}
}

func (t torrent) Close() error {
	return nil
}

type piece struct {
	name   string
	length int64
	*client
}

func (p piece) doAtIoWithBlob(
	atIo func(*sqlite.Blob) func([]byte, int64) (int, error),
	b []byte,
	off int64,
	create bool,
) (n int, err error) {
	p.l.Lock()
	defer p.l.Unlock()
	if p.opts.NoCacheBlobs {
		defer p.forgetBlob()
	}
	blob, err := p.getBlob(create)
	if err != nil {
		err = fmt.Errorf("getting blob: %w", err)
		return
	}
	n, err = atIo(blob)(b, off)
	if err == nil {
		return
	}
	var se sqlite.Error
	if !errors.As(err, &se) {
		return
	}
	// "ABORT" occurs if the row the blob is on is modified elsewhere. "ERROR: invalid blob" occurs
	// if the blob has been closed. We don't forget blobs that are closed by our GC finalizers,
	// because they may be attached to names that have since moved on to another blob.
	if se.Code != sqlite.SQLITE_ABORT && !(p.opts.GcBlobs && se.Code == sqlite.SQLITE_ERROR && se.Msg == "invalid blob") {
		return
	}
	p.forgetBlob()
	// Try again, this time we're guaranteed to get a fresh blob, and so errors are no excuse. It
	// might be possible to skip to this version if we don't cache blobs.
	blob, err = p.getBlob(create)
	if err != nil {
		err = fmt.Errorf("getting blob: %w", err)
		return
	}
	return atIo(blob)(b, off)
}

func (p piece) ReadAt(b []byte, off int64) (n int, err error) {
	return p.doAtIoWithBlob(func(blob *sqlite.Blob) func([]byte, int64) (int, error) {
		return blob.ReadAt
	}, b, off, false)
}

func (p piece) WriteAt(b []byte, off int64) (n int, err error) {
	return p.doAtIoWithBlob(func(blob *sqlite.Blob) func([]byte, int64) (int, error) {
		return blob.WriteAt
	}, b, off, true)
}

func (p piece) MarkComplete() error {
	p.l.Lock()
	defer p.l.Unlock()
	err := sqlitex.Exec(p.conn, "update blob set verified=true where name=?", nil, p.name)
	if err != nil {
		return err
	}
	changes := p.conn.Changes()
	if changes != 1 {
		panic(changes)
	}
	return nil
}

func (p piece) forgetBlob() {
	blob, ok := p.blobs[p.name]
	if !ok {
		return
	}
	blob.Close()
	delete(p.blobs, p.name)
}

func (p piece) MarkNotComplete() error {
	p.l.Lock()
	defer p.l.Unlock()
	return sqlitex.Exec(p.conn, "update blob set verified=false where name=?", nil, p.name)
}

func (p piece) Completion() (ret storage.Completion) {
	p.l.Lock()
	defer p.l.Unlock()
	err := sqlitex.Exec(p.conn, "select verified from blob where name=?", func(stmt *sqlite.Stmt) error {
		ret.Complete = stmt.ColumnInt(0) != 0
		return nil
	}, p.name)
	ret.Ok = err == nil
	if err != nil {
		panic(err)
	}
	return
}

func (p piece) getBlob(create bool) (*sqlite.Blob, error) {
	blob, ok := p.blobs[p.name]
	if !ok {
		rowid, err := rowidForBlob(p.conn, p.name, p.length, create)
		if err != nil {
			return nil, fmt.Errorf("getting rowid for blob: %w", err)
		}
		blob, err = p.conn.OpenBlob("main", "blob", "data", rowid, true)
		if err != nil {
			panic(err)
		}
		if p.opts.GcBlobs {
			herp := new(byte)
			runtime.SetFinalizer(herp, func(*byte) {
				p.l.Lock()
				defer p.l.Unlock()
				// Note there's no guarantee that the finalizer fired while this blob is the same
				// one in the blob cache. It might be possible to rework this so that we check, or
				// strip finalizers as appropriate.
				blob.Close()
			})
		}
		p.blobs[p.name] = blob
	}
	return blob, nil
}
<commit_msg>Begin extracting 'squirrel' from storage/sqlite<commit_after>//go:build cgo
// +build cgo

package sqliteStorage

import (
	"errors"
	"fmt"
	"runtime"
	"sync"
	"time"

	"crawshaw.io/sqlite"
	"crawshaw.io/sqlite/sqlitex"

	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/storage"
)

type NewDirectStorageOpts struct {
	NewConnOpts
	InitDbOpts
	InitConnOpts
	GcBlobs           bool
	NoCacheBlobs      bool
	BlobFlushInterval time.Duration
}

func NewSquirrelCache(opts NewDirectStorageOpts) (_ *SquirrelCache, err error) {
	conn, err := newConn(opts.NewConnOpts)
	if err != nil {
		return
	}
	if opts.PageSize == 0 {
		// The largest size sqlite supports. I think we want this to be the smallest SquirrelBlob size we
		// can expect, which is probably 1<<17.
		opts.PageSize = 1 << 16
	}
	err = initDatabase(conn, opts.InitDbOpts)
	if err != nil {
		conn.Close()
		return
	}
	err = initConn(conn, opts.InitConnOpts)
	if err != nil {
		conn.Close()
		return
	}
	if opts.BlobFlushInterval == 0 && !opts.GcBlobs {
		// This is influenced by typical busy timeouts, of 5-10s. We want to give other connections
		// a few chances at getting a transaction through.
		opts.BlobFlushInterval = time.Second
	}
	cl := &SquirrelCache{
		conn:  conn,
		blobs: make(map[string]*sqlite.Blob),
		opts:  opts,
	}
	// Avoid race with cl.blobFlusherFunc
	cl.l.Lock()
	defer cl.l.Unlock()
	if opts.BlobFlushInterval != 0 {
		cl.blobFlusher = time.AfterFunc(opts.BlobFlushInterval, cl.blobFlusherFunc)
	}
	cl.capacity = cl.getCapacity
	return cl, nil
}

// A convenience function that creates a connection pool, resource provider, and a pieces storage
// ClientImpl and returns them all with a Close attached.
func NewDirectStorage(opts NewDirectStorageOpts) (_ storage.ClientImplCloser, err error) {
	cache, err := NewSquirrelCache(opts)
	if err != nil {
		return
	}
	return &client{cache}, nil
}

func (cl *SquirrelCache) getCapacity() (ret *int64) {
	cl.l.Lock()
	defer cl.l.Unlock()
	err := sqlitex.Exec(cl.conn, "select value from setting where name='capacity'", func(stmt *sqlite.Stmt) error {
		ret = new(int64)
		*ret = stmt.ColumnInt64(0)
		return nil
	})
	if err != nil {
		panic(err)
	}
	return
}

type SquirrelCache struct {
	l           sync.Mutex
	conn        conn
	blobs       map[string]*sqlite.Blob
	blobFlusher *time.Timer
	opts        NewDirectStorageOpts
	closed      bool
	capacity    func() *int64
}

type client struct {
	*SquirrelCache
}

func (c *SquirrelCache) blobFlusherFunc() {
	c.l.Lock()
	defer c.l.Unlock()
	c.flushBlobs()
	if !c.closed {
		c.blobFlusher.Reset(c.opts.BlobFlushInterval)
	}
}

func (c *SquirrelCache) flushBlobs() {
	for key, b := range c.blobs {
		// Need the lock to prevent racing with the GC finalizers.
		b.Close()
		delete(c.blobs, key)
	}
}

func (c *SquirrelCache) OpenTorrent(info *metainfo.Info, infoHash metainfo.Hash) (storage.TorrentImpl, error) {
	t := torrent{c}
	return storage.TorrentImpl{Piece: t.Piece, Close: t.Close, Capacity: &c.capacity}, nil
}

func (c *SquirrelCache) Close() (err error) {
	c.l.Lock()
	defer c.l.Unlock()
	c.flushBlobs()
	if c.opts.BlobFlushInterval != 0 {
		c.blobFlusher.Stop()
	}
	if !c.closed {
		c.closed = true
		err = c.conn.Close()
		c.conn = nil
	}
	return
}

type torrent struct {
	c *SquirrelCache
}

func rowidForBlob(c conn, name string, length int64, create bool) (rowid int64, err error) {
	rowidOk := false
	err = sqlitex.Exec(c, "select rowid from blob where name=?", func(stmt *sqlite.Stmt) error {
		if rowidOk {
			panic("expected at most one row")
		}
		// TODO: How do we know if we got this wrong?
		rowid = stmt.ColumnInt64(0)
		rowidOk = true
		return nil
	}, name)
	if err != nil {
		return
	}
	if rowidOk {
		return
	}
	if !create {
		err = errors.New("no existing row")
		return
	}
	err = sqlitex.Exec(c, "insert into blob(name, data) values(?, zeroblob(?))", nil, name, length)
	if err != nil {
		return
	}
	rowid = c.LastInsertRowID()
	return
}

func (t torrent) Piece(p metainfo.Piece) storage.PieceImpl {
	name := p.Hash().HexString()
	return piece{SquirrelBlob{
		name,
		p.Length(),
		t.c,
	}}
}

func (t torrent) Close() error {
	return nil
}

type SquirrelBlob struct {
	name   string
	length int64
	*SquirrelCache
}

type piece struct {
	SquirrelBlob
}

func (p SquirrelBlob) doAtIoWithBlob(
	atIo func(*sqlite.Blob) func([]byte, int64) (int, error),
	b []byte,
	off int64,
	create bool,
) (n int, err error) {
	p.l.Lock()
	defer p.l.Unlock()
	if p.opts.NoCacheBlobs {
		defer p.forgetBlob()
	}
	blob, err := p.getBlob(create)
	if err != nil {
		err = fmt.Errorf("getting blob: %w", err)
		return
	}
	n, err = atIo(blob)(b, off)
	if err == nil {
		return
	}
	var se sqlite.Error
	if !errors.As(err, &se) {
		return
	}
	// "ABORT" occurs if the row the blob is on is modified elsewhere. "ERROR: invalid blob" occurs
	// if the blob has been closed. We don't forget blobs that are closed by our GC finalizers,
	// because they may be attached to names that have since moved on to another blob.
	if se.Code != sqlite.SQLITE_ABORT && !(p.opts.GcBlobs && se.Code == sqlite.SQLITE_ERROR && se.Msg == "invalid blob") {
		return
	}
	p.forgetBlob()
	// Try again, this time we're guaranteed to get a fresh blob, and so errors are no excuse. It
	// might be possible to skip to this version if we don't cache blobs.
	blob, err = p.getBlob(create)
	if err != nil {
		err = fmt.Errorf("getting blob: %w", err)
		return
	}
	return atIo(blob)(b, off)
}

func (p SquirrelBlob) ReadAt(b []byte, off int64) (n int, err error) {
	return p.doAtIoWithBlob(func(blob *sqlite.Blob) func([]byte, int64) (int, error) {
		return blob.ReadAt
	}, b, off, false)
}

func (p SquirrelBlob) WriteAt(b []byte, off int64) (n int, err error) {
	return p.doAtIoWithBlob(func(blob *sqlite.Blob) func([]byte, int64) (int, error) {
		return blob.WriteAt
	}, b, off, true)
}

func (p piece) MarkComplete() error {
	p.l.Lock()
	defer p.l.Unlock()
	err := sqlitex.Exec(p.conn, "update blob set verified=true where name=?", nil, p.name)
	if err != nil {
		return err
	}
	changes := p.conn.Changes()
	if changes != 1 {
		panic(changes)
	}
	return nil
}

func (p SquirrelBlob) forgetBlob() {
	blob, ok := p.blobs[p.name]
	if !ok {
		return
	}
	blob.Close()
	delete(p.blobs, p.name)
}

func (p piece) MarkNotComplete() error {
	p.l.Lock()
	defer p.l.Unlock()
	return sqlitex.Exec(p.conn, "update blob set verified=false where name=?", nil, p.name)
}

func (p piece) Completion() (ret storage.Completion) {
	p.l.Lock()
	defer p.l.Unlock()
	err := sqlitex.Exec(p.conn, "select verified from blob where name=?", func(stmt *sqlite.Stmt) error {
		ret.Complete = stmt.ColumnInt(0) != 0
		return nil
	}, p.name)
	ret.Ok = err == nil
	if err != nil {
		panic(err)
	}
	return
}

func (p SquirrelBlob) getBlob(create bool) (*sqlite.Blob, error) {
	blob, ok := p.blobs[p.name]
	if !ok {
		rowid, err := rowidForBlob(p.conn, p.name, p.length, create)
		if err != nil {
			return nil, fmt.Errorf("getting rowid for blob: %w", err)
		}
		blob, err = p.conn.OpenBlob("main", "blob", "data", rowid, true)
		if err != nil {
			panic(err)
		}
		if p.opts.GcBlobs {
			herp := new(byte)
			runtime.SetFinalizer(herp, func(*byte) {
				p.l.Lock()
				defer p.l.Unlock()
				// Note there's no guarantee that the finalizer fired while this blob is the same
				// one in the blob cache. It might be possible to rework this so that we check, or
				// strip finalizers as appropriate.
				blob.Close()
			})
		}
		p.blobs[p.name] = blob
	}
	return blob, nil
}
<|endoftext|>


<commit_before>package bigtable

import (
	"flag"
	"fmt"
	"log"
	"time"

	"github.com/rakyll/globalconf"
)

type StoreConfig struct {
	Enabled           bool
	GcpProject        string
	BigtableInstance  string
	TableName         string
	WriteQueueSize    int
	WriteMaxFlushSize int
	WriteConcurrency  int
	ReadConcurrency   int
	MaxChunkSpan      time.Duration
	ReadTimeout       time.Duration
	WriteTimeout      time.Duration
	CreateCF          bool
}

func (cfg *StoreConfig) Validate() error {
	// If we dont have any write threads, then we dont WriteMaxFlushSize and WriteQueueSize
	// are not used.  If we do have write threads, then we need to make sure that
	if cfg.WriteConcurrency > 0 {
		if cfg.WriteMaxFlushSize >= cfg.WriteQueueSize {
			return fmt.Errorf("write-queue-size must be larger then write-max-flush-size")
		}
	}
	return nil
}

// return StoreConfig with default values set.
func NewStoreConfig() *StoreConfig {
	return &StoreConfig{
		Enabled:           false,
		GcpProject:        "default",
		BigtableInstance:  "default",
		TableName:         "metrics",
		WriteQueueSize:    100000,
		WriteMaxFlushSize: 10000,
		WriteConcurrency:  10,
		ReadConcurrency:   20,
		MaxChunkSpan:      time.Hour * 6,
		ReadTimeout:       time.Second * 5,
		WriteTimeout:      time.Second * 5,
		CreateCF:          true,
	}
}

var CliConfig = NewStoreConfig()

func ConfigSetup() {
	btStore := flag.NewFlagSet("bigtable-store", flag.ExitOnError)
	btStore.BoolVar(&CliConfig.Enabled, "enabled", CliConfig.Enabled, "enable the bigtable backend store plugin")
	btStore.StringVar(&CliConfig.GcpProject, "gcp-project", CliConfig.GcpProject, "Name of GCP project the bigtable cluster resides in")
	btStore.StringVar(&CliConfig.BigtableInstance, "bigtable-instance", CliConfig.BigtableInstance, "Name of bigtable instance")
	btStore.StringVar(&CliConfig.TableName, "table-name", CliConfig.TableName, "Name of bigtable table used for chunks")
	btStore.IntVar(&CliConfig.WriteQueueSize, "write-queue-size", CliConfig.WriteQueueSize, "Max number of chunks, per write thread, allowed to be unwritten to bigtable. Must be larger then write-max-flush-size")
	btStore.IntVar(&CliConfig.WriteMaxFlushSize, "write-max-flush-size", CliConfig.WriteMaxFlushSize, "Max number of chunks in each batch write to bigtable")
	btStore.IntVar(&CliConfig.WriteConcurrency, "write-concurrency", CliConfig.WriteConcurrency, "Number of writer threads to use.")
	btStore.IntVar(&CliConfig.ReadConcurrency, "read-concurrency", CliConfig.ReadConcurrency, "Number concurrent reads that can be processed")
	btStore.DurationVar(&CliConfig.MaxChunkSpan, "max-chunkspan", CliConfig.MaxChunkSpan, "Maximum chunkspan size used.")
	btStore.DurationVar(&CliConfig.ReadTimeout, "read-timeout", CliConfig.ReadTimeout, "read timeout")
	btStore.DurationVar(&CliConfig.WriteTimeout, "write-timeout", CliConfig.WriteTimeout, "write timeout")
	btStore.BoolVar(&CliConfig.CreateCF, "create-cf", CliConfig.CreateCF, "enable the creation of the table and column families")

	globalconf.Register("bigtable-store", btStore)
	return
}

func ConfigProcess() {
	if err := CliConfig.Validate(); err != nil {
		log.Fatalf("bigtable-store: Config validation error. %s", err)
	}
}
<commit_msg>ensure write-max-flush-size is <= 100000<commit_after>package bigtable

import (
	"flag"
	"fmt"
	"log"
	"time"

	"github.com/rakyll/globalconf"
)

type StoreConfig struct {
	Enabled           bool
	GcpProject        string
	BigtableInstance  string
	TableName         string
	WriteQueueSize    int
	WriteMaxFlushSize int
	WriteConcurrency  int
	ReadConcurrency   int
	MaxChunkSpan      time.Duration
	ReadTimeout       time.Duration
	WriteTimeout      time.Duration
	CreateCF          bool
}

func (cfg *StoreConfig) Validate() error {
	// If we dont have any write threads, then we dont WriteMaxFlushSize and WriteQueueSize
	// are not used.  If we do have write threads, then we need to make sure that
	if cfg.WriteConcurrency > 0 {
		if cfg.WriteMaxFlushSize > 100000 {
			return fmt.Errorf("write-max-flush-size must be <= 100000.")
		}
		if cfg.WriteMaxFlushSize >= cfg.WriteQueueSize {
			return fmt.Errorf("write-queue-size must be larger then write-max-flush-size")
		}
	}
	return nil
}

// return StoreConfig with default values set.
func NewStoreConfig() *StoreConfig {
	return &StoreConfig{
		Enabled:           false,
		GcpProject:        "default",
		BigtableInstance:  "default",
		TableName:         "metrics",
		WriteQueueSize:    100000,
		WriteMaxFlushSize: 10000,
		WriteConcurrency:  10,
		ReadConcurrency:   20,
		MaxChunkSpan:      time.Hour * 6,
		ReadTimeout:       time.Second * 5,
		WriteTimeout:      time.Second * 5,
		CreateCF:          true,
	}
}

var CliConfig = NewStoreConfig()

func ConfigSetup() {
	btStore := flag.NewFlagSet("bigtable-store", flag.ExitOnError)
	btStore.BoolVar(&CliConfig.Enabled, "enabled", CliConfig.Enabled, "enable the bigtable backend store plugin")
	btStore.StringVar(&CliConfig.GcpProject, "gcp-project", CliConfig.GcpProject, "Name of GCP project the bigtable cluster resides in")
	btStore.StringVar(&CliConfig.BigtableInstance, "bigtable-instance", CliConfig.BigtableInstance, "Name of bigtable instance")
	btStore.StringVar(&CliConfig.TableName, "table-name", CliConfig.TableName, "Name of bigtable table used for chunks")
	btStore.IntVar(&CliConfig.WriteQueueSize, "write-queue-size", CliConfig.WriteQueueSize, "Max number of chunks, per write thread, allowed to be unwritten to bigtable. Must be larger then write-max-flush-size")
	btStore.IntVar(&CliConfig.WriteMaxFlushSize, "write-max-flush-size", CliConfig.WriteMaxFlushSize, "Max number of chunks in each batch write to bigtable")
	btStore.IntVar(&CliConfig.WriteConcurrency, "write-concurrency", CliConfig.WriteConcurrency, "Number of writer threads to use.")
	btStore.IntVar(&CliConfig.ReadConcurrency, "read-concurrency", CliConfig.ReadConcurrency, "Number concurrent reads that can be processed")
	btStore.DurationVar(&CliConfig.MaxChunkSpan, "max-chunkspan", CliConfig.MaxChunkSpan, "Maximum chunkspan size used.")
	btStore.DurationVar(&CliConfig.ReadTimeout, "read-timeout", CliConfig.ReadTimeout, "read timeout")
	btStore.DurationVar(&CliConfig.WriteTimeout, "write-timeout", CliConfig.WriteTimeout, "write timeout")
	btStore.BoolVar(&CliConfig.CreateCF, "create-cf", CliConfig.CreateCF, "enable the creation of the table and column families")

	globalconf.Register("bigtable-store", btStore)
	return
}

func ConfigProcess() {
	if err := CliConfig.Validate(); err != nil {
		log.Fatalf("bigtable-store: Config validation error. %s", err)
	}
}
<|endoftext|>