<commit_before>package weed_server

import (
	"context"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/golang/protobuf/proto"
	"strings"
	"time"
)

func (fs *FilerServer) DownloadToLocal(ctx context.Context, req *filer_pb.DownloadToLocalRequest) (*filer_pb.DownloadToLocalResponse, error) {

	// load all mappings
	mappingEntry, err := fs.filer.FindEntry(ctx, util.JoinPath(filer.DirectoryEtcRemote, filer.REMOTE_STORAGE_MOUNT_FILE))
	if err != nil {
		return nil, err
	}
	mappings, err := filer.UnmarshalRemoteStorageMappings(mappingEntry.Content)
	if err != nil {
		return nil, err
	}

	// find mapping
	var remoteStorageMountedLocation *filer_pb.RemoteStorageLocation
	var localMountedDir string
	for k, loc := range mappings.Mappings {
		if strings.HasPrefix(req.Directory, k) {
			localMountedDir, remoteStorageMountedLocation = k, loc
		}
	}
	if localMountedDir == "" {
		return nil, fmt.Errorf("%s is not mounted", req.Directory)
	}

	// find storage configuration
	storageConfEntry, err := fs.filer.FindEntry(ctx, util.JoinPath(filer.DirectoryEtcRemote, remoteStorageMountedLocation.Name+filer.REMOTE_STORAGE_CONF_SUFFIX))
	if err != nil {
		return nil, err
	}
	storageConf := &filer_pb.RemoteConf{}
	if unMarshalErr := proto.Unmarshal(storageConfEntry.Content, storageConf); unMarshalErr != nil {
		return nil, fmt.Errorf("unmarshal remote storage conf %s/%s: %v", filer.DirectoryEtcRemote, remoteStorageMountedLocation.Name+filer.REMOTE_STORAGE_CONF_SUFFIX, unMarshalErr)
	}

	// find the entry
	entry, err := fs.filer.FindEntry(ctx, util.JoinPath(req.Directory, req.Name))
	if err == filer_pb.ErrNotFound {
		return nil, err
	}

	resp := &filer_pb.DownloadToLocalResponse{}
	if entry.Remote == nil || entry.Remote.RemoteSize == 0 {
		return resp, nil
	}

	// detect storage option
	// replication level is set to "000" to ensure only need to ask one volume server to fetch the data.
	so, err := fs.detectStorageOption(req.Directory, "", "000", 0, "", "", "")
	if err != nil {
		return resp, err
	}
	assignRequest, altRequest := so.ToAssignRequests(1)

	// find a good chunk size
	chunkSize := int64(5 * 1024 * 1024)
	chunkCount := entry.Remote.RemoteSize/chunkSize + 1
	for chunkCount > 1000 && chunkSize < int64(fs.option.MaxMB)*1024*1024/2 {
		chunkSize *= 2
		chunkCount = entry.Remote.RemoteSize/chunkSize + 1
	}

	dest := util.FullPath(remoteStorageMountedLocation.Path).Child(string(util.FullPath(req.Directory).Child(req.Name))[len(localMountedDir):])

	var chunks []*filer_pb.FileChunk

	// FIXME limit on parallel
	for offset := int64(0); offset < entry.Remote.RemoteSize; offset += chunkSize {
		size := chunkSize
		if offset+chunkSize > entry.Remote.RemoteSize {
			size = entry.Remote.RemoteSize - offset
		}

		// assign one volume server
		assignResult, err := operation.Assign(fs.filer.GetMaster, fs.grpcDialOption, assignRequest, altRequest)
		if err != nil {
			return resp, err
		}
		if assignResult.Error != "" {
			return resp, fmt.Errorf("assign: %v", assignResult.Error)
		}
		fileId, parseErr := needle.ParseFileIdFromString(assignResult.Fid)
		if assignResult.Error != "" {
			return resp, fmt.Errorf("unrecognized file id %s: %v", assignResult.Fid, parseErr)
		}

		// tell filer to tell volume server to download into needles
		err = operation.WithVolumeServerClient(assignResult.Url, fs.grpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
			_, fetchAndWriteErr := volumeServerClient.FetchAndWriteNeedle(context.Background(), &volume_server_pb.FetchAndWriteNeedleRequest{
				VolumeId:     uint32(fileId.VolumeId),
				NeedleId:     uint64(fileId.Key),
				Cookie:       uint32(fileId.Cookie),
				Offset:       offset,
				Size:         size,
				RemoteType:   storageConf.Type,
				RemoteName:   storageConf.Name,
				S3AccessKey:  storageConf.S3AccessKey,
				S3SecretKey:  storageConf.S3SecretKey,
				S3Region:     storageConf.S3Region,
				S3Endpoint:   storageConf.S3Endpoint,
				RemoteBucket: remoteStorageMountedLocation.Bucket,
				RemotePath:   string(dest),
			})
			if fetchAndWriteErr != nil {
				return fmt.Errorf("volume server %s fetchAndWrite %s: %v", assignResult.Url, dest, fetchAndWriteErr)
			}
			return nil
		})

		if err != nil {
			return nil, err
		}

		chunks = append(chunks, &filer_pb.FileChunk{
			FileId: assignResult.Fid,
			Offset: offset,
			Size:   uint64(size),
			Mtime:  time.Now().Unix(),
			Fid: &filer_pb.FileId{
				VolumeId: uint32(fileId.VolumeId),
				FileKey:  uint64(fileId.Key),
				Cookie:   uint32(fileId.Cookie),
			},
		})

	}

	garbage := entry.Chunks

	newEntry := entry.ShallowClone()
	newEntry.Chunks = chunks
	newEntry.Remote = proto.Clone(entry.Remote).(*filer_pb.RemoteEntry)
	newEntry.Remote.LocalMtime = time.Now().Unix()

	// this skips meta data log events

	if err := fs.filer.Store.UpdateEntry(context.Background(), newEntry); err != nil {
		return nil, err
	}
	fs.filer.DeleteChunks(garbage)

	fs.filer.NotifyUpdateEvent(ctx, entry, newEntry, true, false, nil)

	resp.Entry = newEntry.ToProtoEntry()

	return resp, nil

}
<commit_msg>parallelize remote content fetching<commit_after>package weed_server

import (
	"context"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/golang/protobuf/proto"
	"strings"
	"time"
)

func (fs *FilerServer) DownloadToLocal(ctx context.Context, req *filer_pb.DownloadToLocalRequest) (*filer_pb.DownloadToLocalResponse, error) {

	// load all mappings
	mappingEntry, err := fs.filer.FindEntry(ctx, util.JoinPath(filer.DirectoryEtcRemote, filer.REMOTE_STORAGE_MOUNT_FILE))
	if err != nil {
		return nil, err
	}
	mappings, err := filer.UnmarshalRemoteStorageMappings(mappingEntry.Content)
	if err != nil {
		return nil, err
	}

	// find mapping
	var remoteStorageMountedLocation *filer_pb.RemoteStorageLocation
	var localMountedDir string
	for k, loc := range mappings.Mappings {
		if strings.HasPrefix(req.Directory, k) {
			localMountedDir, remoteStorageMountedLocation = k, loc
		}
	}
	if localMountedDir == "" {
		return nil, fmt.Errorf("%s is not mounted", req.Directory)
	}

	// find storage configuration
	storageConfEntry, err := fs.filer.FindEntry(ctx, util.JoinPath(filer.DirectoryEtcRemote, remoteStorageMountedLocation.Name+filer.REMOTE_STORAGE_CONF_SUFFIX))
	if err != nil {
		return nil, err
	}
	storageConf := &filer_pb.RemoteConf{}
	if unMarshalErr := proto.Unmarshal(storageConfEntry.Content, storageConf); unMarshalErr != nil {
		return nil, fmt.Errorf("unmarshal remote storage conf %s/%s: %v", filer.DirectoryEtcRemote, remoteStorageMountedLocation.Name+filer.REMOTE_STORAGE_CONF_SUFFIX, unMarshalErr)
	}

	// find the entry
	entry, err := fs.filer.FindEntry(ctx, util.JoinPath(req.Directory, req.Name))
	if err == filer_pb.ErrNotFound {
		return nil, err
	}

	resp := &filer_pb.DownloadToLocalResponse{}
	if entry.Remote == nil || entry.Remote.RemoteSize == 0 {
		return resp, nil
	}

	// detect storage option
	// replication level is set to "000" to ensure only need to ask one volume server to fetch the data.
	so, err := fs.detectStorageOption(req.Directory, "", "000", 0, "", "", "")
	if err != nil {
		return resp, err
	}
	assignRequest, altRequest := so.ToAssignRequests(1)

	// find a good chunk size
	chunkSize := int64(5 * 1024 * 1024)
	chunkCount := entry.Remote.RemoteSize/chunkSize + 1
	for chunkCount > 1000 && chunkSize < int64(fs.option.MaxMB)*1024*1024/2 {
		chunkSize *= 2
		chunkCount = entry.Remote.RemoteSize/chunkSize + 1
	}

	dest := util.FullPath(remoteStorageMountedLocation.Path).Child(string(util.FullPath(req.Directory).Child(req.Name))[len(localMountedDir):])

	var chunks []*filer_pb.FileChunk
	var fetchAndWriteErr error

	limitedConcurrentExecutor := util.NewLimitedConcurrentExecutor(8)
	for offset := int64(0); offset < entry.Remote.RemoteSize; offset += chunkSize {
		localOffset := offset

		limitedConcurrentExecutor.Execute(func() {
			size := chunkSize
			if localOffset+chunkSize > entry.Remote.RemoteSize {
				size = entry.Remote.RemoteSize - localOffset
			}

			// assign one volume server
			assignResult, err := operation.Assign(fs.filer.GetMaster, fs.grpcDialOption, assignRequest, altRequest)
			if err != nil {
				fetchAndWriteErr = err
				return
			}
			if assignResult.Error != "" {
				fetchAndWriteErr = fmt.Errorf("assign: %v", assignResult.Error)
				return
			}
			fileId, parseErr := needle.ParseFileIdFromString(assignResult.Fid)
			if assignResult.Error != "" {
				fetchAndWriteErr = fmt.Errorf("unrecognized file id %s: %v", assignResult.Fid, parseErr)
				return
			}

			// tell filer to tell volume server to download into needles
			err = operation.WithVolumeServerClient(assignResult.Url, fs.grpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
				_, fetchAndWriteErr := volumeServerClient.FetchAndWriteNeedle(context.Background(), &volume_server_pb.FetchAndWriteNeedleRequest{
					VolumeId:     uint32(fileId.VolumeId),
					NeedleId:     uint64(fileId.Key),
					Cookie:       uint32(fileId.Cookie),
					Offset:       localOffset,
					Size:         size,
					RemoteType:   storageConf.Type,
					RemoteName:   storageConf.Name,
					S3AccessKey:  storageConf.S3AccessKey,
					S3SecretKey:  storageConf.S3SecretKey,
					S3Region:     storageConf.S3Region,
					S3Endpoint:   storageConf.S3Endpoint,
					RemoteBucket: remoteStorageMountedLocation.Bucket,
					RemotePath:   string(dest),
				})
				if fetchAndWriteErr != nil {
					return fmt.Errorf("volume server %s fetchAndWrite %s: %v", assignResult.Url, dest, fetchAndWriteErr)
				}
				return nil
			})

			if err != nil {
				fetchAndWriteErr = err
				return
			}

			chunks = append(chunks, &filer_pb.FileChunk{
				FileId: assignResult.Fid,
				Offset: localOffset,
				Size:   uint64(size),
				Mtime:  time.Now().Unix(),
				Fid: &filer_pb.FileId{
					VolumeId: uint32(fileId.VolumeId),
					FileKey:  uint64(fileId.Key),
					Cookie:   uint32(fileId.Cookie),
				},
			})
		})
	}

	garbage := entry.Chunks

	newEntry := entry.ShallowClone()
	newEntry.Chunks = chunks
	newEntry.Remote = proto.Clone(entry.Remote).(*filer_pb.RemoteEntry)
	newEntry.Remote.LocalMtime = time.Now().Unix()

	// this skips meta data log events

	if err := fs.filer.Store.UpdateEntry(context.Background(), newEntry); err != nil {
		return nil, err
	}
	fs.filer.DeleteChunks(garbage)

	fs.filer.NotifyUpdateEvent(ctx, entry, newEntry, true, false, nil)

	resp.Entry = newEntry.ToProtoEntry()

	return resp, nil

}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 vault

import (
	"path"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	cmapi "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha2"
	cmmeta "github.com/jetstack/cert-manager/pkg/apis/meta/v1"
	"github.com/jetstack/cert-manager/test/e2e/framework"
	"github.com/jetstack/cert-manager/test/e2e/framework/addon"
	"github.com/jetstack/cert-manager/test/e2e/framework/addon/vault"
	"github.com/jetstack/cert-manager/test/e2e/suite/conformance/certificates"
)

const (
	intermediateMount      = "intermediate-ca"
	role                   = "kubernetes-vault"
	vaultSecretAppRoleName = "vault-role-"
	authPath               = "approle"
)

var _ = framework.ConformanceDescribe("Certificates", func() {
	var unsupportedFeatures = certificates.NewFeatureSet(
		certificates.KeyUsagesFeature,
	)

	provisioner := new(vaultAppRoleProvisioner)

	(&certificates.Suite{
		Name:                "VaultAppRole Issuer",
		CreateIssuerFunc:    provisioner.createIssuer,
		DeleteIssuerFunc:    provisioner.delete,
		UnsupportedFeatures: unsupportedFeatures,
	}).Define()

	(&certificates.Suite{
		Name:                "VaultAppRole ClusterIssuer",
		CreateIssuerFunc:    provisioner.createClusterIssuer,
		DeleteIssuerFunc:    provisioner.delete,
		UnsupportedFeatures: unsupportedFeatures,
	}).Define()
})

type vaultAppRoleProvisioner struct {
	vault     *vault.Vault
	vaultInit *vault.VaultInitializer

	*vaultSecrets
}

type vaultSecrets struct {
	roleID   string
	secretID string

	secretName      string
	secretNamespace string
}

func (v *vaultAppRoleProvisioner) delete(f *framework.Framework, ref cmmeta.ObjectReference) {
	Expect(v.vaultInit.Clean()).NotTo(HaveOccurred(), "failed to deprovision vault initializer")
	Expect(v.vault.Deprovision()).NotTo(HaveOccurred(), "failed to deprovision vault")

	err := f.KubeClientSet.CoreV1().Secrets(v.secretNamespace).Delete(v.secretName, nil)
	Expect(err).NotTo(HaveOccurred())

	if ref.Kind == "ClusterIssuer" {
		err = f.CertManagerClientSet.CertmanagerV1alpha2().ClusterIssuers().Delete(ref.Name, nil)
		Expect(err).NotTo(HaveOccurred())
	}
}

func (v *vaultAppRoleProvisioner) createIssuer(f *framework.Framework) cmmeta.ObjectReference {
	By("Creating a VaultAppRole Issuer")

	v.vaultSecrets = v.initVault(f)

	sec, err := f.KubeClientSet.CoreV1().Secrets(f.Namespace.Name).Create(vault.NewVaultAppRoleSecret(vaultSecretAppRoleName, v.secretID))
	Expect(err).NotTo(HaveOccurred(), "vault to store app role secret from vault")

	v.secretName = sec.Name
	v.secretNamespace = sec.Namespace

	issuer, err := f.CertManagerClientSet.CertmanagerV1alpha2().Issuers(f.Namespace.Name).Create(&cmapi.Issuer{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "vault-issuer-",
		},
		Spec: v.createIssuerSpec(f),
	})
	Expect(err).NotTo(HaveOccurred(), "failed to create vault issuer")

	return cmmeta.ObjectReference{
		Group: cmapi.SchemeGroupVersion.Group,
		Kind:  cmapi.IssuerKind,
		Name:  issuer.Name,
	}
}

func (v *vaultAppRoleProvisioner) createClusterIssuer(f *framework.Framework) cmmeta.ObjectReference {
	By("Creating a VaultAppRole ClusterIssuer")

	v.vaultSecrets = v.initVault(f)

	sec, err := f.KubeClientSet.CoreV1().Secrets(f.Config.Addons.CertManager.ClusterResourceNamespace).Create(vault.NewVaultAppRoleSecret(vaultSecretAppRoleName, v.secretID))
	Expect(err).NotTo(HaveOccurred(), "vault to store app role secret from vault")

	v.secretName = sec.Name
	v.secretNamespace = sec.Namespace

	issuer, err := f.CertManagerClientSet.CertmanagerV1alpha2().ClusterIssuers().Create(&cmapi.ClusterIssuer{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "vault-cluster-issuer-",
		},
		Spec: v.createIssuerSpec(f),
	})
	Expect(err).NotTo(HaveOccurred(), "failed to create vault issuer")

	return cmmeta.ObjectReference{
		Group: cmapi.SchemeGroupVersion.Group,
		Kind:  cmapi.ClusterIssuerKind,
		Name:  issuer.Name,
	}
}

func (v *vaultAppRoleProvisioner) initVault(f *framework.Framework) *vaultSecrets {
	v.vault = &vault.Vault{
		Base:      addon.Base,
		Namespace: f.Namespace.Name,
		Name:      "cm-e2e-create-vault-issuer",
	}
	Expect(v.vault.Setup(f.Config)).NotTo(HaveOccurred(), "failed to setup vault")
	Expect(v.vault.Provision()).NotTo(HaveOccurred(), "failed to provision vault")

	By("Configuring the VaultAppRole server")
	v.vaultInit = &vault.VaultInitializer{
		Details:           *v.vault.Details(),
		RootMount:         "root-ca",
		IntermediateMount: intermediateMount,
		Role:              role,
		AppRoleAuthPath:   authPath,
	}
	Expect(v.vaultInit.Init()).NotTo(HaveOccurred(), "failed to init vault")
	Expect(v.vaultInit.Setup()).NotTo(HaveOccurred(), "fauled to setup vault")

	roleID, secretID, err := v.vaultInit.CreateAppRole()
	Expect(err).NotTo(HaveOccurred(), "vault to create app role from vault")

	return &vaultSecrets{
		roleID:   roleID,
		secretID: secretID,
	}
}

func (v *vaultAppRoleProvisioner) createIssuerSpec(f *framework.Framework) cmapi.IssuerSpec {
	vaultPath := path.Join(intermediateMount, "sign", role)

	return cmapi.IssuerSpec{
		IssuerConfig: cmapi.IssuerConfig{
			Vault: &cmapi.VaultIssuer{
				Server:   v.vault.Details().Host,
				Path:     vaultPath,
				CABundle: v.vault.Details().VaultCA,
				Auth: cmapi.VaultAuth{
					AppRole: &cmapi.VaultAppRole{
						Path:   authPath,
						RoleId: v.roleID,
						SecretRef: cmmeta.SecretKeySelector{
							Key: "secretkey",
							LocalObjectReference: cmmeta.LocalObjectReference{
								Name: v.secretName,
							},
						},
					},
				},
			},
		},
	}
}
<commit_msg>spelling: failed<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 vault

import (
	"path"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	cmapi "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha2"
	cmmeta "github.com/jetstack/cert-manager/pkg/apis/meta/v1"
	"github.com/jetstack/cert-manager/test/e2e/framework"
	"github.com/jetstack/cert-manager/test/e2e/framework/addon"
	"github.com/jetstack/cert-manager/test/e2e/framework/addon/vault"
	"github.com/jetstack/cert-manager/test/e2e/suite/conformance/certificates"
)

const (
	intermediateMount      = "intermediate-ca"
	role                   = "kubernetes-vault"
	vaultSecretAppRoleName = "vault-role-"
	authPath               = "approle"
)

var _ = framework.ConformanceDescribe("Certificates", func() {
	var unsupportedFeatures = certificates.NewFeatureSet(
		certificates.KeyUsagesFeature,
	)

	provisioner := new(vaultAppRoleProvisioner)

	(&certificates.Suite{
		Name:                "VaultAppRole Issuer",
		CreateIssuerFunc:    provisioner.createIssuer,
		DeleteIssuerFunc:    provisioner.delete,
		UnsupportedFeatures: unsupportedFeatures,
	}).Define()

	(&certificates.Suite{
		Name:                "VaultAppRole ClusterIssuer",
		CreateIssuerFunc:    provisioner.createClusterIssuer,
		DeleteIssuerFunc:    provisioner.delete,
		UnsupportedFeatures: unsupportedFeatures,
	}).Define()
})

type vaultAppRoleProvisioner struct {
	vault     *vault.Vault
	vaultInit *vault.VaultInitializer

	*vaultSecrets
}

type vaultSecrets struct {
	roleID   string
	secretID string

	secretName      string
	secretNamespace string
}

func (v *vaultAppRoleProvisioner) delete(f *framework.Framework, ref cmmeta.ObjectReference) {
	Expect(v.vaultInit.Clean()).NotTo(HaveOccurred(), "failed to deprovision vault initializer")
	Expect(v.vault.Deprovision()).NotTo(HaveOccurred(), "failed to deprovision vault")

	err := f.KubeClientSet.CoreV1().Secrets(v.secretNamespace).Delete(v.secretName, nil)
	Expect(err).NotTo(HaveOccurred())

	if ref.Kind == "ClusterIssuer" {
		err = f.CertManagerClientSet.CertmanagerV1alpha2().ClusterIssuers().Delete(ref.Name, nil)
		Expect(err).NotTo(HaveOccurred())
	}
}

func (v *vaultAppRoleProvisioner) createIssuer(f *framework.Framework) cmmeta.ObjectReference {
	By("Creating a VaultAppRole Issuer")

	v.vaultSecrets = v.initVault(f)

	sec, err := f.KubeClientSet.CoreV1().Secrets(f.Namespace.Name).Create(vault.NewVaultAppRoleSecret(vaultSecretAppRoleName, v.secretID))
	Expect(err).NotTo(HaveOccurred(), "vault to store app role secret from vault")

	v.secretName = sec.Name
	v.secretNamespace = sec.Namespace

	issuer, err := f.CertManagerClientSet.CertmanagerV1alpha2().Issuers(f.Namespace.Name).Create(&cmapi.Issuer{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "vault-issuer-",
		},
		Spec: v.createIssuerSpec(f),
	})
	Expect(err).NotTo(HaveOccurred(), "failed to create vault issuer")

	return cmmeta.ObjectReference{
		Group: cmapi.SchemeGroupVersion.Group,
		Kind:  cmapi.IssuerKind,
		Name:  issuer.Name,
	}
}

func (v *vaultAppRoleProvisioner) createClusterIssuer(f *framework.Framework) cmmeta.ObjectReference {
	By("Creating a VaultAppRole ClusterIssuer")

	v.vaultSecrets = v.initVault(f)

	sec, err := f.KubeClientSet.CoreV1().Secrets(f.Config.Addons.CertManager.ClusterResourceNamespace).Create(vault.NewVaultAppRoleSecret(vaultSecretAppRoleName, v.secretID))
	Expect(err).NotTo(HaveOccurred(), "vault to store app role secret from vault")

	v.secretName = sec.Name
	v.secretNamespace = sec.Namespace

	issuer, err := f.CertManagerClientSet.CertmanagerV1alpha2().ClusterIssuers().Create(&cmapi.ClusterIssuer{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "vault-cluster-issuer-",
		},
		Spec: v.createIssuerSpec(f),
	})
	Expect(err).NotTo(HaveOccurred(), "failed to create vault issuer")

	return cmmeta.ObjectReference{
		Group: cmapi.SchemeGroupVersion.Group,
		Kind:  cmapi.ClusterIssuerKind,
		Name:  issuer.Name,
	}
}

func (v *vaultAppRoleProvisioner) initVault(f *framework.Framework) *vaultSecrets {
	v.vault = &vault.Vault{
		Base:      addon.Base,
		Namespace: f.Namespace.Name,
		Name:      "cm-e2e-create-vault-issuer",
	}
	Expect(v.vault.Setup(f.Config)).NotTo(HaveOccurred(), "failed to setup vault")
	Expect(v.vault.Provision()).NotTo(HaveOccurred(), "failed to provision vault")

	By("Configuring the VaultAppRole server")
	v.vaultInit = &vault.VaultInitializer{
		Details:           *v.vault.Details(),
		RootMount:         "root-ca",
		IntermediateMount: intermediateMount,
		Role:              role,
		AppRoleAuthPath:   authPath,
	}
	Expect(v.vaultInit.Init()).NotTo(HaveOccurred(), "failed to init vault")
	Expect(v.vaultInit.Setup()).NotTo(HaveOccurred(), "failed to setup vault")

	roleID, secretID, err := v.vaultInit.CreateAppRole()
	Expect(err).NotTo(HaveOccurred(), "vault to create app role from vault")

	return &vaultSecrets{
		roleID:   roleID,
		secretID: secretID,
	}
}

func (v *vaultAppRoleProvisioner) createIssuerSpec(f *framework.Framework) cmapi.IssuerSpec {
	vaultPath := path.Join(intermediateMount, "sign", role)

	return cmapi.IssuerSpec{
		IssuerConfig: cmapi.IssuerConfig{
			Vault: &cmapi.VaultIssuer{
				Server:   v.vault.Details().Host,
				Path:     vaultPath,
				CABundle: v.vault.Details().VaultCA,
				Auth: cmapi.VaultAuth{
					AppRole: &cmapi.VaultAppRole{
						Path:   authPath,
						RoleId: v.roleID,
						SecretRef: cmmeta.SecretKeySelector{
							Key: "secretkey",
							LocalObjectReference: cmmeta.LocalObjectReference{
								Name: v.secretName,
							},
						},
					},
				},
			},
		},
	}
}
<|endoftext|>


<commit_before>package hypervisor

import (
	"errors"
	"net"
)

const consoleTypeUnknown = "UNKNOWN ConsoleType"
const stateUnknown = "UNKNOWN State"
const volumeFormatUnknown = "UNKNOWN VolumeFormat"

var (
	consoleTypeToText = map[ConsoleType]string{
		ConsoleNone:  "none",
		ConsoleDummy: "dummy",
		ConsoleVNC:   "vnc",
	}
	textToConsoleType map[string]ConsoleType

	stateToText = map[State]string{
		StateStarting:      "starting",
		StateRunning:       "running",
		StateFailedToStart: "failed to start",
		StateStopping:      "stopping",
		StateStopped:       "stopped",
		StateDestroying:    "destroying",
		StateMigrating:     "migrating",
		StateExporting:     "exporting",
	}
	textToState map[string]State

	volumeFormatToText = map[VolumeFormat]string{
		VolumeFormatRaw:   "raw",
		VolumeFormatQCOW2: "qcow2",
	}
	textToVolumeFormat map[string]VolumeFormat
)

func init() {
	textToConsoleType = make(map[string]ConsoleType, len(consoleTypeToText))
	for consoleType, text := range consoleTypeToText {
		textToConsoleType[text] = consoleType
	}
	textToState = make(map[string]State, len(stateToText))
	for state, text := range stateToText {
		textToState[text] = state
	}
	textToVolumeFormat = make(map[string]VolumeFormat, len(volumeFormatToText))
	for format, text := range volumeFormatToText {
		textToVolumeFormat[text] = format
	}
}

func shrinkIP(netIP net.IP) net.IP {
	switch len(netIP) {
	case 4:
		return netIP
	case 16:
		if ip4 := netIP.To4(); ip4 == nil {
			return netIP
		} else {
			return ip4
		}
	default:
		return netIP
	}
}

func (left *Address) Equal(right *Address) bool {
	if !left.IpAddress.Equal(right.IpAddress) {
		return false
	}
	if left.MacAddress != right.MacAddress {
		return false
	}
	return true
}

func (address *Address) Shrink() {
	address.IpAddress = shrinkIP(address.IpAddress)
}

func stringSlicesEqual(left, right []string) bool {
	if len(left) != len(right) {
		return false
	}
	for index, leftString := range left {
		if leftString != right[index] {
			return false
		}
	}
	return true
}

func (consoleType *ConsoleType) CheckValid() error {
	if _, ok := consoleTypeToText[*consoleType]; !ok {
		return errors.New(consoleTypeUnknown)
	} else {
		return nil
	}
}

func (consoleType ConsoleType) MarshalText() ([]byte, error) {
	if text := consoleType.String(); text == consoleTypeUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (consoleType *ConsoleType) Set(value string) error {
	if val, ok := textToConsoleType[value]; !ok {
		return errors.New(consoleTypeUnknown)
	} else {
		*consoleType = val
		return nil
	}
}

func (consoleType ConsoleType) String() string {
	if str, ok := consoleTypeToText[consoleType]; !ok {
		return consoleTypeUnknown
	} else {
		return str
	}
}

func (consoleType *ConsoleType) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToConsoleType[txt]; ok {
		*consoleType = val
		return nil
	} else {
		return errors.New("unknown ConsoleType: " + txt)
	}
}

func (state State) MarshalText() ([]byte, error) {
	if text := state.String(); text == stateUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (state State) String() string {
	if text, ok := stateToText[state]; ok {
		return text
	} else {
		return stateUnknown
	}
}

func (state *State) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToState[txt]; ok {
		*state = val
		return nil
	} else {
		return errors.New("unknown State: " + txt)
	}
}

func (left *Subnet) Equal(right *Subnet) bool {
	if left.Id != right.Id {
		return false
	}
	if !left.IpGateway.Equal(right.IpGateway) {
		return false
	}
	if !left.IpMask.Equal(right.IpMask) {
		return false
	}
	if left.DomainName != right.DomainName {
		return false
	}
	if !IpListsEqual(left.DomainNameServers, right.DomainNameServers) {
		return false
	}
	if left.Manage != right.Manage {
		return false
	}
	if left.VlanId != right.VlanId {
		return false
	}
	if !stringSlicesEqual(left.AllowedGroups, right.AllowedGroups) {
		return false
	}
	if !stringSlicesEqual(left.AllowedUsers, right.AllowedUsers) {
		return false
	}
	return true
}

func IpListsEqual(left, right []net.IP) bool {
	if len(left) != len(right) {
		return false
	}
	for index, leftIP := range left {
		if !leftIP.Equal(right[index]) {
			return false
		}
	}
	return true
}

func (subnet *Subnet) Shrink() {
	subnet.IpGateway = shrinkIP(subnet.IpGateway)
	subnet.IpMask = shrinkIP(subnet.IpMask)
	for index, ipAddr := range subnet.DomainNameServers {
		subnet.DomainNameServers[index] = shrinkIP(ipAddr)
	}
}

func (left *VmInfo) Equal(right *VmInfo) bool {
	if !left.Address.Equal(&right.Address) {
		return false
	}
	if left.Hostname != right.Hostname {
		return false
	}
	if left.ImageName != right.ImageName {
		return false
	}
	if left.ImageURL != right.ImageURL {
		return false
	}
	if left.MemoryInMiB != right.MemoryInMiB {
		return false
	}
	if left.MilliCPUs != right.MilliCPUs {
		return false
	}
	if !stringSlicesEqual(left.OwnerGroups, right.OwnerGroups) {
		return false
	}
	if !stringSlicesEqual(left.OwnerUsers, right.OwnerUsers) {
		return false
	}
	if left.SpreadVolumes != right.SpreadVolumes {
		return false
	}
	if left.State != right.State {
		return false
	}
	if !left.Tags.Equal(right.Tags) {
		return false
	}
	if len(left.SecondaryAddresses) != len(right.SecondaryAddresses) {
		return false
	}
	for index, leftAddress := range left.SecondaryAddresses {
		if !leftAddress.Equal(&right.SecondaryAddresses[index]) {
			return false
		}
	}
	if !stringSlicesEqual(left.SecondarySubnetIDs, right.SecondarySubnetIDs) {
		return false
	}
	if left.SubnetId != right.SubnetId {
		return false
	}
	if left.Uncommitted != right.Uncommitted {
		return false
	}
	if len(left.Volumes) != len(right.Volumes) {
		return false
	}
	for index, leftVolume := range left.Volumes {
		if leftVolume != right.Volumes[index] {
			return false
		}
	}
	return true
}

func (volumeFormat VolumeFormat) MarshalText() ([]byte, error) {
	if text := volumeFormat.String(); text == volumeFormatUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (volumeFormat VolumeFormat) String() string {
	if text, ok := volumeFormatToText[volumeFormat]; ok {
		return text
	} else {
		return volumeFormatUnknown
	}
}

func (volumeFormat *VolumeFormat) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToVolumeFormat[txt]; ok {
		*volumeFormat = val
		return nil
	} else {
		return errors.New("unknown VolumeFormat: " + txt)
	}
}
<commit_msg>Add ConsoleType and DestroyProtection to proto/hypervisor.vmInfo.Equal().<commit_after>package hypervisor

import (
	"errors"
	"net"
)

const consoleTypeUnknown = "UNKNOWN ConsoleType"
const stateUnknown = "UNKNOWN State"
const volumeFormatUnknown = "UNKNOWN VolumeFormat"

var (
	consoleTypeToText = map[ConsoleType]string{
		ConsoleNone:  "none",
		ConsoleDummy: "dummy",
		ConsoleVNC:   "vnc",
	}
	textToConsoleType map[string]ConsoleType

	stateToText = map[State]string{
		StateStarting:      "starting",
		StateRunning:       "running",
		StateFailedToStart: "failed to start",
		StateStopping:      "stopping",
		StateStopped:       "stopped",
		StateDestroying:    "destroying",
		StateMigrating:     "migrating",
		StateExporting:     "exporting",
	}
	textToState map[string]State

	volumeFormatToText = map[VolumeFormat]string{
		VolumeFormatRaw:   "raw",
		VolumeFormatQCOW2: "qcow2",
	}
	textToVolumeFormat map[string]VolumeFormat
)

func init() {
	textToConsoleType = make(map[string]ConsoleType, len(consoleTypeToText))
	for consoleType, text := range consoleTypeToText {
		textToConsoleType[text] = consoleType
	}
	textToState = make(map[string]State, len(stateToText))
	for state, text := range stateToText {
		textToState[text] = state
	}
	textToVolumeFormat = make(map[string]VolumeFormat, len(volumeFormatToText))
	for format, text := range volumeFormatToText {
		textToVolumeFormat[text] = format
	}
}

func shrinkIP(netIP net.IP) net.IP {
	switch len(netIP) {
	case 4:
		return netIP
	case 16:
		if ip4 := netIP.To4(); ip4 == nil {
			return netIP
		} else {
			return ip4
		}
	default:
		return netIP
	}
}

func (left *Address) Equal(right *Address) bool {
	if !left.IpAddress.Equal(right.IpAddress) {
		return false
	}
	if left.MacAddress != right.MacAddress {
		return false
	}
	return true
}

func (address *Address) Shrink() {
	address.IpAddress = shrinkIP(address.IpAddress)
}

func stringSlicesEqual(left, right []string) bool {
	if len(left) != len(right) {
		return false
	}
	for index, leftString := range left {
		if leftString != right[index] {
			return false
		}
	}
	return true
}

func (consoleType *ConsoleType) CheckValid() error {
	if _, ok := consoleTypeToText[*consoleType]; !ok {
		return errors.New(consoleTypeUnknown)
	} else {
		return nil
	}
}

func (consoleType ConsoleType) MarshalText() ([]byte, error) {
	if text := consoleType.String(); text == consoleTypeUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (consoleType *ConsoleType) Set(value string) error {
	if val, ok := textToConsoleType[value]; !ok {
		return errors.New(consoleTypeUnknown)
	} else {
		*consoleType = val
		return nil
	}
}

func (consoleType ConsoleType) String() string {
	if str, ok := consoleTypeToText[consoleType]; !ok {
		return consoleTypeUnknown
	} else {
		return str
	}
}

func (consoleType *ConsoleType) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToConsoleType[txt]; ok {
		*consoleType = val
		return nil
	} else {
		return errors.New("unknown ConsoleType: " + txt)
	}
}

func (state State) MarshalText() ([]byte, error) {
	if text := state.String(); text == stateUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (state State) String() string {
	if text, ok := stateToText[state]; ok {
		return text
	} else {
		return stateUnknown
	}
}

func (state *State) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToState[txt]; ok {
		*state = val
		return nil
	} else {
		return errors.New("unknown State: " + txt)
	}
}

func (left *Subnet) Equal(right *Subnet) bool {
	if left.Id != right.Id {
		return false
	}
	if !left.IpGateway.Equal(right.IpGateway) {
		return false
	}
	if !left.IpMask.Equal(right.IpMask) {
		return false
	}
	if left.DomainName != right.DomainName {
		return false
	}
	if !IpListsEqual(left.DomainNameServers, right.DomainNameServers) {
		return false
	}
	if left.Manage != right.Manage {
		return false
	}
	if left.VlanId != right.VlanId {
		return false
	}
	if !stringSlicesEqual(left.AllowedGroups, right.AllowedGroups) {
		return false
	}
	if !stringSlicesEqual(left.AllowedUsers, right.AllowedUsers) {
		return false
	}
	return true
}

func IpListsEqual(left, right []net.IP) bool {
	if len(left) != len(right) {
		return false
	}
	for index, leftIP := range left {
		if !leftIP.Equal(right[index]) {
			return false
		}
	}
	return true
}

func (subnet *Subnet) Shrink() {
	subnet.IpGateway = shrinkIP(subnet.IpGateway)
	subnet.IpMask = shrinkIP(subnet.IpMask)
	for index, ipAddr := range subnet.DomainNameServers {
		subnet.DomainNameServers[index] = shrinkIP(ipAddr)
	}
}

func (left *VmInfo) Equal(right *VmInfo) bool {
	if !left.Address.Equal(&right.Address) {
		return false
	}
	if left.ConsoleType != right.ConsoleType {
		return false
	}
	if left.DestroyProtection != right.DestroyProtection {
		return false
	}
	if left.Hostname != right.Hostname {
		return false
	}
	if left.ImageName != right.ImageName {
		return false
	}
	if left.ImageURL != right.ImageURL {
		return false
	}
	if left.MemoryInMiB != right.MemoryInMiB {
		return false
	}
	if left.MilliCPUs != right.MilliCPUs {
		return false
	}
	if !stringSlicesEqual(left.OwnerGroups, right.OwnerGroups) {
		return false
	}
	if !stringSlicesEqual(left.OwnerUsers, right.OwnerUsers) {
		return false
	}
	if left.SpreadVolumes != right.SpreadVolumes {
		return false
	}
	if left.State != right.State {
		return false
	}
	if !left.Tags.Equal(right.Tags) {
		return false
	}
	if len(left.SecondaryAddresses) != len(right.SecondaryAddresses) {
		return false
	}
	for index, leftAddress := range left.SecondaryAddresses {
		if !leftAddress.Equal(&right.SecondaryAddresses[index]) {
			return false
		}
	}
	if !stringSlicesEqual(left.SecondarySubnetIDs, right.SecondarySubnetIDs) {
		return false
	}
	if left.SubnetId != right.SubnetId {
		return false
	}
	if left.Uncommitted != right.Uncommitted {
		return false
	}
	if len(left.Volumes) != len(right.Volumes) {
		return false
	}
	for index, leftVolume := range left.Volumes {
		if leftVolume != right.Volumes[index] {
			return false
		}
	}
	return true
}

func (volumeFormat VolumeFormat) MarshalText() ([]byte, error) {
	if text := volumeFormat.String(); text == volumeFormatUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (volumeFormat VolumeFormat) String() string {
	if text, ok := volumeFormatToText[volumeFormat]; ok {
		return text
	} else {
		return volumeFormatUnknown
	}
}

func (volumeFormat *VolumeFormat) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToVolumeFormat[txt]; ok {
		*volumeFormat = val
		return nil
	} else {
		return errors.New("unknown VolumeFormat: " + txt)
	}
}
<|endoftext|>


<commit_before>// The digitalocean package contains a packer.Builder implementation
// that builds DigitalOcean images (snapshots).

package digitalocean

import (
	"log"
	"time"

	"github.com/digitalocean/godo"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
	"golang.org/x/oauth2"
)

// The unique id for the builder
const BuilderId = "pearkes.digitalocean"

type Builder struct {
	config Config
	runner multistep.Runner
}

func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	c, warnings, errs := NewConfig(raws...)
	if errs != nil {
		return warnings, errs
	}
	b.config = *c

	return nil, nil
}

func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
	client := godo.NewClient(oauth2.NewClient(oauth2.NoContext, &apiTokenSource{
		AccessToken: b.config.APIToken,
	}))

	// Set up the state
	state := new(multistep.BasicStateBag)
	state.Put("config", b.config)
	state.Put("client", client)
	state.Put("hook", hook)
	state.Put("ui", ui)

	// Build the steps
	steps := []multistep.Step{
		&stepCreateSSHKey{
			Debug: b.config.PackerDebug,
			DebugKeyPath: fmt.Sprintf("do_%s.pem", b.config.PackerBuildName),
		},
		new(stepCreateDroplet),
		new(stepDropletInfo),
		&common.StepConnectSSH{
			SSHAddress:     sshAddress,
			SSHConfig:      sshConfig,
			SSHWaitTimeout: 5 * time.Minute,
		},
		new(common.StepProvision),
		new(stepShutdown),
		new(stepPowerOff),
		new(stepSnapshot),
	}

	// Run the steps
	if b.config.PackerDebug {
		b.runner = &multistep.DebugRunner{
			Steps:   steps,
			PauseFn: common.MultistepDebugFn(ui),
		}
	} else {
		b.runner = &multistep.BasicRunner{Steps: steps}
	}

	b.runner.Run(state)

	// If there was an error, return that
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	if _, ok := state.GetOk("snapshot_name"); !ok {
		log.Println("Failed to find snapshot_name in state. Bug?")
		return nil, nil
	}

	artifact := &Artifact{
		snapshotName: state.Get("snapshot_name").(string),
		snapshotId:   state.Get("snapshot_image_id").(int),
		regionName:   state.Get("region").(string),
		client:       client,
	}

	return artifact, nil
}

func (b *Builder) Cancel() {
	if b.runner != nil {
		log.Println("Cancelling the step runner...")
		b.runner.Cancel()
	}
}
<commit_msg>builder/digitalocean: fix build<commit_after>// The digitalocean package contains a packer.Builder implementation
// that builds DigitalOcean images (snapshots).

package digitalocean

import (
	"fmt"
	"log"
	"time"

	"github.com/digitalocean/godo"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
	"golang.org/x/oauth2"
)

// The unique id for the builder
const BuilderId = "pearkes.digitalocean"

type Builder struct {
	config Config
	runner multistep.Runner
}

func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	c, warnings, errs := NewConfig(raws...)
	if errs != nil {
		return warnings, errs
	}
	b.config = *c

	return nil, nil
}

func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
	client := godo.NewClient(oauth2.NewClient(oauth2.NoContext, &apiTokenSource{
		AccessToken: b.config.APIToken,
	}))

	// Set up the state
	state := new(multistep.BasicStateBag)
	state.Put("config", b.config)
	state.Put("client", client)
	state.Put("hook", hook)
	state.Put("ui", ui)

	// Build the steps
	steps := []multistep.Step{
		&stepCreateSSHKey{
			Debug:        b.config.PackerDebug,
			DebugKeyPath: fmt.Sprintf("do_%s.pem", b.config.PackerBuildName),
		},
		new(stepCreateDroplet),
		new(stepDropletInfo),
		&common.StepConnectSSH{
			SSHAddress:     sshAddress,
			SSHConfig:      sshConfig,
			SSHWaitTimeout: 5 * time.Minute,
		},
		new(common.StepProvision),
		new(stepShutdown),
		new(stepPowerOff),
		new(stepSnapshot),
	}

	// Run the steps
	if b.config.PackerDebug {
		b.runner = &multistep.DebugRunner{
			Steps:   steps,
			PauseFn: common.MultistepDebugFn(ui),
		}
	} else {
		b.runner = &multistep.BasicRunner{Steps: steps}
	}

	b.runner.Run(state)

	// If there was an error, return that
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	if _, ok := state.GetOk("snapshot_name"); !ok {
		log.Println("Failed to find snapshot_name in state. Bug?")
		return nil, nil
	}

	artifact := &Artifact{
		snapshotName: state.Get("snapshot_name").(string),
		snapshotId:   state.Get("snapshot_image_id").(int),
		regionName:   state.Get("region").(string),
		client:       client,
	}

	return artifact, nil
}

func (b *Builder) Cancel() {
	if b.runner != nil {
		log.Println("Cancelling the step runner...")
		b.runner.Cancel()
	}
}
<|endoftext|>


<commit_before>package common

import (
	"bytes"
	"fmt"
	"log"
	"os/exec"
	"runtime"
	"strconv"
	"strings"

	"github.com/mitchellh/multistep"
)

// A driver is able to talk to VMware, control virtual machines, etc.
type Driver interface {
	// Clone clones the VMX and the disk to the destination path. The
	// destination is a path to the VMX file. The disk will be copied
	// to that same directory.
	Clone(dst string, src string) error

	// CompactDisk compacts a virtual disk.
	CompactDisk(string) error

	// CreateDisk creates a virtual disk with the given size.
	CreateDisk(string, string, string) error

	// Checks if the VMX file at the given path is running.
	IsRunning(string) (bool, error)

	// SSHAddress returns the SSH address for the VM that is being
	// managed by this driver.
	SSHAddress(multistep.StateBag) (string, error)

	// Start starts a VM specified by the path to the VMX given.
	Start(string, bool) error

	// Stop stops a VM specified by the path to the VMX given.
	Stop(string) error

	// SuppressMessages modifies the VMX or surrounding directory so that
	// VMware doesn't show any annoying messages.
	SuppressMessages(string) error

	// Get the path to the VMware ISO for the given flavor.
	ToolsIsoPath(string) string

	// Attach the VMware tools ISO
	ToolsInstall() error

	// Get the path to the DHCP leases file for the given device.
	DhcpLeasesPath(string) string

	// Verify checks to make sure that this driver should function
	// properly. This should check that all the files it will use
	// appear to exist and so on. If everything is okay, this doesn't
	// return an error. Otherwise, this returns an error.
	Verify() error
}

// NewDriver returns a new driver implementation for this operating
// system, or an error if the driver couldn't be initialized.
func NewDriver(dconfig *DriverConfig, config *SSHConfig) (Driver, error) {
	drivers := []Driver{}

	switch runtime.GOOS {
	case "darwin":
		drivers = []Driver{
			&Fusion6Driver{
				Fusion5Driver: Fusion5Driver{
					AppPath:   dconfig.FusionAppPath,
					SSHConfig: config,
				},
			},
			&Fusion5Driver{
				AppPath:   dconfig.FusionAppPath,
				SSHConfig: config,
			},
		}
	case "linux":
		fallthrough
	case "windows":
		drivers = []Driver{
			&Workstation10Driver{
				Workstation9Driver: Workstation9Driver{
					SSHConfig: config,
				},
			},
			&Workstation9Driver{
				SSHConfig: config,
			},
			&Player6Driver{
				Player5Driver: Player5Driver{
					SSHConfig: config,
				},
			},
			&Player5Driver{
				SSHConfig: config,
			},
		}
	default:
		return nil, fmt.Errorf("can't find driver for OS: %s", runtime.GOOS)
	}

	errs := ""
	for _, driver := range drivers {
		err := driver.Verify()
		if err == nil {
			return driver, nil
		}
		errs += "* " + err.Error() + "\n"
	}

	return nil, fmt.Errorf(
		"Unable to initialize any driver for this platform. The errors\n"+
			"from each driver are shown below. Please fix at least one driver\n"+
			"to continue:\n%s", errs)
}

func runAndLog(cmd *exec.Cmd) (string, string, error) {
	var stdout, stderr bytes.Buffer

	log.Printf("Executing: %s %v", cmd.Path, cmd.Args[1:])
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()

	stdoutString := strings.TrimSpace(stdout.String())
	stderrString := strings.TrimSpace(stderr.String())

	if _, ok := err.(*exec.ExitError); ok {
		message := stderrString
		if message == "" {
			message = stdoutString
		}

		err = fmt.Errorf("VMware error: %s", message)
	}

	log.Printf("stdout: %s", stdoutString)
	log.Printf("stderr: %s", stderrString)

	// Replace these for Windows, we only want to deal with Unix
	// style line endings.
	returnStdout := strings.Replace(stdout.String(), "\r\n", "\n", -1)
	returnStderr := strings.Replace(stderr.String(), "\r\n", "\n", -1)

	return returnStdout, returnStderr, err
}

func normalizeVersion(version string) (string, error) {
	i, err := strconv.Atoi(version)
	if err != nil {
		return "", fmt.Errorf(
			"VMware version '%s' is not numeric", version)
	}

	return fmt.Sprintf("%02d", i), nil
}

func compareVersions(versionFound string, versionWanted string, product string) error {
	found, err := normalizeVersion(versionFound)
	if err != nil {
		return err
	}

	wanted, err := normalizeVersion(versionWanted)
	if err != nil {
		return err
	}

	if found < wanted {
		return fmt.Errorf(
			"VMware %s version %s, or greater, is required. Found version: %s", product, versionWanted, versionFound)
	}

	return nil
}
<commit_msg>vmware/common: detect Vmware 'unknown error' and show better message<commit_after>package common

import (
	"bytes"
	"fmt"
	"log"
	"os/exec"
	"regexp"
	"runtime"
	"strconv"
	"strings"

	"github.com/mitchellh/multistep"
)

// A driver is able to talk to VMware, control virtual machines, etc.
type Driver interface {
	// Clone clones the VMX and the disk to the destination path. The
	// destination is a path to the VMX file. The disk will be copied
	// to that same directory.
	Clone(dst string, src string) error

	// CompactDisk compacts a virtual disk.
	CompactDisk(string) error

	// CreateDisk creates a virtual disk with the given size.
	CreateDisk(string, string, string) error

	// Checks if the VMX file at the given path is running.
	IsRunning(string) (bool, error)

	// SSHAddress returns the SSH address for the VM that is being
	// managed by this driver.
	SSHAddress(multistep.StateBag) (string, error)

	// Start starts a VM specified by the path to the VMX given.
	Start(string, bool) error

	// Stop stops a VM specified by the path to the VMX given.
	Stop(string) error

	// SuppressMessages modifies the VMX or surrounding directory so that
	// VMware doesn't show any annoying messages.
	SuppressMessages(string) error

	// Get the path to the VMware ISO for the given flavor.
	ToolsIsoPath(string) string

	// Attach the VMware tools ISO
	ToolsInstall() error

	// Get the path to the DHCP leases file for the given device.
	DhcpLeasesPath(string) string

	// Verify checks to make sure that this driver should function
	// properly. This should check that all the files it will use
	// appear to exist and so on. If everything is okay, this doesn't
	// return an error. Otherwise, this returns an error.
	Verify() error
}

// NewDriver returns a new driver implementation for this operating
// system, or an error if the driver couldn't be initialized.
func NewDriver(dconfig *DriverConfig, config *SSHConfig) (Driver, error) {
	drivers := []Driver{}

	switch runtime.GOOS {
	case "darwin":
		drivers = []Driver{
			&Fusion6Driver{
				Fusion5Driver: Fusion5Driver{
					AppPath:   dconfig.FusionAppPath,
					SSHConfig: config,
				},
			},
			&Fusion5Driver{
				AppPath:   dconfig.FusionAppPath,
				SSHConfig: config,
			},
		}
	case "linux":
		fallthrough
	case "windows":
		drivers = []Driver{
			&Workstation10Driver{
				Workstation9Driver: Workstation9Driver{
					SSHConfig: config,
				},
			},
			&Workstation9Driver{
				SSHConfig: config,
			},
			&Player6Driver{
				Player5Driver: Player5Driver{
					SSHConfig: config,
				},
			},
			&Player5Driver{
				SSHConfig: config,
			},
		}
	default:
		return nil, fmt.Errorf("can't find driver for OS: %s", runtime.GOOS)
	}

	errs := ""
	for _, driver := range drivers {
		err := driver.Verify()
		if err == nil {
			return driver, nil
		}
		errs += "* " + err.Error() + "\n"
	}

	return nil, fmt.Errorf(
		"Unable to initialize any driver for this platform. The errors\n"+
			"from each driver are shown below. Please fix at least one driver\n"+
			"to continue:\n%s", errs)
}

func runAndLog(cmd *exec.Cmd) (string, string, error) {
	var stdout, stderr bytes.Buffer

	log.Printf("Executing: %s %v", cmd.Path, cmd.Args[1:])
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()

	stdoutString := strings.TrimSpace(stdout.String())
	stderrString := strings.TrimSpace(stderr.String())

	if _, ok := err.(*exec.ExitError); ok {
		message := stderrString
		if message == "" {
			message = stdoutString
		}

		err = fmt.Errorf("VMware error: %s", message)

		// If "unknown error" is in there, add some additional notes
		re := regexp.MustCompile(`(?i)unknown error`)
		if re.MatchString(message) {
			err = fmt.Errorf(
				"%s\n\n%s", err,
				"Packer detected a VMware 'Unknown Error'. Unfortunately VMware\n"+
					"often has extremely vague error messages such as this and Packer\n"+
					"itself can't do much about that. Please check the vmware.log files\n"+
					"created by VMware when a VM is started (in the directory of the\n"+
					"vmx file), which often contains more detailed error information.")
		}
	}

	log.Printf("stdout: %s", stdoutString)
	log.Printf("stderr: %s", stderrString)

	// Replace these for Windows, we only want to deal with Unix
	// style line endings.
	returnStdout := strings.Replace(stdout.String(), "\r\n", "\n", -1)
	returnStderr := strings.Replace(stderr.String(), "\r\n", "\n", -1)

	return returnStdout, returnStderr, err
}

func normalizeVersion(version string) (string, error) {
	i, err := strconv.Atoi(version)
	if err != nil {
		return "", fmt.Errorf(
			"VMware version '%s' is not numeric", version)
	}

	return fmt.Sprintf("%02d", i), nil
}

func compareVersions(versionFound string, versionWanted string, product string) error {
	found, err := normalizeVersion(versionFound)
	if err != nil {
		return err
	}

	wanted, err := normalizeVersion(versionWanted)
	if err != nil {
		return err
	}

	if found < wanted {
		return fmt.Errorf(
			"VMware %s version %s, or greater, is required. Found version: %s", product, versionWanted, versionFound)
	}

	return nil
}
<|endoftext|>


<commit_before>package handlers

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"testing"

	"github.com/wellington/spritewell"
	cx "github.com/wellington/wellington/context"
)

func init() {
	os.MkdirAll("../test/build/img", 0777)
}

func wrapCallback(sc cx.SassCallback, ch chan cx.UnionSassValue) cx.SassCallback {
	return func(c *cx.Context, usv cx.UnionSassValue) cx.UnionSassValue {
		usv = sc(c, usv)
		ch <- usv
		return usv
	}
}

func testSprite(ctx *cx.Context) {
	// Generate test sprite
	imgs := spritewell.ImageList{
		ImageDir:  ctx.ImageDir,
		BuildDir:  ctx.BuildDir,
		GenImgDir: ctx.GenImgDir,
	}
	glob := "*.png"
	err := imgs.Decode(glob)
	if err != nil {
		panic(err)
	}
	_, err = imgs.Combine()
	if err != nil {
		panic(err)
	}
}

func setupCtx(r io.Reader, out io.Writer, cookies ...cx.Cookie) (*cx.Context, cx.UnionSassValue, error) {
	var usv cx.UnionSassValue
	ctx := cx.NewContext()
	ctx.OutputStyle = cx.NESTED_STYLE
	ctx.IncludePaths = make([]string, 0)
	ctx.BuildDir = "../test/build"
	ctx.ImageDir = "../test/img"
	ctx.FontDir = "../test/font"
	ctx.GenImgDir = "../test/build/img"
	ctx.Out = ""

	testSprite(ctx)
	cc := make(chan cx.UnionSassValue, len(cookies))
	// If callbacks were made, add them to the context
	// and create channels for communicating with them.
	if len(cookies) > 0 {
		cs := make([]cx.Cookie, len(cookies))
		for i, c := range cookies {
			cs[i] = cx.Cookie{
				Sign: c.Sign,
				Fn:   wrapCallback(c.Fn, cc),
				Ctx:  ctx,
			}
		}
		usv = <-cc
	}
	err := ctx.Compile(r, out)
	return ctx, usv, err
}

func TestFuncImageURL(t *testing.T) {
	ctx := cx.Context{
		BuildDir: "test/build",
		ImageDir: "test/img",
	}

	usv, _ := cx.Marshal([]string{"image.png"})
	usv = ImageURL(&ctx, usv)
	var path string
	cx.Unmarshal(usv, &path)
	if e := "url('../img/image.png')"; e != path {
		t.Errorf("got: %s wanted: %s", path, e)
	}

	// Test sending invalid date to imageURL
	usv, _ = cx.Marshal(cx.SassNumber{Value: 1, Unit: "px"})
	_ = usv
	errusv := ImageURL(&ctx, usv)
	var s string
	merr := cx.Unmarshal(errusv, &s)
	if merr != nil {
		t.Error(merr)
	}

	e := "Sassvalue is type context.SassNumber and has value {1 px} but expected slice"

	if e != s {
		t.Errorf("got:\n%s\nwanted:\n%s", s, e)
	}

}

func TestFuncSpriteMap(t *testing.T) {
	ctx := cx.NewContext()
	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"

	// Add real arguments when sass lists can be [un]marshalled
	lst := []interface{}{"*.png", cx.SassNumber{Value: 5, Unit: "px"}}
	usv, _ := cx.Marshal(lst)
	usv = SpriteMap(ctx, usv)
	var path string
	err := cx.Unmarshal(usv, &path)
	if err != nil {
		t.Error(err)
	}

	if e := "*.png5"; e != path {
		t.Errorf("got: %s wanted: %s", path, e)
	}
}

func TestFuncSpriteFile(t *testing.T) {
	ctx := cx.NewContext()
	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"

	// Add real arguments when sass lists can be [un]marshalled
	lst := []interface{}{"*.png", "139"}
	usv, _ := cx.Marshal(lst)
	usv = SpriteFile(ctx, usv)
	var glob, path string
	err := cx.Unmarshal(usv, &glob, &path)
	if err != nil {
		t.Error(err)
	}

	if e := "*.png"; e != glob {
		t.Errorf("got: %s wanted: %s", e, glob)
	}

	if e := "139"; e != path {
		t.Errorf("got: %s wanted: %s", e, path)
	}

}

func TestCompileSpriteMap(t *testing.T) {
	in := bytes.NewBufferString(`
$aritymap: sprite-map("*.png", 0px); // Optional arguments
$map: sprite-map("*.png"); // One argument
$paddedmap: sprite-map("*.png", 1px); // One argument
div {
width: $map;
height: $aritymap;
line-height: $paddedmap;
}`)

	ctx := cx.NewContext()

	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"
	var out bytes.Buffer
	err := ctx.Compile(in, &out)
	if err != nil {
		t.Error(err)
	}
	exp := `div {
  width: *.png0;
  height: *.png0;
  line-height: *.png1; }
`

	if exp != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), exp)
	}
}

func TestCompileSpritePaddingMap(t *testing.T) {
	in := bytes.NewBufferString(`$map: sprite-map("dual/*.png",10px);
div {
  content: $map;
}`)

	ctx := cx.NewContext()

	ctx.ImageDir = "../test/img"
	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"

	var out bytes.Buffer
	err := ctx.Compile(in, &out)
	if err != nil {
		t.Error(err)
	}
	exp := `div {
  content: dual/*.png10; }
`
	if exp != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), exp)
	}
}

func TestFuncImageHeight(t *testing.T) {
	in := bytes.NewBufferString(`div {
    height: image-height("139");
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)

	if err != nil {
		t.Error(err)
	}

	e := `div {
  height: 139px; }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegImageWidth(t *testing.T) {
	in := bytes.NewBufferString(`div {
    height: image-width("139");
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  height: 96px; }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegSpriteImageHeight(t *testing.T) {
	in := bytes.NewBufferString(`$map: sprite-map("*.png");
div {
  height: image-height(sprite-file($map,"139"));
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  height: 139px; }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegSpriteImageWidth(t *testing.T) {
	in := bytes.NewBufferString(`$map: sprite-map("*.png");
div {
  width: image-width(sprite-file($map,"139"));
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  width: 96px; }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegImageURL(t *testing.T) {
	in := bytes.NewBufferString(`
div {
    background: image-url("139.png");
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  background: url('../img/139.png'); }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegInlineImage(t *testing.T) {
	in := bytes.NewBufferString(`
div {
    background: inline-image("pixel/1x1.png");
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  background: url('data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAMAAAAoyzS7AAAAA1BMVEX/TQBcNTh/AAAAAXRSTlMz/za5cAAAAA5JREFUeJxiYgAEAAD//wAGAAP60FmuAAAAAElFTkSuQmCC'); }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegInlineImageFail(t *testing.T) {
	var f *os.File
	old := os.Stdout
	os.Stdout = f
	defer func() { os.Stdout = old }()
	in := bytes.NewBufferString(`
div {
    background: inline-image("image.svg");
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  background: inline-image: image.svg filetype .svg is not supported; }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestFontURLFail(t *testing.T) {
	r, w, _ := os.Pipe()
	old := os.Stdout
	defer func() { os.Stdout = old }()
	os.Stdout = w
	in := bytes.NewBufferString(`@font-face {
  src: font-url("arial.eot");
}`)
	var out bytes.Buffer
	ctx := cx.Context{}
	err := ctx.Compile(in, &out)

	if err != nil {
		t.Error(err)
	}

	outC := make(chan string)
	go func(r *os.File) {
		var buf bytes.Buffer
		io.Copy(&buf, r)
		outC <- buf.String()
	}(r)

	w.Close()
	stdout := <-outC

	if e := "font-url: font path not set\n"; e != stdout {
		t.Errorf("got:\n%s\nwanted:\n%s\n", stdout, e)
	}

}

func ExampleFontURL() {
	in := bytes.NewBufferString(`
$path: font-url("arial.eot", true);
@font-face {
  src: font-url("arial.eot");
  src: url("#{$path}");
}`)

	_, _, err := setupCtx(in, os.Stdout)
	if err != nil {
		fmt.Println(err)
	}

	// Output:
	// @font-face {
	//   src: url("../font/arial.eot");
	//   src: url("../font/arial.eot"); }
}

func ExampleSprite() {
	in := bytes.NewBufferString(`
$map: sprite-map("dual/*.png", 10px); // One argument
div {
  background: sprite($map, "140");
}`)

	ctx := cx.NewContext()

	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"
	var out bytes.Buffer
	err := ctx.Compile(in, &out)
	if err != nil {
		fmt.Println(err)
	}

	io.Copy(os.Stdout, &out)

	// Output:
	// div {
	//   background: url("img/img-b798ab.png") -0px -149px; }

}

func TestSprite(t *testing.T) {
	in := bytes.NewBufferString(`
$map: sprite-map("dual/*.png", 10px);
div {
  background: sprite($map, "140", 0, 0);
}`)

	ctx := cx.NewContext()

	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"
	var out bytes.Buffer
	err := ctx.Compile(in, &out)

	e := `Error > stdin:4
error in C function sprite: Please specify unit for offset ie. (2px)
Backtrace:
	stdin:4, in function ` + "`sprite`" + `
	stdin:4

$map: sprite-map("dual/*.png", 10px);
div {
  background: sprite($map, "140", 0, 0);
}
`
	if e != err.Error() {
		t.Errorf("got:\n~%s~\nwanted:\n~%s~\n", err.Error(), e)
	}

}

func BenchmarkSprite(b *testing.B) {
	ctx := cx.NewContext()
	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"
	// Add real arguments when sass lists can be [un]marshalled
	lst := []interface{}{"*.png", cx.SassNumber{Value: 5, Unit: "px"}}
	usv, _ := cx.Marshal(lst)

	for i := 0; i < b.N; i++ {
		usv = SpriteMap(ctx, usv)
	}
	// Debug if needed
	// var s string
	// Unmarshal(usv, &s)
	// fmt.Println(s)
}
<commit_msg>Update test for new spritewell<commit_after>package handlers

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"testing"

	"github.com/wellington/spritewell"
	cx "github.com/wellington/wellington/context"
)

func init() {
	os.MkdirAll("../test/build/img", 0777)
}

func wrapCallback(sc cx.SassCallback, ch chan cx.UnionSassValue) cx.SassCallback {
	return func(c *cx.Context, usv cx.UnionSassValue) cx.UnionSassValue {
		usv = sc(c, usv)
		ch <- usv
		return usv
	}
}

func testSprite(ctx *cx.Context) {
	// Generate test sprite
	imgs := spritewell.ImageList{
		ImageDir:  ctx.ImageDir,
		BuildDir:  ctx.BuildDir,
		GenImgDir: ctx.GenImgDir,
	}
	glob := "*.png"
	err := imgs.Decode(glob)
	if err != nil {
		panic(err)
	}
	_, err = imgs.Combine()
	if err != nil {
		panic(err)
	}
}

func setupCtx(r io.Reader, out io.Writer, cookies ...cx.Cookie) (*cx.Context, cx.UnionSassValue, error) {
	var usv cx.UnionSassValue
	ctx := cx.NewContext()
	ctx.OutputStyle = cx.NESTED_STYLE
	ctx.IncludePaths = make([]string, 0)
	ctx.BuildDir = "../test/build"
	ctx.ImageDir = "../test/img"
	ctx.FontDir = "../test/font"
	ctx.GenImgDir = "../test/build/img"
	ctx.Out = ""

	testSprite(ctx)
	cc := make(chan cx.UnionSassValue, len(cookies))
	// If callbacks were made, add them to the context
	// and create channels for communicating with them.
	if len(cookies) > 0 {
		cs := make([]cx.Cookie, len(cookies))
		for i, c := range cookies {
			cs[i] = cx.Cookie{
				Sign: c.Sign,
				Fn:   wrapCallback(c.Fn, cc),
				Ctx:  ctx,
			}
		}
		usv = <-cc
	}
	err := ctx.Compile(r, out)
	return ctx, usv, err
}

func TestFuncImageURL(t *testing.T) {
	ctx := cx.Context{
		BuildDir: "test/build",
		ImageDir: "test/img",
	}

	usv, _ := cx.Marshal([]string{"image.png"})
	usv = ImageURL(&ctx, usv)
	var path string
	cx.Unmarshal(usv, &path)
	if e := "url('../img/image.png')"; e != path {
		t.Errorf("got: %s wanted: %s", path, e)
	}

	// Test sending invalid date to imageURL
	usv, _ = cx.Marshal(cx.SassNumber{Value: 1, Unit: "px"})
	_ = usv
	errusv := ImageURL(&ctx, usv)
	var s string
	merr := cx.Unmarshal(errusv, &s)
	if merr != nil {
		t.Error(merr)
	}

	e := "Sassvalue is type context.SassNumber and has value {1 px} but expected slice"

	if e != s {
		t.Errorf("got:\n%s\nwanted:\n%s", s, e)
	}

}

func TestFuncSpriteMap(t *testing.T) {
	ctx := cx.NewContext()
	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"

	// Add real arguments when sass lists can be [un]marshalled
	lst := []interface{}{"*.png", cx.SassNumber{Value: 5, Unit: "px"}}
	usv, _ := cx.Marshal(lst)
	usv = SpriteMap(ctx, usv)
	var path string
	err := cx.Unmarshal(usv, &path)
	if err != nil {
		t.Error(err)
	}

	if e := "*.png5"; e != path {
		t.Errorf("got: %s wanted: %s", path, e)
	}
}

func TestFuncSpriteFile(t *testing.T) {
	ctx := cx.NewContext()
	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"

	// Add real arguments when sass lists can be [un]marshalled
	lst := []interface{}{"*.png", "139"}
	usv, _ := cx.Marshal(lst)
	usv = SpriteFile(ctx, usv)
	var glob, path string
	err := cx.Unmarshal(usv, &glob, &path)
	if err != nil {
		t.Error(err)
	}

	if e := "*.png"; e != glob {
		t.Errorf("got: %s wanted: %s", e, glob)
	}

	if e := "139"; e != path {
		t.Errorf("got: %s wanted: %s", e, path)
	}

}

func TestCompileSpriteMap(t *testing.T) {
	in := bytes.NewBufferString(`
$aritymap: sprite-map("*.png", 0px); // Optional arguments
$map: sprite-map("*.png"); // One argument
$paddedmap: sprite-map("*.png", 1px); // One argument
div {
width: $map;
height: $aritymap;
line-height: $paddedmap;
}`)

	ctx := cx.NewContext()

	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"
	var out bytes.Buffer
	err := ctx.Compile(in, &out)
	if err != nil {
		t.Error(err)
	}
	exp := `div {
  width: *.png0;
  height: *.png0;
  line-height: *.png1; }
`

	if exp != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), exp)
	}
}

func TestCompileSpritePaddingMap(t *testing.T) {
	in := bytes.NewBufferString(`$map: sprite-map("dual/*.png",10px);
div {
  content: $map;
}`)

	ctx := cx.NewContext()

	ctx.ImageDir = "../test/img"
	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"

	var out bytes.Buffer
	err := ctx.Compile(in, &out)
	if err != nil {
		t.Error(err)
	}
	exp := `div {
  content: dual/*.png10; }
`
	if exp != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), exp)
	}
}

func TestFuncImageHeight(t *testing.T) {
	in := bytes.NewBufferString(`div {
    height: image-height("139");
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)

	if err != nil {
		t.Error(err)
	}

	e := `div {
  height: 139px; }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegImageWidth(t *testing.T) {
	in := bytes.NewBufferString(`div {
    height: image-width("139");
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  height: 96px; }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegSpriteImageHeight(t *testing.T) {
	in := bytes.NewBufferString(`$map: sprite-map("*.png");
div {
  height: image-height(sprite-file($map,"139"));
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  height: 139px; }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegSpriteImageWidth(t *testing.T) {
	in := bytes.NewBufferString(`$map: sprite-map("*.png");
div {
  width: image-width(sprite-file($map,"139"));
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  width: 96px; }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegImageURL(t *testing.T) {
	in := bytes.NewBufferString(`
div {
    background: image-url("139.png");
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  background: url('../img/139.png'); }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegInlineImage(t *testing.T) {
	in := bytes.NewBufferString(`
div {
    background: inline-image("pixel/1x1.png");
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  background: url('data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAMAAAAoyzS7AAAAA1BMVEX/TQBcNTh/AAAAAXRSTlMz/za5cAAAAA5JREFUeJxiYgAEAAD//wAGAAP60FmuAAAAAElFTkSuQmCC'); }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestRegInlineImageFail(t *testing.T) {
	var f *os.File
	old := os.Stdout
	os.Stdout = f
	defer func() { os.Stdout = old }()
	in := bytes.NewBufferString(`
div {
    background: inline-image("image.svg");
}`)
	var out bytes.Buffer
	_, _, err := setupCtx(in, &out)
	if err != nil {
		t.Error(err)
	}
	e := `div {
  background: inline-image: image.svg filetype .svg is not supported; }
`
	if e != out.String() {
		t.Errorf("got:\n%s\nwanted:\n%s", out.String(), e)
	}
}

func TestFontURLFail(t *testing.T) {
	r, w, _ := os.Pipe()
	old := os.Stdout
	defer func() { os.Stdout = old }()
	os.Stdout = w
	in := bytes.NewBufferString(`@font-face {
  src: font-url("arial.eot");
}`)
	var out bytes.Buffer
	ctx := cx.Context{}
	err := ctx.Compile(in, &out)

	if err != nil {
		t.Error(err)
	}

	outC := make(chan string)
	go func(r *os.File) {
		var buf bytes.Buffer
		io.Copy(&buf, r)
		outC <- buf.String()
	}(r)

	w.Close()
	stdout := <-outC

	if e := "font-url: font path not set\n"; e != stdout {
		t.Errorf("got:\n%s\nwanted:\n%s\n", stdout, e)
	}

}

func ExampleFontURL() {
	in := bytes.NewBufferString(`
$path: font-url("arial.eot", true);
@font-face {
  src: font-url("arial.eot");
  src: url("#{$path}");
}`)

	_, _, err := setupCtx(in, os.Stdout)
	if err != nil {
		fmt.Println(err)
	}

	// Output:
	// @font-face {
	//   src: url("../font/arial.eot");
	//   src: url("../font/arial.eot"); }
}

func ExampleSprite() {
	in := bytes.NewBufferString(`
$map: sprite-map("dual/*.png", 10px); // One argument
div {
  background: sprite($map, "140");
}`)

	ctx := cx.NewContext()

	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"
	var out bytes.Buffer
	err := ctx.Compile(in, &out)
	if err != nil {
		fmt.Println(err)
	}

	io.Copy(os.Stdout, &out)

	// Output:
	// div {
	//   background: url("img/b798ab.png") -0px -149px; }

}

func TestSprite(t *testing.T) {
	in := bytes.NewBufferString(`
$map: sprite-map("dual/*.png", 10px);
div {
  background: sprite($map, "140", 0, 0);
}`)

	ctx := cx.NewContext()

	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"
	var out bytes.Buffer
	err := ctx.Compile(in, &out)

	e := `Error > stdin:4
error in C function sprite: Please specify unit for offset ie. (2px)
Backtrace:
	stdin:4, in function ` + "`sprite`" + `
	stdin:4

$map: sprite-map("dual/*.png", 10px);
div {
  background: sprite($map, "140", 0, 0);
}
`
	if e != err.Error() {
		t.Errorf("got:\n~%s~\nwanted:\n~%s~\n", err.Error(), e)
	}

}

func BenchmarkSprite(b *testing.B) {
	ctx := cx.NewContext()
	ctx.BuildDir = "../test/build"
	ctx.GenImgDir = "../test/build/img"
	ctx.ImageDir = "../test/img"
	// Add real arguments when sass lists can be [un]marshalled
	lst := []interface{}{"*.png", cx.SassNumber{Value: 5, Unit: "px"}}
	usv, _ := cx.Marshal(lst)

	for i := 0; i < b.N; i++ {
		usv = SpriteMap(ctx, usv)
	}
	// Debug if needed
	// var s string
	// Unmarshal(usv, &s)
	// fmt.Println(s)
}
<|endoftext|>


<commit_before>// This file is adapted from code in the net/http/httputil
// package of the Go standard library, which is by the
// Go Authors, and bears this copyright and license info:
//
//   Copyright 2011 The Go Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style
//   license that can be found in the LICENSE file.
//
// This file has been modified from the standard lib to
// meet the needs of the application.

package proxy

import (
	"crypto/tls"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

var bufferPool = sync.Pool{New: createBuffer}

func createBuffer() interface{} {
	return make([]byte, 32*1024)
}

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// Director must be a function which modifies
	// the request into a new request to be sent
	// using Transport. Its response is then copied
	// back to the original client unmodified.
	Director func(*http.Request)

	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

// Though the relevant directive prefix is just "unix:", url.Parse
// will - assuming the regular URL scheme - add additional slashes
// as if "unix" was a request protocol.
// What we need is just the path, so if "unix:/var/run/www.socket"
// was the proxy directive, the parsed hostName would be
// "unix:///var/run/www.socket", hence the ambiguous trimming.
func socketDial(hostName string) func(network, addr string) (conn net.Conn, err error) {
	return func(network, addr string) (conn net.Conn, err error) {
		return net.Dial("unix", hostName[len("unix://"):])
	}
}

// NewSingleHostReverseProxy returns a new ReverseProxy that rewrites
// URLs to the scheme, host, and base path provided in target. If the
// target's path is "/base" and the incoming request was for "/dir",
// the target request will be for /base/dir.
// Without logic: target's path is "/", incoming is "/api/messages",
// without is "/api", then the target request will be for /messages.
func NewSingleHostReverseProxy(target *url.URL, without string, keepalive int) *ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		if target.Scheme == "unix" {
			// to make Dial work with unix URL,
			// scheme and host have to be faked
			req.URL.Scheme = "http"
			req.URL.Host = "socket"
		} else {
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
		}
		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
		// Trims the path of the socket from the URL path.
		// This is done because req.URL passed to your proxied service
		// will have the full path of the socket file prefixed to it.
		// Calling /test on a server that proxies requests to
		// unix:/var/run/www.socket will thus set the requested path
		// to /var/run/www.socket/test, rendering paths useless.
		if target.Scheme == "unix" {
			// See comment on socketDial for the trim
			socketPrefix := target.String()[len("unix://"):]
			req.URL.Path = strings.TrimPrefix(req.URL.Path, socketPrefix)
		}
		// We are then safe to remove the `without` prefix.
		if without != "" {
			req.URL.Path = strings.TrimPrefix(req.URL.Path, without)
		}
	}
	rp := &ReverseProxy{Director: director, FlushInterval: 250 * time.Millisecond} // flushing good for streaming & server-sent events
	if target.Scheme == "unix" {
		rp.Transport = &http.Transport{
			Dial: socketDial(target.String()),
		}
	} else if keepalive != http.DefaultMaxIdleConnsPerHost {
		// if keepalive is equal to the default,
		// just use default transport, to avoid creating
		// a brand new transport
		rp.Transport = &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			Dial: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
			}).Dial,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		}
		if keepalive == 0 {
			rp.Transport.(*http.Transport).DisableKeepAlives = true
		} else {
			rp.Transport.(*http.Transport).MaxIdleConnsPerHost = keepalive
		}
	}
	return rp
}

// UseInsecureTransport is used to facilitate HTTPS proxying
// when it is OK for upstream to be using a bad certificate,
// since this transport skips verification.
func (rp *ReverseProxy) UseInsecureTransport() {
	if rp.Transport == nil {
		rp.Transport = &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			Dial: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 10 * time.Second,
			TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
		}
	} else if transport, ok := rp.Transport.(*http.Transport); ok {
		transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
	}
}

// ServeHTTP serves the proxied request to the upstream by performing a roundtrip.
// It is designed to handle websocket connection upgrades as well.
func (rp *ReverseProxy) ServeHTTP(rw http.ResponseWriter, outreq *http.Request, respUpdateFn respUpdateFn) error {
	transport := rp.Transport
	if requestIsWebsocket(outreq) {
		transport = newConnHijackerTransport(transport)
	} else if transport == nil {
		transport = http.DefaultTransport
	}

	rp.Director(outreq)
	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		return err
	}

	if respUpdateFn != nil {
		respUpdateFn(res)
	}
	if res.StatusCode == http.StatusSwitchingProtocols && strings.ToLower(res.Header.Get("Upgrade")) == "websocket" {
		res.Body.Close()
		hj, ok := rw.(http.Hijacker)
		if !ok {
			return nil
		}

		conn, _, err := hj.Hijack()
		if err != nil {
			return err
		}
		defer conn.Close()

		var backendConn net.Conn
		if hj, ok := transport.(*connHijackerTransport); ok {
			backendConn = hj.Conn
			if _, err := conn.Write(hj.Replay); err != nil {
				return err
			}
			bufferPool.Put(hj.Replay)
		} else {
			backendConn, err = net.Dial("tcp", outreq.URL.Host)
			if err != nil {
				return err
			}
			outreq.Write(backendConn)
		}
		defer backendConn.Close()

		go func() {
			io.Copy(backendConn, conn) // write tcp stream to backend.
		}()
		io.Copy(conn, backendConn) // read tcp stream from backend.
	} else {
		defer res.Body.Close()
		for _, h := range hopHeaders {
			res.Header.Del(h)
		}
		copyHeader(rw.Header(), res.Header)
		rw.WriteHeader(res.StatusCode)
		rp.copyResponse(rw, res.Body)
	}

	return nil
}

func (rp *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	buf := bufferPool.Get()
	defer bufferPool.Put(buf)

	if rp.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: rp.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}
	io.CopyBuffer(dst, src, buf.([]byte))
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

type respUpdateFn func(resp *http.Response)

type hijackedConn struct {
	net.Conn
	hj *connHijackerTransport
}

func (c *hijackedConn) Read(b []byte) (n int, err error) {
	n, err = c.Conn.Read(b)
	c.hj.Replay = append(c.hj.Replay, b[:n]...)
	return
}

func (c *hijackedConn) Close() error {
	return nil
}

type connHijackerTransport struct {
	*http.Transport
	Conn   net.Conn
	Replay []byte
}

func newConnHijackerTransport(base http.RoundTripper) *connHijackerTransport {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial,
		TLSHandshakeTimeout: 10 * time.Second,
		MaxIdleConnsPerHost: -1,
	}
	if base != nil {
		if baseTransport, ok := base.(*http.Transport); ok {
			transport.Proxy = baseTransport.Proxy
			transport.TLSClientConfig = baseTransport.TLSClientConfig
			transport.TLSHandshakeTimeout = baseTransport.TLSHandshakeTimeout
			transport.Dial = baseTransport.Dial
			transport.DialTLS = baseTransport.DialTLS
			transport.MaxIdleConnsPerHost = -1
		}
	}
	hjTransport := &connHijackerTransport{transport, nil, bufferPool.Get().([]byte)[:0]}
	oldDial := transport.Dial
	oldDialTLS := transport.DialTLS
	if oldDial == nil {
		oldDial = (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial
	}
	hjTransport.Dial = func(network, addr string) (net.Conn, error) {
		c, err := oldDial(network, addr)
		hjTransport.Conn = c
		return &hijackedConn{c, hjTransport}, err
	}
	if oldDialTLS != nil {
		hjTransport.DialTLS = func(network, addr string) (net.Conn, error) {
			c, err := oldDialTLS(network, addr)
			hjTransport.Conn = c
			return &hijackedConn{c, hjTransport}, err
		}
	}
	return hjTransport
}

func requestIsWebsocket(req *http.Request) bool {
	return !(strings.ToLower(req.Header.Get("Upgrade")) != "websocket" || !strings.Contains(strings.ToLower(req.Header.Get("Connection")), "upgrade"))
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<commit_msg>Keep quic protocol headers only between one hop<commit_after>// This file is adapted from code in the net/http/httputil
// package of the Go standard library, which is by the
// Go Authors, and bears this copyright and license info:
//
//   Copyright 2011 The Go Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style
//   license that can be found in the LICENSE file.
//
// This file has been modified from the standard lib to
// meet the needs of the application.

package proxy

import (
	"crypto/tls"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

var bufferPool = sync.Pool{New: createBuffer}

func createBuffer() interface{} {
	return make([]byte, 32*1024)
}

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// Director must be a function which modifies
	// the request into a new request to be sent
	// using Transport. Its response is then copied
	// back to the original client unmodified.
	Director func(*http.Request)

	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

// Though the relevant directive prefix is just "unix:", url.Parse
// will - assuming the regular URL scheme - add additional slashes
// as if "unix" was a request protocol.
// What we need is just the path, so if "unix:/var/run/www.socket"
// was the proxy directive, the parsed hostName would be
// "unix:///var/run/www.socket", hence the ambiguous trimming.
func socketDial(hostName string) func(network, addr string) (conn net.Conn, err error) {
	return func(network, addr string) (conn net.Conn, err error) {
		return net.Dial("unix", hostName[len("unix://"):])
	}
}

// NewSingleHostReverseProxy returns a new ReverseProxy that rewrites
// URLs to the scheme, host, and base path provided in target. If the
// target's path is "/base" and the incoming request was for "/dir",
// the target request will be for /base/dir.
// Without logic: target's path is "/", incoming is "/api/messages",
// without is "/api", then the target request will be for /messages.
func NewSingleHostReverseProxy(target *url.URL, without string, keepalive int) *ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		if target.Scheme == "unix" {
			// to make Dial work with unix URL,
			// scheme and host have to be faked
			req.URL.Scheme = "http"
			req.URL.Host = "socket"
		} else {
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
		}
		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
		// Trims the path of the socket from the URL path.
		// This is done because req.URL passed to your proxied service
		// will have the full path of the socket file prefixed to it.
		// Calling /test on a server that proxies requests to
		// unix:/var/run/www.socket will thus set the requested path
		// to /var/run/www.socket/test, rendering paths useless.
		if target.Scheme == "unix" {
			// See comment on socketDial for the trim
			socketPrefix := target.String()[len("unix://"):]
			req.URL.Path = strings.TrimPrefix(req.URL.Path, socketPrefix)
		}
		// We are then safe to remove the `without` prefix.
		if without != "" {
			req.URL.Path = strings.TrimPrefix(req.URL.Path, without)
		}
	}
	rp := &ReverseProxy{Director: director, FlushInterval: 250 * time.Millisecond} // flushing good for streaming & server-sent events
	if target.Scheme == "unix" {
		rp.Transport = &http.Transport{
			Dial: socketDial(target.String()),
		}
	} else if keepalive != http.DefaultMaxIdleConnsPerHost {
		// if keepalive is equal to the default,
		// just use default transport, to avoid creating
		// a brand new transport
		rp.Transport = &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			Dial: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
			}).Dial,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		}
		if keepalive == 0 {
			rp.Transport.(*http.Transport).DisableKeepAlives = true
		} else {
			rp.Transport.(*http.Transport).MaxIdleConnsPerHost = keepalive
		}
	}
	return rp
}

// UseInsecureTransport is used to facilitate HTTPS proxying
// when it is OK for upstream to be using a bad certificate,
// since this transport skips verification.
func (rp *ReverseProxy) UseInsecureTransport() {
	if rp.Transport == nil {
		rp.Transport = &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			Dial: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 10 * time.Second,
			TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
		}
	} else if transport, ok := rp.Transport.(*http.Transport); ok {
		transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
	}
}

// ServeHTTP serves the proxied request to the upstream by performing a roundtrip.
// It is designed to handle websocket connection upgrades as well.
func (rp *ReverseProxy) ServeHTTP(rw http.ResponseWriter, outreq *http.Request, respUpdateFn respUpdateFn) error {
	transport := rp.Transport
	if requestIsWebsocket(outreq) {
		transport = newConnHijackerTransport(transport)
	} else if transport == nil {
		transport = http.DefaultTransport
	}

	rp.Director(outreq)
	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		return err
	}

	if respUpdateFn != nil {
		respUpdateFn(res)
	}
	if res.StatusCode == http.StatusSwitchingProtocols && strings.ToLower(res.Header.Get("Upgrade")) == "websocket" {
		res.Body.Close()
		hj, ok := rw.(http.Hijacker)
		if !ok {
			return nil
		}

		conn, _, err := hj.Hijack()
		if err != nil {
			return err
		}
		defer conn.Close()

		var backendConn net.Conn
		if hj, ok := transport.(*connHijackerTransport); ok {
			backendConn = hj.Conn
			if _, err := conn.Write(hj.Replay); err != nil {
				return err
			}
			bufferPool.Put(hj.Replay)
		} else {
			backendConn, err = net.Dial("tcp", outreq.URL.Host)
			if err != nil {
				return err
			}
			outreq.Write(backendConn)
		}
		defer backendConn.Close()

		go func() {
			io.Copy(backendConn, conn) // write tcp stream to backend.
		}()
		io.Copy(conn, backendConn) // read tcp stream from backend.
	} else {
		defer res.Body.Close()
		for _, h := range hopHeaders {
			res.Header.Del(h)
		}
		copyHeader(rw.Header(), res.Header)
		rw.WriteHeader(res.StatusCode)
		rp.copyResponse(rw, res.Body)
	}

	return nil
}

func (rp *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	buf := bufferPool.Get()
	defer bufferPool.Put(buf)

	if rp.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: rp.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}
	io.CopyBuffer(dst, src, buf.([]byte))
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
	"Alternate-Protocol",
	"Alt-Svc",
}

type respUpdateFn func(resp *http.Response)

type hijackedConn struct {
	net.Conn
	hj *connHijackerTransport
}

func (c *hijackedConn) Read(b []byte) (n int, err error) {
	n, err = c.Conn.Read(b)
	c.hj.Replay = append(c.hj.Replay, b[:n]...)
	return
}

func (c *hijackedConn) Close() error {
	return nil
}

type connHijackerTransport struct {
	*http.Transport
	Conn   net.Conn
	Replay []byte
}

func newConnHijackerTransport(base http.RoundTripper) *connHijackerTransport {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial,
		TLSHandshakeTimeout: 10 * time.Second,
		MaxIdleConnsPerHost: -1,
	}
	if base != nil {
		if baseTransport, ok := base.(*http.Transport); ok {
			transport.Proxy = baseTransport.Proxy
			transport.TLSClientConfig = baseTransport.TLSClientConfig
			transport.TLSHandshakeTimeout = baseTransport.TLSHandshakeTimeout
			transport.Dial = baseTransport.Dial
			transport.DialTLS = baseTransport.DialTLS
			transport.MaxIdleConnsPerHost = -1
		}
	}
	hjTransport := &connHijackerTransport{transport, nil, bufferPool.Get().([]byte)[:0]}
	oldDial := transport.Dial
	oldDialTLS := transport.DialTLS
	if oldDial == nil {
		oldDial = (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial
	}
	hjTransport.Dial = func(network, addr string) (net.Conn, error) {
		c, err := oldDial(network, addr)
		hjTransport.Conn = c
		return &hijackedConn{c, hjTransport}, err
	}
	if oldDialTLS != nil {
		hjTransport.DialTLS = func(network, addr string) (net.Conn, error) {
			c, err := oldDialTLS(network, addr)
			hjTransport.Conn = c
			return &hijackedConn{c, hjTransport}, err
		}
	}
	return hjTransport
}

func requestIsWebsocket(req *http.Request) bool {
	return !(strings.ToLower(req.Header.Get("Upgrade")) != "websocket" || !strings.Contains(strings.ToLower(req.Header.Get("Connection")), "upgrade"))
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<|endoftext|>


<commit_before>package db

// country.go - Country geolocation database lookup.

import (
	"a2sapi/src/constants"
	"a2sapi/src/logger"
	"a2sapi/src/models"
	"fmt"
	"net"
	"runtime"

	"github.com/oschwald/maxminddb-golang"
)

// This is an intermediate struct to represent the MaxMind DB format, not for JSON
type mmdbformat struct {
	Country struct {
		IsoCode string            `maxminddb:"iso_code"`
		Names   map[string]string `maxminddb:"names"`
	} `maxminddb:"country"`
	Continent struct {
		Names map[string]string `maxminddb:"names"`
	} `maxminddb:"continent"`
	Subdivisions []struct {
		IsoCode string `maxminddb:"iso_code"`
	} `maxminddb:"subdivisions"`
}

func getDefaultCountryData() *models.DbCountry {
	return &models.DbCountry{
		CountryName: "Unknown",
		CountryCode: "Unknown",
		Continent:   "Unknown",
		State:       "Unknown",
	}
}

// OpenCountryDB opens the country lookup database for reading. The caller of
// this function will be responsinble for calling a .Close() on the Reader pointer
// returned by this function.
func OpenCountryDB() (*maxminddb.Reader, error) {
	// Note: the caller of this function needs to handle db.Close()
	db, err := maxminddb.Open(constants.CountryDbFilePath)
	if err != nil {
		dir := "build_nix"
		if runtime.GOOS == "windows" {
			dir = "build_win"
		}
		logger.LogAppError(err)
		panic(
			fmt.Sprintf(
				`Unable to open country database! Use the get_countrydb script in the %s
directory to get the country DB file, or download it from:
http://geolite.maxmind.com/download/geoip/database/GeoLite2-City.mmdb.gz and
extract the "GeoLite2-City.mmdb" file into a directory called "db" in the same
directory as the a2sapi executable. Error: %s`, dir, err))
	}
	return db, nil
}

// GetCountryInfo attempts to retrieve the country information for a given IP,
// returning the result as a country model object over the corresponding result channel.
func GetCountryInfo(ch chan<- *models.DbCountry, db *maxminddb.Reader, ipstr string) {
	ip := net.ParseIP(ipstr)
	c := &mmdbformat{}
	err := db.Lookup(ip, c)
	if err != nil {
		ch <- getDefaultCountryData()
		return
	}

	countrydata := &models.DbCountry{
		CountryName: c.Country.Names["en"],
		CountryCode: c.Country.IsoCode,
		Continent:   c.Continent.Names["en"],
	}
	if c.Country.IsoCode == "US" {
		if len(c.Subdivisions) > 0 {
			countrydata.State = c.Subdivisions[0].IsoCode
		} else {
			countrydata.State = "Unknown"
		}
	} else {
		countrydata.State = "None"
	}
	ch <- countrydata
}
<commit_msg>Use default country information when country name or ISO code is empty<commit_after>package db

// country.go - Country geolocation database lookup.

import (
	"a2sapi/src/constants"
	"a2sapi/src/logger"
	"a2sapi/src/models"
	"fmt"
	"net"
	"runtime"

	"github.com/oschwald/maxminddb-golang"
)

// This is an intermediate struct to represent the MaxMind DB format, not for JSON
type mmdbformat struct {
	Country struct {
		IsoCode string            `maxminddb:"iso_code"`
		Names   map[string]string `maxminddb:"names"`
	} `maxminddb:"country"`
	Continent struct {
		Names map[string]string `maxminddb:"names"`
	} `maxminddb:"continent"`
	Subdivisions []struct {
		IsoCode string `maxminddb:"iso_code"`
	} `maxminddb:"subdivisions"`
}

func getDefaultCountryData() *models.DbCountry {
	return &models.DbCountry{
		CountryName: "Unknown",
		CountryCode: "Unknown",
		Continent:   "Unknown",
		State:       "Unknown",
	}
}

// OpenCountryDB opens the country lookup database for reading. The caller of
// this function will be responsinble for calling a .Close() on the Reader pointer
// returned by this function.
func OpenCountryDB() (*maxminddb.Reader, error) {
	// Note: the caller of this function needs to handle db.Close()
	db, err := maxminddb.Open(constants.CountryDbFilePath)
	if err != nil {
		dir := "build_nix"
		if runtime.GOOS == "windows" {
			dir = "build_win"
		}
		logger.LogAppError(err)
		panic(
			fmt.Sprintf(
				`Unable to open country database! Use the get_countrydb script in the %s
directory to get the country DB file, or download it from:
http://geolite.maxmind.com/download/geoip/database/GeoLite2-City.mmdb.gz and
extract the "GeoLite2-City.mmdb" file into a directory called "db" in the same
directory as the a2sapi executable. Error: %s`, dir, err))
	}
	return db, nil
}

// GetCountryInfo attempts to retrieve the country information for a given IP,
// returning the result as a country model object over the corresponding result channel.
func GetCountryInfo(ch chan<- *models.DbCountry, db *maxminddb.Reader, ipstr string) {
	ip := net.ParseIP(ipstr)
	c := &mmdbformat{}
	err := db.Lookup(ip, c)
	if err != nil {
		ch <- getDefaultCountryData()
		return
	}
	if c.Country.Names["en"] == "" || c.Country.IsoCode == "" {
		ch <- getDefaultCountryData()
		return
	}

	countrydata := &models.DbCountry{
		CountryName: c.Country.Names["en"],
		CountryCode: c.Country.IsoCode,
		Continent:   c.Continent.Names["en"],
	}
	if c.Country.IsoCode == "US" {
		if len(c.Subdivisions) > 0 {
			countrydata.State = c.Subdivisions[0].IsoCode
		} else {
			countrydata.State = "Unknown"
		}
	} else {
		countrydata.State = "None"
	}
	ch <- countrydata
}
<|endoftext|>