<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

// Package imagepolicy contains an admission controller that configures a webhook to which policy
// decisions are delegated.
package imagepolicy

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/golang/glog"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime/schema"
	kubeschema "k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/apiserver/pkg/admission"
	"k8s.io/apiserver/pkg/util/cache"
	"k8s.io/apiserver/pkg/util/webhook"
	"k8s.io/client-go/rest"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/apis/imagepolicy/v1alpha1"
	kubeapiserveradmission "k8s.io/kubernetes/pkg/kubeapiserver/admission"

	// install the clientgo image policy API for use with api registry
	_ "k8s.io/kubernetes/pkg/apis/imagepolicy/install"
)

var (
	groupVersions = []schema.GroupVersion{v1alpha1.SchemeGroupVersion}
)

func init() {
	kubeapiserveradmission.Plugins.Register("ImagePolicyWebhook", func(config io.Reader) (admission.Interface, error) {
		newImagePolicyWebhook, err := NewImagePolicyWebhook(config)
		if err != nil {
			return nil, err
		}
		return newImagePolicyWebhook, nil
	})
}

// imagePolicyWebhook is an implementation of admission.Interface.
type imagePolicyWebhook struct {
	*admission.Handler
	webhook       *webhook.GenericWebhook
	responseCache *cache.LRUExpireCache
	allowTTL      time.Duration
	denyTTL       time.Duration
	retryBackoff  time.Duration
	defaultAllow  bool
}

func (a *imagePolicyWebhook) statusTTL(status v1alpha1.ImageReviewStatus) time.Duration {
	if status.Allowed {
		return a.allowTTL
	}
	return a.denyTTL
}

// Filter out annotations that don't match *.image-policy.k8s.io/*
func (a *imagePolicyWebhook) filterAnnotations(allAnnotations map[string]string) map[string]string {
	annotations := make(map[string]string)
	for k, v := range allAnnotations {
		if strings.Contains(k, ".image-policy.k8s.io/") {
			annotations[k] = v
		}
	}
	return annotations
}

// Function to call on webhook failure; behavior determined by defaultAllow flag
func (a *imagePolicyWebhook) webhookError(attributes admission.Attributes, err error) error {
	if err != nil {
		glog.V(2).Infof("error contacting webhook backend: %s", err)
		if a.defaultAllow {
			glog.V(2).Infof("resource allowed in spite of webhook backend failure")
			return nil
		}
		glog.V(2).Infof("resource not allowed due to webhook backend failure ")
		return admission.NewForbidden(attributes, err)
	}
	return nil
}

func (a *imagePolicyWebhook) Admit(attributes admission.Attributes) (err error) {
	// Ignore all calls to subresources or resources other than pods.
	allowedResources := map[kubeschema.GroupResource]bool{
		api.Resource("pods"): true,
	}

	if len(attributes.GetSubresource()) != 0 || !allowedResources[attributes.GetResource().GroupResource()] {
		return nil
	}

	pod, ok := attributes.GetObject().(*api.Pod)
	if !ok {
		return apierrors.NewBadRequest("Resource was marked with kind Pod but was unable to be converted")
	}

	// Build list of ImageReviewContainerSpec
	var imageReviewContainerSpecs []v1alpha1.ImageReviewContainerSpec
	containers := make([]api.Container, 0, len(pod.Spec.Containers)+len(pod.Spec.InitContainers))
	containers = append(containers, pod.Spec.Containers...)
	containers = append(containers, pod.Spec.InitContainers...)
	for _, c := range containers {
		imageReviewContainerSpecs = append(imageReviewContainerSpecs, v1alpha1.ImageReviewContainerSpec{
			Image: c.Image,
		})
	}
	imageReview := v1alpha1.ImageReview{
		Spec: v1alpha1.ImageReviewSpec{
			Containers:  imageReviewContainerSpecs,
			Annotations: a.filterAnnotations(pod.Annotations),
			Namespace:   attributes.GetNamespace(),
		},
	}
	if err := a.admitPod(attributes, &imageReview); err != nil {
		return admission.NewForbidden(attributes, err)
	}
	return nil
}

func (a *imagePolicyWebhook) admitPod(attributes admission.Attributes, review *v1alpha1.ImageReview) error {
	cacheKey, err := json.Marshal(review.Spec)
	if err != nil {
		return err
	}
	if entry, ok := a.responseCache.Get(string(cacheKey)); ok {
		review.Status = entry.(v1alpha1.ImageReviewStatus)
	} else {
		result := a.webhook.WithExponentialBackoff(func() rest.Result {
			return a.webhook.RestClient.Post().Body(review).Do()
		})

		if err := result.Error(); err != nil {
			return a.webhookError(attributes, err)
		}
		var statusCode int
		if result.StatusCode(&statusCode); statusCode < 200 || statusCode >= 300 {
			return a.webhookError(attributes, fmt.Errorf("Error contacting webhook: %d", statusCode))
		}

		if err := result.Into(review); err != nil {
			return a.webhookError(attributes, err)
		}

		a.responseCache.Add(string(cacheKey), review.Status, a.statusTTL(review.Status))
	}

	if !review.Status.Allowed {
		if len(review.Status.Reason) > 0 {
			return fmt.Errorf("image policy webook backend denied one or more images: %s", review.Status.Reason)
		}
		return errors.New("one or more images rejected by webhook backend")
	}

	return nil
}

// NewImagePolicyWebhook a new imagePolicyWebhook from the provided config file.
// The config file is specified by --admission-controller-config-file and has the
// following format for a webhook:
//
//   {
//     "imagePolicy": {
//        "kubeConfigFile": "path/to/kubeconfig/for/backend",
//        "allowTTL": 30,           # time in s to cache approval
//        "denyTTL": 30,            # time in s to cache denial
//        "retryBackoff": 500,      # time in ms to wait between retries
//        "defaultAllow": true      # determines behavior if the webhook backend fails
//     }
//   }
//
// The config file may be json or yaml.
//
// The kubeconfig property refers to another file in the kubeconfig format which
// specifies how to connect to the webhook backend.
//
// The kubeconfig's cluster field is used to refer to the remote service, user refers to the returned authorizer.
//
//     # clusters refers to the remote service.
//     clusters:
//     - name: name-of-remote-imagepolicy-service
//       cluster:
//         certificate-authority: /path/to/ca.pem      # CA for verifying the remote service.
//         server: https://images.example.com/policy # URL of remote service to query. Must use 'https'.
//
//     # users refers to the API server's webhook configuration.
//     users:
//     - name: name-of-api-server
//       user:
//         client-certificate: /path/to/cert.pem # cert for the webhook plugin to use
//         client-key: /path/to/key.pem          # key matching the cert
//
// For additional HTTP configuration, refer to the kubeconfig documentation
// http://kubernetes.io/v1.1/docs/user-guide/kubeconfig-file.html.
func NewImagePolicyWebhook(configFile io.Reader) (admission.Interface, error) {
	// TODO: move this to a versioned configuration file format
	var config AdmissionConfig
	d := yaml.NewYAMLOrJSONDecoder(configFile, 4096)
	err := d.Decode(&config)
	if err != nil {
		return nil, err
	}

	whConfig := config.ImagePolicyWebhook
	if err := normalizeWebhookConfig(&whConfig); err != nil {
		return nil, err
	}

	gw, err := webhook.NewGenericWebhook(api.Registry, api.Codecs, whConfig.KubeConfigFile, groupVersions, whConfig.RetryBackoff)
	if err != nil {
		return nil, err
	}
	return &imagePolicyWebhook{
		Handler:       admission.NewHandler(admission.Create, admission.Update),
		webhook:       gw,
		responseCache: cache.NewLRUExpireCache(1024),
		allowTTL:      whConfig.AllowTTL,
		denyTTL:       whConfig.DenyTTL,
		defaultAllow:  whConfig.DefaultAllow,
	}, nil
}
<commit_msg>imagepolicy admission.go: typo fix<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

// Package imagepolicy contains an admission controller that configures a webhook to which policy
// decisions are delegated.
package imagepolicy

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/golang/glog"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime/schema"
	kubeschema "k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/apiserver/pkg/admission"
	"k8s.io/apiserver/pkg/util/cache"
	"k8s.io/apiserver/pkg/util/webhook"
	"k8s.io/client-go/rest"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/apis/imagepolicy/v1alpha1"
	kubeapiserveradmission "k8s.io/kubernetes/pkg/kubeapiserver/admission"

	// install the clientgo image policy API for use with api registry
	_ "k8s.io/kubernetes/pkg/apis/imagepolicy/install"
)

var (
	groupVersions = []schema.GroupVersion{v1alpha1.SchemeGroupVersion}
)

func init() {
	kubeapiserveradmission.Plugins.Register("ImagePolicyWebhook", func(config io.Reader) (admission.Interface, error) {
		newImagePolicyWebhook, err := NewImagePolicyWebhook(config)
		if err != nil {
			return nil, err
		}
		return newImagePolicyWebhook, nil
	})
}

// imagePolicyWebhook is an implementation of admission.Interface.
type imagePolicyWebhook struct {
	*admission.Handler
	webhook       *webhook.GenericWebhook
	responseCache *cache.LRUExpireCache
	allowTTL      time.Duration
	denyTTL       time.Duration
	retryBackoff  time.Duration
	defaultAllow  bool
}

func (a *imagePolicyWebhook) statusTTL(status v1alpha1.ImageReviewStatus) time.Duration {
	if status.Allowed {
		return a.allowTTL
	}
	return a.denyTTL
}

// Filter out annotations that don't match *.image-policy.k8s.io/*
func (a *imagePolicyWebhook) filterAnnotations(allAnnotations map[string]string) map[string]string {
	annotations := make(map[string]string)
	for k, v := range allAnnotations {
		if strings.Contains(k, ".image-policy.k8s.io/") {
			annotations[k] = v
		}
	}
	return annotations
}

// Function to call on webhook failure; behavior determined by defaultAllow flag
func (a *imagePolicyWebhook) webhookError(attributes admission.Attributes, err error) error {
	if err != nil {
		glog.V(2).Infof("error contacting webhook backend: %s", err)
		if a.defaultAllow {
			glog.V(2).Infof("resource allowed in spite of webhook backend failure")
			return nil
		}
		glog.V(2).Infof("resource not allowed due to webhook backend failure ")
		return admission.NewForbidden(attributes, err)
	}
	return nil
}

func (a *imagePolicyWebhook) Admit(attributes admission.Attributes) (err error) {
	// Ignore all calls to subresources or resources other than pods.
	allowedResources := map[kubeschema.GroupResource]bool{
		api.Resource("pods"): true,
	}

	if len(attributes.GetSubresource()) != 0 || !allowedResources[attributes.GetResource().GroupResource()] {
		return nil
	}

	pod, ok := attributes.GetObject().(*api.Pod)
	if !ok {
		return apierrors.NewBadRequest("Resource was marked with kind Pod but was unable to be converted")
	}

	// Build list of ImageReviewContainerSpec
	var imageReviewContainerSpecs []v1alpha1.ImageReviewContainerSpec
	containers := make([]api.Container, 0, len(pod.Spec.Containers)+len(pod.Spec.InitContainers))
	containers = append(containers, pod.Spec.Containers...)
	containers = append(containers, pod.Spec.InitContainers...)
	for _, c := range containers {
		imageReviewContainerSpecs = append(imageReviewContainerSpecs, v1alpha1.ImageReviewContainerSpec{
			Image: c.Image,
		})
	}
	imageReview := v1alpha1.ImageReview{
		Spec: v1alpha1.ImageReviewSpec{
			Containers:  imageReviewContainerSpecs,
			Annotations: a.filterAnnotations(pod.Annotations),
			Namespace:   attributes.GetNamespace(),
		},
	}
	if err := a.admitPod(attributes, &imageReview); err != nil {
		return admission.NewForbidden(attributes, err)
	}
	return nil
}

func (a *imagePolicyWebhook) admitPod(attributes admission.Attributes, review *v1alpha1.ImageReview) error {
	cacheKey, err := json.Marshal(review.Spec)
	if err != nil {
		return err
	}
	if entry, ok := a.responseCache.Get(string(cacheKey)); ok {
		review.Status = entry.(v1alpha1.ImageReviewStatus)
	} else {
		result := a.webhook.WithExponentialBackoff(func() rest.Result {
			return a.webhook.RestClient.Post().Body(review).Do()
		})

		if err := result.Error(); err != nil {
			return a.webhookError(attributes, err)
		}
		var statusCode int
		if result.StatusCode(&statusCode); statusCode < 200 || statusCode >= 300 {
			return a.webhookError(attributes, fmt.Errorf("Error contacting webhook: %d", statusCode))
		}

		if err := result.Into(review); err != nil {
			return a.webhookError(attributes, err)
		}

		a.responseCache.Add(string(cacheKey), review.Status, a.statusTTL(review.Status))
	}

	if !review.Status.Allowed {
		if len(review.Status.Reason) > 0 {
			return fmt.Errorf("image policy webhook backend denied one or more images: %s", review.Status.Reason)
		}
		return errors.New("one or more images rejected by webhook backend")
	}

	return nil
}

// NewImagePolicyWebhook a new imagePolicyWebhook from the provided config file.
// The config file is specified by --admission-controller-config-file and has the
// following format for a webhook:
//
//   {
//     "imagePolicy": {
//        "kubeConfigFile": "path/to/kubeconfig/for/backend",
//        "allowTTL": 30,           # time in s to cache approval
//        "denyTTL": 30,            # time in s to cache denial
//        "retryBackoff": 500,      # time in ms to wait between retries
//        "defaultAllow": true      # determines behavior if the webhook backend fails
//     }
//   }
//
// The config file may be json or yaml.
//
// The kubeconfig property refers to another file in the kubeconfig format which
// specifies how to connect to the webhook backend.
//
// The kubeconfig's cluster field is used to refer to the remote service, user refers to the returned authorizer.
//
//     # clusters refers to the remote service.
//     clusters:
//     - name: name-of-remote-imagepolicy-service
//       cluster:
//         certificate-authority: /path/to/ca.pem      # CA for verifying the remote service.
//         server: https://images.example.com/policy # URL of remote service to query. Must use 'https'.
//
//     # users refers to the API server's webhook configuration.
//     users:
//     - name: name-of-api-server
//       user:
//         client-certificate: /path/to/cert.pem # cert for the webhook plugin to use
//         client-key: /path/to/key.pem          # key matching the cert
//
// For additional HTTP configuration, refer to the kubeconfig documentation
// http://kubernetes.io/v1.1/docs/user-guide/kubeconfig-file.html.
func NewImagePolicyWebhook(configFile io.Reader) (admission.Interface, error) {
	// TODO: move this to a versioned configuration file format
	var config AdmissionConfig
	d := yaml.NewYAMLOrJSONDecoder(configFile, 4096)
	err := d.Decode(&config)
	if err != nil {
		return nil, err
	}

	whConfig := config.ImagePolicyWebhook
	if err := normalizeWebhookConfig(&whConfig); err != nil {
		return nil, err
	}

	gw, err := webhook.NewGenericWebhook(api.Registry, api.Codecs, whConfig.KubeConfigFile, groupVersions, whConfig.RetryBackoff)
	if err != nil {
		return nil, err
	}
	return &imagePolicyWebhook{
		Handler:       admission.NewHandler(admission.Create, admission.Update),
		webhook:       gw,
		responseCache: cache.NewLRUExpireCache(1024),
		allowTTL:      whConfig.AllowTTL,
		denyTTL:       whConfig.DenyTTL,
		defaultAllow:  whConfig.DefaultAllow,
	}, nil
}
<|endoftext|>


<commit_before>package system

import (
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
	"unsafe"

	winio "github.com/Microsoft/go-winio"
	"golang.org/x/sys/windows"
)

const (
	// SddlAdministratorsLocalSystem is local administrators plus NT AUTHORITY\System
	SddlAdministratorsLocalSystem = "D:P(A;OICI;GA;;;BA)(A;OICI;GA;;;SY)"
	// SddlNtvmAdministratorsLocalSystem is NT VIRTUAL MACHINE\Virtual Machines plus local administrators plus NT AUTHORITY\System
	SddlNtvmAdministratorsLocalSystem = "D:P(A;OICI;GA;;;S-1-5-83-0)(A;OICI;GA;;;BA)(A;OICI;GA;;;SY)"
)

// MkdirAllWithACL is a wrapper for MkdirAll that creates a directory
// with an appropriate SDDL defined ACL.
func MkdirAllWithACL(path string, perm os.FileMode, sddl string) error {
	return mkdirall(path, true, sddl)
}

// MkdirAll implementation that is volume path aware for Windows.
func MkdirAll(path string, _ os.FileMode, sddl string) error {
	return mkdirall(path, false, sddl)
}

// mkdirall is a custom version of os.MkdirAll modified for use on Windows
// so that it is both volume path aware, and can create a directory with
// a DACL.
func mkdirall(path string, applyACL bool, sddl string) error {
	if re := regexp.MustCompile(`^\\\\\?\\Volume{[a-z0-9-]+}$`); re.MatchString(path) {
		return nil
	}

	// The rest of this method is largely copied from os.MkdirAll and should be kept
	// as-is to ensure compatibility.

	// Fast path: if we can tell whether path is a directory or file, stop with success or error.
	dir, err := os.Stat(path)
	if err == nil {
		if dir.IsDir() {
			return nil
		}
		return &os.PathError{
			Op:   "mkdir",
			Path: path,
			Err:  syscall.ENOTDIR,
		}
	}

	// Slow path: make sure parent exists and then call Mkdir for path.
	i := len(path)
	for i > 0 && os.IsPathSeparator(path[i-1]) { // Skip trailing path separator.
		i--
	}

	j := i
	for j > 0 && !os.IsPathSeparator(path[j-1]) { // Scan backward over element.
		j--
	}

	if j > 1 {
		// Create parent
		err = mkdirall(path[0:j-1], false, sddl)
		if err != nil {
			return err
		}
	}

	// Parent now exists; invoke os.Mkdir or mkdirWithACL and use its result.
	if applyACL {
		err = mkdirWithACL(path, sddl)
	} else {
		err = os.Mkdir(path, 0)
	}

	if err != nil {
		// Handle arguments like "foo/." by
		// double-checking that directory doesn't exist.
		dir, err1 := os.Lstat(path)
		if err1 == nil && dir.IsDir() {
			return nil
		}
		return err
	}
	return nil
}

// mkdirWithACL creates a new directory. If there is an error, it will be of
// type *PathError. .
//
// This is a modified and combined version of os.Mkdir and windows.Mkdir
// in golang to cater for creating a directory am ACL permitting full
// access, with inheritance, to any subfolder/file for Built-in Administrators
// and Local System.
func mkdirWithACL(name string, sddl string) error {
	sa := windows.SecurityAttributes{Length: 0}
	sd, err := winio.SddlToSecurityDescriptor(sddl)
	if err != nil {
		return &os.PathError{Op: "mkdir", Path: name, Err: err}
	}
	sa.Length = uint32(unsafe.Sizeof(sa))
	sa.InheritHandle = 1
	sa.SecurityDescriptor = uintptr(unsafe.Pointer(&sd[0]))

	namep, err := windows.UTF16PtrFromString(name)
	if err != nil {
		return &os.PathError{Op: "mkdir", Path: name, Err: err}
	}

	e := windows.CreateDirectory(namep, &sa)
	if e != nil {
		return &os.PathError{Op: "mkdir", Path: name, Err: e}
	}
	return nil
}

// IsAbs is a platform-specific wrapper for filepath.IsAbs. On Windows,
// golang filepath.IsAbs does not consider a path \windows\system32 as absolute
// as it doesn't start with a drive-letter/colon combination. However, in
// docker we need to verify things such as WORKDIR /windows/system32 in
// a Dockerfile (which gets translated to \windows\system32 when being processed
// by the daemon. This SHOULD be treated as absolute from a docker processing
// perspective.
func IsAbs(path string) bool {
	if !filepath.IsAbs(path) {
		if !strings.HasPrefix(path, string(os.PathSeparator)) {
			return false
		}
	}
	return true
}

// The origin of the functions below here are the golang OS and windows packages,
// slightly modified to only cope with files, not directories due to the
// specific use case.
//
// The alteration is to allow a file on Windows to be opened with
// FILE_FLAG_SEQUENTIAL_SCAN (particular for docker load), to avoid eating
// the standby list, particularly when accessing large files such as layer.tar.

// CreateSequential creates the named file with mode 0666 (before umask), truncating
// it if it already exists. If successful, methods on the returned
// File can be used for I/O; the associated file descriptor has mode
// O_RDWR.
// If there is an error, it will be of type *PathError.
func CreateSequential(name string) (*os.File, error) {
	return OpenFileSequential(name, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0)
}

// OpenSequential opens the named file for reading. If successful, methods on
// the returned file can be used for reading; the associated file
// descriptor has mode O_RDONLY.
// If there is an error, it will be of type *PathError.
func OpenSequential(name string) (*os.File, error) {
	return OpenFileSequential(name, os.O_RDONLY, 0)
}

// OpenFileSequential is the generalized open call; most users will use Open
// or Create instead.
// If there is an error, it will be of type *PathError.
func OpenFileSequential(name string, flag int, _ os.FileMode) (*os.File, error) {
	if name == "" {
		return nil, &os.PathError{Op: "open", Path: name, Err: syscall.ENOENT}
	}
	r, errf := windowsOpenFileSequential(name, flag, 0)
	if errf == nil {
		return r, nil
	}
	return nil, &os.PathError{Op: "open", Path: name, Err: errf}
}

func windowsOpenFileSequential(name string, flag int, _ os.FileMode) (file *os.File, err error) {
	r, e := windowsOpenSequential(name, flag|windows.O_CLOEXEC, 0)
	if e != nil {
		return nil, e
	}
	return os.NewFile(uintptr(r), name), nil
}

func makeInheritSa() *windows.SecurityAttributes {
	var sa windows.SecurityAttributes
	sa.Length = uint32(unsafe.Sizeof(sa))
	sa.InheritHandle = 1
	return &sa
}

func windowsOpenSequential(path string, mode int, _ uint32) (fd windows.Handle, err error) {
	if len(path) == 0 {
		return windows.InvalidHandle, windows.ERROR_FILE_NOT_FOUND
	}
	pathp, err := windows.UTF16PtrFromString(path)
	if err != nil {
		return windows.InvalidHandle, err
	}
	var access uint32
	switch mode & (windows.O_RDONLY | windows.O_WRONLY | windows.O_RDWR) {
	case windows.O_RDONLY:
		access = windows.GENERIC_READ
	case windows.O_WRONLY:
		access = windows.GENERIC_WRITE
	case windows.O_RDWR:
		access = windows.GENERIC_READ | windows.GENERIC_WRITE
	}
	if mode&windows.O_CREAT != 0 {
		access |= windows.GENERIC_WRITE
	}
	if mode&windows.O_APPEND != 0 {
		access &^= windows.GENERIC_WRITE
		access |= windows.FILE_APPEND_DATA
	}
	sharemode := uint32(windows.FILE_SHARE_READ | windows.FILE_SHARE_WRITE)
	var sa *windows.SecurityAttributes
	if mode&windows.O_CLOEXEC == 0 {
		sa = makeInheritSa()
	}
	var createmode uint32
	switch {
	case mode&(windows.O_CREAT|windows.O_EXCL) == (windows.O_CREAT | windows.O_EXCL):
		createmode = windows.CREATE_NEW
	case mode&(windows.O_CREAT|windows.O_TRUNC) == (windows.O_CREAT | windows.O_TRUNC):
		createmode = windows.CREATE_ALWAYS
	case mode&windows.O_CREAT == windows.O_CREAT:
		createmode = windows.OPEN_ALWAYS
	case mode&windows.O_TRUNC == windows.O_TRUNC:
		createmode = windows.TRUNCATE_EXISTING
	default:
		createmode = windows.OPEN_EXISTING
	}
	// Use FILE_FLAG_SEQUENTIAL_SCAN rather than FILE_ATTRIBUTE_NORMAL as implemented in golang.
	//https://msdn.microsoft.com/en-us/library/windows/desktop/aa363858(v=vs.85).aspx
	const fileFlagSequentialScan = 0x08000000 // FILE_FLAG_SEQUENTIAL_SCAN
	h, e := windows.CreateFile(pathp, access, sharemode, sa, createmode, fileFlagSequentialScan, 0)
	return h, e
}

// Helpers for TempFileSequential
var rand uint32
var randmu sync.Mutex

func reseed() uint32 {
	return uint32(time.Now().UnixNano() + int64(os.Getpid()))
}
func nextSuffix() string {
	randmu.Lock()
	r := rand
	if r == 0 {
		r = reseed()
	}
	r = r*1664525 + 1013904223 // constants from Numerical Recipes
	rand = r
	randmu.Unlock()
	return strconv.Itoa(int(1e9 + r%1e9))[1:]
}

// TempFileSequential is a copy of ioutil.TempFile, modified to use sequential
// file access. Below is the original comment from golang:
// TempFile creates a new temporary file in the directory dir
// with a name beginning with prefix, opens the file for reading
// and writing, and returns the resulting *os.File.
// If dir is the empty string, TempFile uses the default directory
// for temporary files (see os.TempDir).
// Multiple programs calling TempFile simultaneously
// will not choose the same file. The caller can use f.Name()
// to find the pathname of the file. It is the caller's responsibility
// to remove the file when no longer needed.
func TempFileSequential(dir, prefix string) (f *os.File, err error) {
	if dir == "" {
		dir = os.TempDir()
	}

	nconflict := 0
	for i := 0; i < 10000; i++ {
		name := filepath.Join(dir, prefix+nextSuffix())
		f, err = OpenFileSequential(name, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0600)
		if os.IsExist(err) {
			if nconflict++; nconflict > 10 {
				randmu.Lock()
				rand = reseed()
				randmu.Unlock()
			}
			continue
		}
		break
	}
	return
}
<commit_msg>Add canonical import comment<commit_after>package system // import "github.com/docker/docker/pkg/system"

import (
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
	"unsafe"

	winio "github.com/Microsoft/go-winio"
	"golang.org/x/sys/windows"
)

const (
	// SddlAdministratorsLocalSystem is local administrators plus NT AUTHORITY\System
	SddlAdministratorsLocalSystem = "D:P(A;OICI;GA;;;BA)(A;OICI;GA;;;SY)"
	// SddlNtvmAdministratorsLocalSystem is NT VIRTUAL MACHINE\Virtual Machines plus local administrators plus NT AUTHORITY\System
	SddlNtvmAdministratorsLocalSystem = "D:P(A;OICI;GA;;;S-1-5-83-0)(A;OICI;GA;;;BA)(A;OICI;GA;;;SY)"
)

// MkdirAllWithACL is a wrapper for MkdirAll that creates a directory
// with an appropriate SDDL defined ACL.
func MkdirAllWithACL(path string, perm os.FileMode, sddl string) error {
	return mkdirall(path, true, sddl)
}

// MkdirAll implementation that is volume path aware for Windows.
func MkdirAll(path string, _ os.FileMode, sddl string) error {
	return mkdirall(path, false, sddl)
}

// mkdirall is a custom version of os.MkdirAll modified for use on Windows
// so that it is both volume path aware, and can create a directory with
// a DACL.
func mkdirall(path string, applyACL bool, sddl string) error {
	if re := regexp.MustCompile(`^\\\\\?\\Volume{[a-z0-9-]+}$`); re.MatchString(path) {
		return nil
	}

	// The rest of this method is largely copied from os.MkdirAll and should be kept
	// as-is to ensure compatibility.

	// Fast path: if we can tell whether path is a directory or file, stop with success or error.
	dir, err := os.Stat(path)
	if err == nil {
		if dir.IsDir() {
			return nil
		}
		return &os.PathError{
			Op:   "mkdir",
			Path: path,
			Err:  syscall.ENOTDIR,
		}
	}

	// Slow path: make sure parent exists and then call Mkdir for path.
	i := len(path)
	for i > 0 && os.IsPathSeparator(path[i-1]) { // Skip trailing path separator.
		i--
	}

	j := i
	for j > 0 && !os.IsPathSeparator(path[j-1]) { // Scan backward over element.
		j--
	}

	if j > 1 {
		// Create parent
		err = mkdirall(path[0:j-1], false, sddl)
		if err != nil {
			return err
		}
	}

	// Parent now exists; invoke os.Mkdir or mkdirWithACL and use its result.
	if applyACL {
		err = mkdirWithACL(path, sddl)
	} else {
		err = os.Mkdir(path, 0)
	}

	if err != nil {
		// Handle arguments like "foo/." by
		// double-checking that directory doesn't exist.
		dir, err1 := os.Lstat(path)
		if err1 == nil && dir.IsDir() {
			return nil
		}
		return err
	}
	return nil
}

// mkdirWithACL creates a new directory. If there is an error, it will be of
// type *PathError. .
//
// This is a modified and combined version of os.Mkdir and windows.Mkdir
// in golang to cater for creating a directory am ACL permitting full
// access, with inheritance, to any subfolder/file for Built-in Administrators
// and Local System.
func mkdirWithACL(name string, sddl string) error {
	sa := windows.SecurityAttributes{Length: 0}
	sd, err := winio.SddlToSecurityDescriptor(sddl)
	if err != nil {
		return &os.PathError{Op: "mkdir", Path: name, Err: err}
	}
	sa.Length = uint32(unsafe.Sizeof(sa))
	sa.InheritHandle = 1
	sa.SecurityDescriptor = uintptr(unsafe.Pointer(&sd[0]))

	namep, err := windows.UTF16PtrFromString(name)
	if err != nil {
		return &os.PathError{Op: "mkdir", Path: name, Err: err}
	}

	e := windows.CreateDirectory(namep, &sa)
	if e != nil {
		return &os.PathError{Op: "mkdir", Path: name, Err: e}
	}
	return nil
}

// IsAbs is a platform-specific wrapper for filepath.IsAbs. On Windows,
// golang filepath.IsAbs does not consider a path \windows\system32 as absolute
// as it doesn't start with a drive-letter/colon combination. However, in
// docker we need to verify things such as WORKDIR /windows/system32 in
// a Dockerfile (which gets translated to \windows\system32 when being processed
// by the daemon. This SHOULD be treated as absolute from a docker processing
// perspective.
func IsAbs(path string) bool {
	if !filepath.IsAbs(path) {
		if !strings.HasPrefix(path, string(os.PathSeparator)) {
			return false
		}
	}
	return true
}

// The origin of the functions below here are the golang OS and windows packages,
// slightly modified to only cope with files, not directories due to the
// specific use case.
//
// The alteration is to allow a file on Windows to be opened with
// FILE_FLAG_SEQUENTIAL_SCAN (particular for docker load), to avoid eating
// the standby list, particularly when accessing large files such as layer.tar.

// CreateSequential creates the named file with mode 0666 (before umask), truncating
// it if it already exists. If successful, methods on the returned
// File can be used for I/O; the associated file descriptor has mode
// O_RDWR.
// If there is an error, it will be of type *PathError.
func CreateSequential(name string) (*os.File, error) {
	return OpenFileSequential(name, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0)
}

// OpenSequential opens the named file for reading. If successful, methods on
// the returned file can be used for reading; the associated file
// descriptor has mode O_RDONLY.
// If there is an error, it will be of type *PathError.
func OpenSequential(name string) (*os.File, error) {
	return OpenFileSequential(name, os.O_RDONLY, 0)
}

// OpenFileSequential is the generalized open call; most users will use Open
// or Create instead.
// If there is an error, it will be of type *PathError.
func OpenFileSequential(name string, flag int, _ os.FileMode) (*os.File, error) {
	if name == "" {
		return nil, &os.PathError{Op: "open", Path: name, Err: syscall.ENOENT}
	}
	r, errf := windowsOpenFileSequential(name, flag, 0)
	if errf == nil {
		return r, nil
	}
	return nil, &os.PathError{Op: "open", Path: name, Err: errf}
}

func windowsOpenFileSequential(name string, flag int, _ os.FileMode) (file *os.File, err error) {
	r, e := windowsOpenSequential(name, flag|windows.O_CLOEXEC, 0)
	if e != nil {
		return nil, e
	}
	return os.NewFile(uintptr(r), name), nil
}

func makeInheritSa() *windows.SecurityAttributes {
	var sa windows.SecurityAttributes
	sa.Length = uint32(unsafe.Sizeof(sa))
	sa.InheritHandle = 1
	return &sa
}

func windowsOpenSequential(path string, mode int, _ uint32) (fd windows.Handle, err error) {
	if len(path) == 0 {
		return windows.InvalidHandle, windows.ERROR_FILE_NOT_FOUND
	}
	pathp, err := windows.UTF16PtrFromString(path)
	if err != nil {
		return windows.InvalidHandle, err
	}
	var access uint32
	switch mode & (windows.O_RDONLY | windows.O_WRONLY | windows.O_RDWR) {
	case windows.O_RDONLY:
		access = windows.GENERIC_READ
	case windows.O_WRONLY:
		access = windows.GENERIC_WRITE
	case windows.O_RDWR:
		access = windows.GENERIC_READ | windows.GENERIC_WRITE
	}
	if mode&windows.O_CREAT != 0 {
		access |= windows.GENERIC_WRITE
	}
	if mode&windows.O_APPEND != 0 {
		access &^= windows.GENERIC_WRITE
		access |= windows.FILE_APPEND_DATA
	}
	sharemode := uint32(windows.FILE_SHARE_READ | windows.FILE_SHARE_WRITE)
	var sa *windows.SecurityAttributes
	if mode&windows.O_CLOEXEC == 0 {
		sa = makeInheritSa()
	}
	var createmode uint32
	switch {
	case mode&(windows.O_CREAT|windows.O_EXCL) == (windows.O_CREAT | windows.O_EXCL):
		createmode = windows.CREATE_NEW
	case mode&(windows.O_CREAT|windows.O_TRUNC) == (windows.O_CREAT | windows.O_TRUNC):
		createmode = windows.CREATE_ALWAYS
	case mode&windows.O_CREAT == windows.O_CREAT:
		createmode = windows.OPEN_ALWAYS
	case mode&windows.O_TRUNC == windows.O_TRUNC:
		createmode = windows.TRUNCATE_EXISTING
	default:
		createmode = windows.OPEN_EXISTING
	}
	// Use FILE_FLAG_SEQUENTIAL_SCAN rather than FILE_ATTRIBUTE_NORMAL as implemented in golang.
	//https://msdn.microsoft.com/en-us/library/windows/desktop/aa363858(v=vs.85).aspx
	const fileFlagSequentialScan = 0x08000000 // FILE_FLAG_SEQUENTIAL_SCAN
	h, e := windows.CreateFile(pathp, access, sharemode, sa, createmode, fileFlagSequentialScan, 0)
	return h, e
}

// Helpers for TempFileSequential
var rand uint32
var randmu sync.Mutex

func reseed() uint32 {
	return uint32(time.Now().UnixNano() + int64(os.Getpid()))
}
func nextSuffix() string {
	randmu.Lock()
	r := rand
	if r == 0 {
		r = reseed()
	}
	r = r*1664525 + 1013904223 // constants from Numerical Recipes
	rand = r
	randmu.Unlock()
	return strconv.Itoa(int(1e9 + r%1e9))[1:]
}

// TempFileSequential is a copy of ioutil.TempFile, modified to use sequential
// file access. Below is the original comment from golang:
// TempFile creates a new temporary file in the directory dir
// with a name beginning with prefix, opens the file for reading
// and writing, and returns the resulting *os.File.
// If dir is the empty string, TempFile uses the default directory
// for temporary files (see os.TempDir).
// Multiple programs calling TempFile simultaneously
// will not choose the same file. The caller can use f.Name()
// to find the pathname of the file. It is the caller's responsibility
// to remove the file when no longer needed.
func TempFileSequential(dir, prefix string) (f *os.File, err error) {
	if dir == "" {
		dir = os.TempDir()
	}

	nconflict := 0
	for i := 0; i < 10000; i++ {
		name := filepath.Join(dir, prefix+nextSuffix())
		f, err = OpenFileSequential(name, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0600)
		if os.IsExist(err) {
			if nconflict++; nconflict > 10 {
				randmu.Lock()
				rand = reseed()
				randmu.Unlock()
			}
			continue
		}
		break
	}
	return
}
<|endoftext|>


<commit_before>package cache

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"math/big"
	"sync"

	"github.com/golang/glog"
	"github.com/google/trillian"
	"github.com/google/trillian/merkle"
	"github.com/google/trillian/storage"
)

// GetSubtreeFunc describes a function which can return a Subtree from storage.
type GetSubtreeFunc func(id storage.NodeID) (*storage.SubtreeProto, error)

// SetSubtreeFunc describes a function which can store a Subtree into storage.
type SetSubtreeFunc func(s *storage.SubtreeProto) error

// SubtreeCache provides a caching access to Subtree storage.
type SubtreeCache struct {
	// subtrees contains the Subtree data read from storage, and is updated by
	// calls to SetNodeHash.
	subtrees map[string]*storage.SubtreeProto
	// dirtyPrefixes keeps track of all Subtrees which need to be written back
	// to storage.
	dirtyPrefixes map[string]bool
	// mutex guards access to the maps above.
	mutex *sync.RWMutex

	populateSubtree storage.PopulateSubtreeFunc
}

// Suffix represents the tail of a NodeID, indexing into the Subtree which
// corresponds to the prefix of the NodeID.
type Suffix struct {
	// bits is the number of bits in the node ID suffix.
	bits byte
	// path is the suffix itself.
	path byte
}

func (s Suffix) serialize() string {
	r := make([]byte, 2)
	r[0] = s.bits
	r[1] = s.path
	return base64.StdEncoding.EncodeToString(r)
}

const (
	// strataDepth is the depth of Subtree.
	strataDepth = 8
)

// NewSubtreeCache returns a newly intialised cache ready for use.
// populateSubtree is a function which knows how to populate a subtree's
// internal nodes given its leaves, and will be called for each subtree loaded
// from storage.
// TODO(al): consider supporting different sized subtrees - for now everything's subtrees of 8 levels.
func NewSubtreeCache(populateSubtree storage.PopulateSubtreeFunc) SubtreeCache {
	return SubtreeCache{
		subtrees:        make(map[string]*storage.SubtreeProto),
		dirtyPrefixes:   make(map[string]bool),
		mutex:           new(sync.RWMutex),
		populateSubtree: populateSubtree,
	}
}

// splitNodeID breaks a NodeID out into its prefix and suffix parts.
// unless ID is 0 bits long, Suffix must always contain at least one bit.
func splitNodeID(id storage.NodeID) ([]byte, Suffix) {
	if id.PrefixLenBits == 0 {
		return []byte{}, Suffix{bits: 0, path: 0}
	}
	prefixSplit := (id.PrefixLenBits - 1) / strataDepth
	s := Suffix{
		bits: byte((id.PrefixLenBits-1)%strataDepth) + 1,
		path: id.Path[prefixSplit],
	}
	s.path &= ((0x01 << s.bits) - 1) << uint(8-s.bits)

	// TODO(al): is all this copying actually necessary?
	r := make([]byte, prefixSplit)
	copy(r, id.Path[:prefixSplit])
	return r, s
}

// GetNodeHash retrieves the previously written hash and corresponding tree
// revision for the given node ID.
func (s *SubtreeCache) GetNodeHash(id storage.NodeID, getSubtree GetSubtreeFunc) (trillian.Hash, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	return s.getNodeHashUnderLock(id, getSubtree)
}

// getNodeHashUnderLock must be called with s.mutex locked.
func (s *SubtreeCache) getNodeHashUnderLock(id storage.NodeID, getSubtree GetSubtreeFunc) (trillian.Hash, error) {
	px, sx := splitNodeID(id)
	prefixKey := string(px)
	c := s.subtrees[prefixKey]
	if c == nil {
		// Cache miss, so we'll try to fetch from storage.
		subID := id
		subID.PrefixLenBits = len(px) * 8
		var err error
		c, err = getSubtree(subID)
		if err != nil {
			return nil, err
		}
		if c == nil {
			// storage didn't have one for us, so we'll store an empty proto here
			// incase we try to update it later on (we won't flush it back to
			// storage unless it's been written to.)
			c = &storage.SubtreeProto{
				Prefix:        px,
				Depth:         strataDepth,
				Leaves:        make(map[string][]byte),
				InternalNodes: make(map[string][]byte),
			}
		} else {
			if err := s.populateSubtree(c); err != nil {
				return nil, err
			}
		}
		if c.Prefix == nil {
			panic(fmt.Errorf("GetNodeHash nil prefix on %v for id %v with px %#v", c, id.String(), px))
		}

		s.subtrees[prefixKey] = c
	}

	// finally look for the particular node within the subtree so we can return
	// the hash & revision.
	var nh trillian.Hash

	// Look up the hash in the appropriate map.
	// The leaf hashes are stored in a separate map to the internal nodes so that
	// we can easily dump (and later reconstruct) the internal nodes.
	// Since the subtrees are fixed to a depth of 8, any suffix with 8
	// significant bits must be a leaf hash.
	if sx.bits == 8 {
		nh = c.Leaves[sx.serialize()]
	} else {
		nh = c.InternalNodes[sx.serialize()]
	}
	if nh == nil {
		return nil, nil
	}
	return nh, nil
}

// SetNodeHash sets a node hash in the cache.
func (s *SubtreeCache) SetNodeHash(id storage.NodeID, h trillian.Hash, getSubtree GetSubtreeFunc) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	px, sx := splitNodeID(id)
	prefixKey := string(px)
	c := s.subtrees[prefixKey]
	if c == nil {
		// TODO(al): This is ok, IFF *all* leaves in the subtree are being set,
		// verify that this is the case when it happens.
		// For now, just read from storage if we don't already have it.
		glog.Infof("attempting to write to unread subtree for %v, reading now", id.String())
		// We hold the lock so can call this directly:
		_, err := s.getNodeHashUnderLock(id, getSubtree)
		if err != nil {
			return err
		}
		// There must be a subtree present in the cache now, even if storage didn't have anything for us.
		c = s.subtrees[prefixKey]
		if c == nil {
			return fmt.Errorf("internal error, subtree cache for %v is nil after a read attempt", id.String())
		}
	}
	if c.Prefix == nil {
		panic(fmt.Errorf("nil prefix for %v (key %v)", id.String(), prefixKey))
	}
	s.dirtyPrefixes[prefixKey] = true
	// Determine whether we're being asked to store a leaf node, or an internal
	// node, and store it accordingly.
	if sx.bits == 8 {
		c.Leaves[sx.serialize()] = h
	} else {
		c.InternalNodes[sx.serialize()] = h
	}
	return nil
}

// Flush causes the cache to write all dirty Subtrees back to storage.
func (s *SubtreeCache) Flush(setSubtree SetSubtreeFunc) error {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	for k, v := range s.subtrees {
		if s.dirtyPrefixes[k] {
			bk := []byte(k)
			if !bytes.Equal(bk, v.Prefix) {
				return fmt.Errorf("inconsistent cache: prefix key is %v, but cached object claims %v", bk, v.Prefix)
			}
			// TODO(al): Do actually write this one once we're storing the updated
			// subtree root value here during tree update calculations.
			v.RootHash = nil

			if len(v.Leaves) > 0 {
				// clear the internal node cache; we don't want to write that.
				v.InternalNodes = nil
				if err := setSubtree(v); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

// makeSuffixKey creates a suffix key for indexing into the subtree's Leaves and
// InternalNodes maps.
func makeSuffixKey(depth int, index int64) (string, error) {
	// TODO(al): only supports 8 bit subtree sizes currently
	if depth > 8 || depth < 0 {
		return "", fmt.Errorf("found depth of %d, but we only support positive depths of up to and including 8 currently", depth)
	}
	if index > 255 || index < 0 {
		return "", fmt.Errorf("got unsupported index of %d, 0 <= index < 256", index)
	}
	sfx := Suffix{byte(depth), byte(index)}
	return sfx.serialize(), nil
}

// PopulateMapSubtreeNodes re-creates Map subtree's InternalNodes from the
// subtree Leaves map.
//
// This uses HStar2 to repopulate internal nodes.
func PopulateMapSubtreeNodes(treeHasher merkle.TreeHasher) storage.PopulateSubtreeFunc {
	return func(st *storage.SubtreeProto) error {
		st.InternalNodes = make(map[string][]byte)
		rootID := storage.NewNodeIDFromHash(st.Prefix)
		fullTreeDepth := treeHasher.Size() * 8
		leaves := make([]merkle.HStar2LeafHash, 0, len(st.Leaves))
		for k64, v := range st.Leaves {
			k, err := base64.StdEncoding.DecodeString(k64)
			if err != nil {
				return err
			}
			if k[0] != 8 {
				return fmt.Errorf("unexpected non-leaf suffix found: %x", k)
			}
			leaves = append(leaves, merkle.HStar2LeafHash{
				LeafHash: v,
				Index:    big.NewInt(int64(k[1])),
			})
		}
		hs2 := merkle.NewHStar2(treeHasher)
		offset := fullTreeDepth - rootID.PrefixLenBits - int(st.Depth)
		root, err := hs2.HStar2Nodes(int(st.Depth), offset, leaves,
			func(depth int, index *big.Int) (trillian.Hash, error) {
				return nil, nil
			},
			func(depth int, index *big.Int, h trillian.Hash) error {
				i := index.Int64()
				sfx, err := makeSuffixKey(depth, i)
				if err != nil {
					return err
				}
				st.InternalNodes[sfx] = h
				return nil
			})
		if err != nil {
			return err
		}
		st.RootHash = root
		return err
	}
}

// PopulateLogSubtreeNodes re-creates a Log subtree's InternalNodes from the
// subtree Leaves map.
//
// This uses the CompactMerkleTree to repopulate internal nodes, and so will
// handle imperfect (but left-hand dense) subtrees.
func PopulateLogSubtreeNodes(treeHasher merkle.TreeHasher) storage.PopulateSubtreeFunc {
	return func(st *storage.SubtreeProto) error {
		st.InternalNodes = make(map[string][]byte)
		cmt := merkle.NewCompactMerkleTree(treeHasher)
		for leafIndex := int64(0); leafIndex < int64(len(st.Leaves)); leafIndex++ {
			sfx, err := makeSuffixKey(8, leafIndex)
			if err != nil {
				return err
			}
			h := st.Leaves[sfx]
			if h == nil {
				return fmt.Errorf("unexpectedly got nil for subtree leaf suffix %s", sfx)
			}
			seq := cmt.AddLeafHash(h, func(depth int, index int64, h trillian.Hash) {
				if depth == 8 && index == 0 {
					// no space for the root in the node cache
					return
				}
				key, err := makeSuffixKey(8-depth, index<<uint(depth))
				if err != nil {
					// TODO(al): Don't panic Mr. Mainwaring.
					panic(err)
				}
				if depth > 0 {
					st.InternalNodes[key] = h
				}
			})
			if got, expected := seq, leafIndex; got != expected {
				return fmt.Errorf("got seq of %d, but expected %d", got, expected)
			}
		}
		st.RootHash = cmt.CurrentRoot()
		return nil
	}
}
<commit_msg>Quieten noisy subtree cache logging (#162)<commit_after>package cache

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"math/big"
	"sync"

	"github.com/golang/glog"
	"github.com/google/trillian"
	"github.com/google/trillian/merkle"
	"github.com/google/trillian/storage"
)

// GetSubtreeFunc describes a function which can return a Subtree from storage.
type GetSubtreeFunc func(id storage.NodeID) (*storage.SubtreeProto, error)

// SetSubtreeFunc describes a function which can store a Subtree into storage.
type SetSubtreeFunc func(s *storage.SubtreeProto) error

// SubtreeCache provides a caching access to Subtree storage.
type SubtreeCache struct {
	// subtrees contains the Subtree data read from storage, and is updated by
	// calls to SetNodeHash.
	subtrees map[string]*storage.SubtreeProto
	// dirtyPrefixes keeps track of all Subtrees which need to be written back
	// to storage.
	dirtyPrefixes map[string]bool
	// mutex guards access to the maps above.
	mutex *sync.RWMutex

	populateSubtree storage.PopulateSubtreeFunc
}

// Suffix represents the tail of a NodeID, indexing into the Subtree which
// corresponds to the prefix of the NodeID.
type Suffix struct {
	// bits is the number of bits in the node ID suffix.
	bits byte
	// path is the suffix itself.
	path byte
}

func (s Suffix) serialize() string {
	r := make([]byte, 2)
	r[0] = s.bits
	r[1] = s.path
	return base64.StdEncoding.EncodeToString(r)
}

const (
	// strataDepth is the depth of Subtree.
	strataDepth = 8
)

// NewSubtreeCache returns a newly intialised cache ready for use.
// populateSubtree is a function which knows how to populate a subtree's
// internal nodes given its leaves, and will be called for each subtree loaded
// from storage.
// TODO(al): consider supporting different sized subtrees - for now everything's subtrees of 8 levels.
func NewSubtreeCache(populateSubtree storage.PopulateSubtreeFunc) SubtreeCache {
	return SubtreeCache{
		subtrees:        make(map[string]*storage.SubtreeProto),
		dirtyPrefixes:   make(map[string]bool),
		mutex:           new(sync.RWMutex),
		populateSubtree: populateSubtree,
	}
}

// splitNodeID breaks a NodeID out into its prefix and suffix parts.
// unless ID is 0 bits long, Suffix must always contain at least one bit.
func splitNodeID(id storage.NodeID) ([]byte, Suffix) {
	if id.PrefixLenBits == 0 {
		return []byte{}, Suffix{bits: 0, path: 0}
	}
	prefixSplit := (id.PrefixLenBits - 1) / strataDepth
	s := Suffix{
		bits: byte((id.PrefixLenBits-1)%strataDepth) + 1,
		path: id.Path[prefixSplit],
	}
	s.path &= ((0x01 << s.bits) - 1) << uint(8-s.bits)

	// TODO(al): is all this copying actually necessary?
	r := make([]byte, prefixSplit)
	copy(r, id.Path[:prefixSplit])
	return r, s
}

// GetNodeHash retrieves the previously written hash and corresponding tree
// revision for the given node ID.
func (s *SubtreeCache) GetNodeHash(id storage.NodeID, getSubtree GetSubtreeFunc) (trillian.Hash, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	return s.getNodeHashUnderLock(id, getSubtree)
}

// getNodeHashUnderLock must be called with s.mutex locked.
func (s *SubtreeCache) getNodeHashUnderLock(id storage.NodeID, getSubtree GetSubtreeFunc) (trillian.Hash, error) {
	px, sx := splitNodeID(id)
	prefixKey := string(px)
	c := s.subtrees[prefixKey]
	if c == nil {
		// Cache miss, so we'll try to fetch from storage.
		subID := id
		subID.PrefixLenBits = len(px) * 8
		var err error
		c, err = getSubtree(subID)
		if err != nil {
			return nil, err
		}
		if c == nil {
			// storage didn't have one for us, so we'll store an empty proto here
			// incase we try to update it later on (we won't flush it back to
			// storage unless it's been written to.)
			c = &storage.SubtreeProto{
				Prefix:        px,
				Depth:         strataDepth,
				Leaves:        make(map[string][]byte),
				InternalNodes: make(map[string][]byte),
			}
		} else {
			if err := s.populateSubtree(c); err != nil {
				return nil, err
			}
		}
		if c.Prefix == nil {
			panic(fmt.Errorf("GetNodeHash nil prefix on %v for id %v with px %#v", c, id.String(), px))
		}

		s.subtrees[prefixKey] = c
	}

	// finally look for the particular node within the subtree so we can return
	// the hash & revision.
	var nh trillian.Hash

	// Look up the hash in the appropriate map.
	// The leaf hashes are stored in a separate map to the internal nodes so that
	// we can easily dump (and later reconstruct) the internal nodes.
	// Since the subtrees are fixed to a depth of 8, any suffix with 8
	// significant bits must be a leaf hash.
	if sx.bits == 8 {
		nh = c.Leaves[sx.serialize()]
	} else {
		nh = c.InternalNodes[sx.serialize()]
	}
	if nh == nil {
		return nil, nil
	}
	return nh, nil
}

// SetNodeHash sets a node hash in the cache.
func (s *SubtreeCache) SetNodeHash(id storage.NodeID, h trillian.Hash, getSubtree GetSubtreeFunc) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	px, sx := splitNodeID(id)
	prefixKey := string(px)
	c := s.subtrees[prefixKey]
	if c == nil {
		// TODO(al): This is ok, IFF *all* leaves in the subtree are being set,
		// verify that this is the case when it happens.
		// For now, just read from storage if we don't already have it.
		glog.V(1).Infof("attempting to write to unread subtree for %v, reading now", id.String())
		// We hold the lock so can call this directly:
		_, err := s.getNodeHashUnderLock(id, getSubtree)
		if err != nil {
			return err
		}
		// There must be a subtree present in the cache now, even if storage didn't have anything for us.
		c = s.subtrees[prefixKey]
		if c == nil {
			return fmt.Errorf("internal error, subtree cache for %v is nil after a read attempt", id.String())
		}
	}
	if c.Prefix == nil {
		panic(fmt.Errorf("nil prefix for %v (key %v)", id.String(), prefixKey))
	}
	s.dirtyPrefixes[prefixKey] = true
	// Determine whether we're being asked to store a leaf node, or an internal
	// node, and store it accordingly.
	if sx.bits == 8 {
		c.Leaves[sx.serialize()] = h
	} else {
		c.InternalNodes[sx.serialize()] = h
	}
	return nil
}

// Flush causes the cache to write all dirty Subtrees back to storage.
func (s *SubtreeCache) Flush(setSubtree SetSubtreeFunc) error {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	for k, v := range s.subtrees {
		if s.dirtyPrefixes[k] {
			bk := []byte(k)
			if !bytes.Equal(bk, v.Prefix) {
				return fmt.Errorf("inconsistent cache: prefix key is %v, but cached object claims %v", bk, v.Prefix)
			}
			// TODO(al): Do actually write this one once we're storing the updated
			// subtree root value here during tree update calculations.
			v.RootHash = nil

			if len(v.Leaves) > 0 {
				// clear the internal node cache; we don't want to write that.
				v.InternalNodes = nil
				if err := setSubtree(v); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

// makeSuffixKey creates a suffix key for indexing into the subtree's Leaves and
// InternalNodes maps.
func makeSuffixKey(depth int, index int64) (string, error) {
	// TODO(al): only supports 8 bit subtree sizes currently
	if depth > 8 || depth < 0 {
		return "", fmt.Errorf("found depth of %d, but we only support positive depths of up to and including 8 currently", depth)
	}
	if index > 255 || index < 0 {
		return "", fmt.Errorf("got unsupported index of %d, 0 <= index < 256", index)
	}
	sfx := Suffix{byte(depth), byte(index)}
	return sfx.serialize(), nil
}

// PopulateMapSubtreeNodes re-creates Map subtree's InternalNodes from the
// subtree Leaves map.
//
// This uses HStar2 to repopulate internal nodes.
func PopulateMapSubtreeNodes(treeHasher merkle.TreeHasher) storage.PopulateSubtreeFunc {
	return func(st *storage.SubtreeProto) error {
		st.InternalNodes = make(map[string][]byte)
		rootID := storage.NewNodeIDFromHash(st.Prefix)
		fullTreeDepth := treeHasher.Size() * 8
		leaves := make([]merkle.HStar2LeafHash, 0, len(st.Leaves))
		for k64, v := range st.Leaves {
			k, err := base64.StdEncoding.DecodeString(k64)
			if err != nil {
				return err
			}
			if k[0] != 8 {
				return fmt.Errorf("unexpected non-leaf suffix found: %x", k)
			}
			leaves = append(leaves, merkle.HStar2LeafHash{
				LeafHash: v,
				Index:    big.NewInt(int64(k[1])),
			})
		}
		hs2 := merkle.NewHStar2(treeHasher)
		offset := fullTreeDepth - rootID.PrefixLenBits - int(st.Depth)
		root, err := hs2.HStar2Nodes(int(st.Depth), offset, leaves,
			func(depth int, index *big.Int) (trillian.Hash, error) {
				return nil, nil
			},
			func(depth int, index *big.Int, h trillian.Hash) error {
				i := index.Int64()
				sfx, err := makeSuffixKey(depth, i)
				if err != nil {
					return err
				}
				st.InternalNodes[sfx] = h
				return nil
			})
		if err != nil {
			return err
		}
		st.RootHash = root
		return err
	}
}

// PopulateLogSubtreeNodes re-creates a Log subtree's InternalNodes from the
// subtree Leaves map.
//
// This uses the CompactMerkleTree to repopulate internal nodes, and so will
// handle imperfect (but left-hand dense) subtrees.
func PopulateLogSubtreeNodes(treeHasher merkle.TreeHasher) storage.PopulateSubtreeFunc {
	return func(st *storage.SubtreeProto) error {
		st.InternalNodes = make(map[string][]byte)
		cmt := merkle.NewCompactMerkleTree(treeHasher)
		for leafIndex := int64(0); leafIndex < int64(len(st.Leaves)); leafIndex++ {
			sfx, err := makeSuffixKey(8, leafIndex)
			if err != nil {
				return err
			}
			h := st.Leaves[sfx]
			if h == nil {
				return fmt.Errorf("unexpectedly got nil for subtree leaf suffix %s", sfx)
			}
			seq := cmt.AddLeafHash(h, func(depth int, index int64, h trillian.Hash) {
				if depth == 8 && index == 0 {
					// no space for the root in the node cache
					return
				}
				key, err := makeSuffixKey(8-depth, index<<uint(depth))
				if err != nil {
					// TODO(al): Don't panic Mr. Mainwaring.
					panic(err)
				}
				if depth > 0 {
					st.InternalNodes[key] = h
				}
			})
			if got, expected := seq, leafIndex; got != expected {
				return fmt.Errorf("got seq of %d, but expected %d", got, expected)
			}
		}
		st.RootHash = cmt.CurrentRoot()
		return nil
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"

	"github.com/aws/aws-sdk-go/aws"
	events "github.com/aws/aws-sdk-go/service/cloudwatchevents"
)

func resourceAwsCloudWatchEventRule() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsCloudWatchEventRuleCreate,
		Read:   resourceAwsCloudWatchEventRuleRead,
		Update: resourceAwsCloudWatchEventRuleUpdate,
		Delete: resourceAwsCloudWatchEventRuleDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateCloudWatchEventRuleName,
			},
			"schedule_expression": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateMaxLength(256),
			},
			"event_pattern": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateMaxLength(2048),
				StateFunc:    normalizeJson,
			},
			"description": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateMaxLength(512),
			},
			"role_arn": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateMaxLength(1600),
			},
			"is_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsCloudWatchEventRuleCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudwatcheventsconn

	input := buildPutRuleInputStruct(d)
	log.Printf("[DEBUG] Creating CloudWatch Event Rule: %s", input)

	// IAM Roles take some time to propagate
	var out *events.PutRuleOutput
	err := resource.Retry(30*time.Second, func() *resource.RetryError {
		var err error
		out, err = conn.PutRule(input)
		pattern := regexp.MustCompile("cannot be assumed by principal '[a-z]+\\.amazonaws\\.com'\\.$")
		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok {
				if awsErr.Code() == "ValidationException" && pattern.MatchString(awsErr.Message()) {
					log.Printf("[DEBUG] Retrying creation of CloudWatch Event Rule %q", *input.Name)
					return resource.RetryableError(err)
				}
			}
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("Creating CloudWatch Event Rule failed: %s", err)
	}

	d.Set("arn", out.RuleArn)
	d.SetId(d.Get("name").(string))

	log.Printf("[INFO] CloudWatch Event Rule %q created", *out.RuleArn)

	return resourceAwsCloudWatchEventRuleUpdate(d, meta)
}

func resourceAwsCloudWatchEventRuleRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudwatcheventsconn

	input := events.DescribeRuleInput{
		Name: aws.String(d.Id()),
	}
	log.Printf("[DEBUG] Reading CloudWatch Event Rule: %s", input)
	out, err := conn.DescribeRule(&input)
	if awsErr, ok := err.(awserr.Error); ok {
		if awsErr.Code() == "ResourceNotFoundException" {
			log.Printf("[WARN] Removing CloudWatch Event Rule %q because it's gone.", d.Id())
			d.SetId("")
			return nil
		}
	}
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Found Event Rule: %s", out)

	d.Set("arn", out.Arn)
	d.Set("description", out.Description)
	if out.EventPattern != nil {
		d.Set("event_pattern", normalizeJson(*out.EventPattern))
	}
	d.Set("name", out.Name)
	d.Set("role_arn", out.RoleArn)
	d.Set("schedule_expression", out.ScheduleExpression)

	boolState, err := getBooleanStateFromString(*out.State)
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Setting boolean state: %t", boolState)
	d.Set("is_enabled", boolState)

	return nil
}

func resourceAwsCloudWatchEventRuleUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudwatcheventsconn

	if d.HasChange("is_enabled") && d.Get("is_enabled").(bool) {
		log.Printf("[DEBUG] Enabling CloudWatch Event Rule %q", d.Id())
		_, err := conn.EnableRule(&events.EnableRuleInput{
			Name: aws.String(d.Id()),
		})
		if err != nil {
			return err
		}
		log.Printf("[DEBUG] CloudWatch Event Rule (%q) enabled", d.Id())
	}

	input := buildPutRuleInputStruct(d)
	log.Printf("[DEBUG] Updating CloudWatch Event Rule: %s", input)

	// IAM Roles take some time to propagate
	var out *events.PutRuleOutput
	err := resource.Retry(30*time.Second, func() *resource.RetryError {
		var err error
		out, err = conn.PutRule(input)
		pattern := regexp.MustCompile("cannot be assumed by principal '[a-z]+\\.amazonaws\\.com'\\.$")
		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok {
				if awsErr.Code() == "ValidationException" && pattern.MatchString(awsErr.Message()) {
					log.Printf("[DEBUG] Retrying update of CloudWatch Event Rule %q", *input.Name)
					return resource.RetryableError(err)
				}
			}
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("Updating CloudWatch Event Rule failed: %s", err)
	}

	if d.HasChange("is_enabled") && !d.Get("is_enabled").(bool) {
		log.Printf("[DEBUG] Disabling CloudWatch Event Rule %q", d.Id())
		_, err := conn.DisableRule(&events.DisableRuleInput{
			Name: aws.String(d.Id()),
		})
		if err != nil {
			return err
		}
		log.Printf("[DEBUG] CloudWatch Event Rule (%q) disabled", d.Id())
	}

	return resourceAwsCloudWatchEventRuleRead(d, meta)
}

func resourceAwsCloudWatchEventRuleDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudwatcheventsconn

	log.Printf("[INFO] Deleting CloudWatch Event Rule: %s", d.Id())
	_, err := conn.DeleteRule(&events.DeleteRuleInput{
		Name: aws.String(d.Id()),
	})
	if err != nil {
		return fmt.Errorf("Error deleting CloudWatch Event Rule: %s", err)
	}
	log.Println("[INFO] CloudWatch Event Rule deleted")

	d.SetId("")

	return nil
}

func buildPutRuleInputStruct(d *schema.ResourceData) *events.PutRuleInput {
	input := events.PutRuleInput{
		Name: aws.String(d.Get("name").(string)),
	}
	if v, ok := d.GetOk("description"); ok {
		input.Description = aws.String(v.(string))
	}
	if v, ok := d.GetOk("event_pattern"); ok {
		input.EventPattern = aws.String(v.(string))
	}
	if v, ok := d.GetOk("role_arn"); ok {
		input.RoleArn = aws.String(v.(string))
	}
	if v, ok := d.GetOk("schedule_expression"); ok {
		input.ScheduleExpression = aws.String(v.(string))
	}

	input.State = aws.String(getStringStateFromBoolean(d.Get("is_enabled").(bool)))

	return &input
}

// State is represented as (ENABLED|DISABLED) in the API
func getBooleanStateFromString(state string) (bool, error) {
	if state == "ENABLED" {
		return true, nil
	} else if state == "DISABLED" {
		return false, nil
	}
	// We don't just blindly trust AWS as they tend to return
	// unexpected values in similar cases (different casing etc.)
	return false, fmt.Errorf("Failed converting state %q into boolean", state)
}

// State is represented as (ENABLED|DISABLED) in the API
func getStringStateFromBoolean(isEnabled bool) string {
	if isEnabled {
		return "ENABLED"
	}
	return "DISABLED"
}
<commit_msg>aws provider: normalize json of cloudwatch event_pattern<commit_after>package aws

import (
	"fmt"
	"log"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"

	"github.com/aws/aws-sdk-go/aws"
	events "github.com/aws/aws-sdk-go/service/cloudwatchevents"
)

func resourceAwsCloudWatchEventRule() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsCloudWatchEventRuleCreate,
		Read:   resourceAwsCloudWatchEventRuleRead,
		Update: resourceAwsCloudWatchEventRuleUpdate,
		Delete: resourceAwsCloudWatchEventRuleDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateCloudWatchEventRuleName,
			},
			"schedule_expression": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateMaxLength(256),
			},
			"event_pattern": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateMaxLength(2048),
				StateFunc:    normalizeJson,
			},
			"description": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateMaxLength(512),
			},
			"role_arn": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateMaxLength(1600),
			},
			"is_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsCloudWatchEventRuleCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudwatcheventsconn

	input := buildPutRuleInputStruct(d)
	log.Printf("[DEBUG] Creating CloudWatch Event Rule: %s", input)

	// IAM Roles take some time to propagate
	var out *events.PutRuleOutput
	err := resource.Retry(30*time.Second, func() *resource.RetryError {
		var err error
		out, err = conn.PutRule(input)
		pattern := regexp.MustCompile("cannot be assumed by principal '[a-z]+\\.amazonaws\\.com'\\.$")
		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok {
				if awsErr.Code() == "ValidationException" && pattern.MatchString(awsErr.Message()) {
					log.Printf("[DEBUG] Retrying creation of CloudWatch Event Rule %q", *input.Name)
					return resource.RetryableError(err)
				}
			}
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("Creating CloudWatch Event Rule failed: %s", err)
	}

	d.Set("arn", out.RuleArn)
	d.SetId(d.Get("name").(string))

	log.Printf("[INFO] CloudWatch Event Rule %q created", *out.RuleArn)

	return resourceAwsCloudWatchEventRuleUpdate(d, meta)
}

func resourceAwsCloudWatchEventRuleRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudwatcheventsconn

	input := events.DescribeRuleInput{
		Name: aws.String(d.Id()),
	}
	log.Printf("[DEBUG] Reading CloudWatch Event Rule: %s", input)
	out, err := conn.DescribeRule(&input)
	if awsErr, ok := err.(awserr.Error); ok {
		if awsErr.Code() == "ResourceNotFoundException" {
			log.Printf("[WARN] Removing CloudWatch Event Rule %q because it's gone.", d.Id())
			d.SetId("")
			return nil
		}
	}
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Found Event Rule: %s", out)

	d.Set("arn", out.Arn)
	d.Set("description", out.Description)
	if out.EventPattern != nil {
		d.Set("event_pattern", normalizeJson(*out.EventPattern))
	}
	d.Set("name", out.Name)
	d.Set("role_arn", out.RoleArn)
	d.Set("schedule_expression", out.ScheduleExpression)

	boolState, err := getBooleanStateFromString(*out.State)
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Setting boolean state: %t", boolState)
	d.Set("is_enabled", boolState)

	return nil
}

func resourceAwsCloudWatchEventRuleUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudwatcheventsconn

	if d.HasChange("is_enabled") && d.Get("is_enabled").(bool) {
		log.Printf("[DEBUG] Enabling CloudWatch Event Rule %q", d.Id())
		_, err := conn.EnableRule(&events.EnableRuleInput{
			Name: aws.String(d.Id()),
		})
		if err != nil {
			return err
		}
		log.Printf("[DEBUG] CloudWatch Event Rule (%q) enabled", d.Id())
	}

	input := buildPutRuleInputStruct(d)
	log.Printf("[DEBUG] Updating CloudWatch Event Rule: %s", input)

	// IAM Roles take some time to propagate
	var out *events.PutRuleOutput
	err := resource.Retry(30*time.Second, func() *resource.RetryError {
		var err error
		out, err = conn.PutRule(input)
		pattern := regexp.MustCompile("cannot be assumed by principal '[a-z]+\\.amazonaws\\.com'\\.$")
		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok {
				if awsErr.Code() == "ValidationException" && pattern.MatchString(awsErr.Message()) {
					log.Printf("[DEBUG] Retrying update of CloudWatch Event Rule %q", *input.Name)
					return resource.RetryableError(err)
				}
			}
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("Updating CloudWatch Event Rule failed: %s", err)
	}

	if d.HasChange("is_enabled") && !d.Get("is_enabled").(bool) {
		log.Printf("[DEBUG] Disabling CloudWatch Event Rule %q", d.Id())
		_, err := conn.DisableRule(&events.DisableRuleInput{
			Name: aws.String(d.Id()),
		})
		if err != nil {
			return err
		}
		log.Printf("[DEBUG] CloudWatch Event Rule (%q) disabled", d.Id())
	}

	return resourceAwsCloudWatchEventRuleRead(d, meta)
}

func resourceAwsCloudWatchEventRuleDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudwatcheventsconn

	log.Printf("[INFO] Deleting CloudWatch Event Rule: %s", d.Id())
	_, err := conn.DeleteRule(&events.DeleteRuleInput{
		Name: aws.String(d.Id()),
	})
	if err != nil {
		return fmt.Errorf("Error deleting CloudWatch Event Rule: %s", err)
	}
	log.Println("[INFO] CloudWatch Event Rule deleted")

	d.SetId("")

	return nil
}

func buildPutRuleInputStruct(d *schema.ResourceData) *events.PutRuleInput {
	input := events.PutRuleInput{
		Name: aws.String(d.Get("name").(string)),
	}
	if v, ok := d.GetOk("description"); ok {
		input.Description = aws.String(v.(string))
	}
	if v, ok := d.GetOk("event_pattern"); ok {
		input.EventPattern = aws.String(normalizeJson(v.(string)))
	}
	if v, ok := d.GetOk("role_arn"); ok {
		input.RoleArn = aws.String(v.(string))
	}
	if v, ok := d.GetOk("schedule_expression"); ok {
		input.ScheduleExpression = aws.String(v.(string))
	}

	input.State = aws.String(getStringStateFromBoolean(d.Get("is_enabled").(bool)))

	return &input
}

// State is represented as (ENABLED|DISABLED) in the API
func getBooleanStateFromString(state string) (bool, error) {
	if state == "ENABLED" {
		return true, nil
	} else if state == "DISABLED" {
		return false, nil
	}
	// We don't just blindly trust AWS as they tend to return
	// unexpected values in similar cases (different casing etc.)
	return false, fmt.Errorf("Failed converting state %q into boolean", state)
}

// State is represented as (ENABLED|DISABLED) in the API
func getStringStateFromBoolean(isEnabled bool) string {
	if isEnabled {
		return "ENABLED"
	}
	return "DISABLED"
}
<|endoftext|>


<commit_before>package workday

import (
	"errors"
	"os"
	"path/filepath"
	"time"

	"github.com/nightlyone/lockfile"
)

var (
	lfile          lockfile.Lockfile
	ErrLockTimeout = errors.New("lock timeout")
)

func tryLockFor(l lockfile.Lockfile, d time.Duration) error {
	start := time.Now()
	for time.Now().Sub(start) < d {
		if err := l.TryLock(); err == nil {
			return nil
		}
		time.Sleep(1)
	}
	return ErrLockTimeout
}

func createLock(fpath string) (lockfile.Lockfile, error) {
	// create the file if it doesn't exist
	if exists, err := fileExists(fpath); err != nil {
		return "", err
	} else if exists {
		return "", nil
	}
	f, err := os.Create(fpath)
	if err != nil {
		return "", err
	}
	if err := f.Close(); err != nil {
		return "", err
	}
	return lockfile.New(fpath)
}

func LockDataDir() error {
	var (
		err   error
		lpath = filepath.Join(DataDir, "lock")
	)
	lfile, err = createLock(lpath)
	if err != nil {
		return err
	}
	return tryLockFor(lfile, time.Second)
}

func UnlockDataDir() error {
	return lfile.Unlock()
}
<commit_msg>create lock when it exists<commit_after>package workday

import (
	"errors"
	"os"
	"path/filepath"
	"time"

	"github.com/nightlyone/lockfile"
)

var (
	lfile          lockfile.Lockfile
	ErrLockTimeout = errors.New("lock timeout")
)

func tryLockFor(l lockfile.Lockfile, d time.Duration) error {
	start := time.Now()
	for time.Now().Sub(start) < d {
		if err := l.TryLock(); err == nil {
			return nil
		}
		time.Sleep(1)
	}
	return ErrLockTimeout
}

func createLock(fpath string) (lockfile.Lockfile, error) {
	// create the file if it doesn't exist
	if exists, err := fileExists(fpath); err != nil {
		return "", err
	} else if !exists {
		f, err := os.Create(fpath)
		if err != nil {
			return "", err
		}
		if err := f.Close(); err != nil {
			return "", err
		}
	}
	return lockfile.New(fpath)
}

func LockDataDir() error {
	var (
		err   error
		lpath = filepath.Join(DataDir, "lock")
	)
	lfile, err = createLock(lpath)
	if err != nil {
		return err
	}
	return tryLockFor(lfile, time.Second)
}

func UnlockDataDir() error {
	return lfile.Unlock()
}
<|endoftext|>


<commit_before>package storage

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

	"github.com/garyburd/redigo/redis"
)

type Storage struct {
	pool *redis.Pool
}

func (stg *Storage) AddSubscriber(appID string, channelID string, subscriberIDs []string) {
	conn := stg.pool.Get()
	defer conn.Close()

	stg.AddChannel(appID, channelID)
	subscribersKey := addPrefix("apps." + appID + ".channels." + channelID + ".subscribers")
	tmpParams := append([]string{subscribersKey}, subscriberIDs...)

	params := make([]interface{}, len(tmpParams))
	for i, v := range tmpParams {
		params[i] = v
	}

	conn.Do("SADD", params...)
}

func (stg *Storage) AddChannel(appID string, channelID string) {
	conn := stg.pool.Get()
	defer conn.Close()

	channelsKey := addPrefix("apps." + appID + ".channels")
	conn.Do("SADD", channelsKey, channelID)
}

func (stg *Storage) DeleteChannel(appID string, channelID string) {
	conn := stg.pool.Get()
	defer conn.Close()

	channelsKey := addPrefix("apps." + appID + ".channels")
	conn.Do("SREM", channelsKey, channelID)

	channelKey := channelsKey + "." + channelID + ".subscribers"
	conn.Do("DEL", channelKey)
}

func (stg *Storage) AddSubscriberDevice(appID string, subscriberID string, device *Device) error {
	conn := stg.pool.Get()
	defer conn.Close()

	subscribersKey := addPrefix("apps." + appID + ".subscribers")
	conn.Do("SADD", subscribersKey, subscriberID)

	devicesKey := subscribersKey + "." + subscriberID + ".devices"
	jstring, _ := json.Marshal(device)
	// todo: multiple devices with same platform and token should not be added
	_, err := conn.Do("HSET", devicesKey, device.Token, jstring)

	if err != nil {
		return err
	}
	return nil
}

func (stg *Storage) UpdateDeviceToken(appID string, subscriberID string, oldDeviceToken string, newDeviceToken string) error {
	conn := stg.pool.Get()
	defer conn.Close()

	key := addPrefix("apps." + appID + ".subscribers." + subscriberID + ".devices")
	deviceData, err := redis.String(conn.Do("HGET", key, oldDeviceToken))
	if err != nil {
		return err
	}

	conn.Do("HDEL", key, oldDeviceToken)

	var device Device
	decoder := json.NewDecoder(strings.NewReader(deviceData))

	if err := decoder.Decode(&device); err != nil {
		return err
	}

	device.Token = newDeviceToken
	err = stg.AddSubscriberDevice(appID, subscriberID, &device)

	if err != nil {
		return err
	}

	return nil
}

func (stg *Storage) GetSubscriberDevices(appID string, subscriberID string) ([]Device, error) {
	conn := stg.pool.Get()
	defer conn.Close()

	key := addPrefix("apps." + appID + ".subscribers." + subscriberID + ".devices")

	var devices map[string]string
	devices, err := redis.StringMap(conn.Do("HGETALL", key))
	if err != nil {
		return nil, err
	}

	var device Device
	var response []Device
	for _, deviceData := range devices {
		decoder := json.NewDecoder(strings.NewReader(deviceData))
		decoder.Decode(&device)
		response = append(response, device)
	}

	return response, nil
}

func (stg *Storage) AppExists(appID string) bool {
	conn := stg.pool.Get()
	defer conn.Close()

	status, err := redis.Int(conn.Do("SISMEMBER", addPrefix("apps"), appID))

	if status == 0 || err != nil {
		return false
	}

	return true
}

func (stg *Storage) CreateApp(appID string, appData string) {
	conn := stg.pool.Get()
	defer conn.Close()

	appsKey := addPrefix("apps")
	conn.Do("SADD", appsKey, appID)
	appKey := appsKey + "." + appID
	conn.Do("SET", appKey, appData)
}

func (stg *Storage) GetApp(appID string) (string, error) {
	conn := stg.pool.Get()
	defer conn.Close()

	appKey := addPrefix("apps") + "." + appID
	value, err := redis.String(conn.Do("GET", appKey))
	if err != nil {
		return "", err
	}
	return value, nil
}

func addPrefix(key string) string {
	return fmt.Sprintf("srv.push.%s", key)
}

func Init(conf *RedisConfig) *Storage {
	stg := new(Storage)
	pool := &redis.Pool{
		MaxIdle:     conf.MaxIdle,
		MaxActive:   conf.MaxActive,
		IdleTimeout: time.Duration(conf.IdleTimeout) * time.Second,
		Wait:        conf.Wait,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial(conf.Network, conf.Addr)
			if err != nil {
				return nil, err
			}

			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
	stg.pool = pool
	return stg
}
<commit_msg>Add GetChannelSubscribers method to storage<commit_after>package storage

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

	"github.com/garyburd/redigo/redis"
)

type Storage struct {
	pool *redis.Pool
}

func (stg *Storage) AddSubscriber(appID string, channelID string, subscriberIDs []string) {
	conn := stg.pool.Get()
	defer conn.Close()

	stg.AddChannel(appID, channelID)
	subscribersKey := addPrefix("apps." + appID + ".channels." + channelID + ".subscribers")
	tmpParams := append([]string{subscribersKey}, subscriberIDs...)

	params := make([]interface{}, len(tmpParams))
	for i, v := range tmpParams {
		params[i] = v
	}

	conn.Do("SADD", params...)
}

func (stg *Storage) AddChannel(appID string, channelID string) {
	conn := stg.pool.Get()
	defer conn.Close()

	channelsKey := addPrefix("apps." + appID + ".channels")
	conn.Do("SADD", channelsKey, channelID)
}

func (stg *Storage) DeleteChannel(appID string, channelID string) {
	conn := stg.pool.Get()
	defer conn.Close()

	channelsKey := addPrefix("apps." + appID + ".channels")
	conn.Do("SREM", channelsKey, channelID)

	channelKey := channelsKey + "." + channelID + ".subscribers"
	conn.Do("DEL", channelKey)
}

func (stg *Storage) AddSubscriberDevice(appID string, subscriberID string, device *Device) error {
	conn := stg.pool.Get()
	defer conn.Close()

	subscribersKey := addPrefix("apps." + appID + ".subscribers")
	conn.Do("SADD", subscribersKey, subscriberID)

	devicesKey := subscribersKey + "." + subscriberID + ".devices"
	jstring, _ := json.Marshal(device)
	// todo: multiple devices with same platform and token should not be added
	_, err := conn.Do("HSET", devicesKey, device.Token, jstring)

	if err != nil {
		return err
	}
	return nil
}

func (stg *Storage) UpdateDeviceToken(appID string, subscriberID string, oldDeviceToken string, newDeviceToken string) error {
	conn := stg.pool.Get()
	defer conn.Close()

	key := addPrefix("apps." + appID + ".subscribers." + subscriberID + ".devices")
	deviceData, err := redis.String(conn.Do("HGET", key, oldDeviceToken))
	if err != nil {
		return err
	}

	conn.Do("HDEL", key, oldDeviceToken)

	var device Device
	decoder := json.NewDecoder(strings.NewReader(deviceData))

	if err := decoder.Decode(&device); err != nil {
		return err
	}

	device.Token = newDeviceToken
	err = stg.AddSubscriberDevice(appID, subscriberID, &device)

	if err != nil {
		return err
	}

	return nil
}

func (stg *Storage) GetChannelSubscribers(appID string, channelID string) ([]string, error) {
	conn := stg.pool.Get()
	defer conn.Close()

	key := addPrefix("apps." + appID + ".channels." + channelID + ".subscribers")
	subscribers, err := redis.Strings(conn.Do("SMEMBERS", key))

	if err != nil {
		return nil, err
	}

	return subscribers, nil
}

func (stg *Storage) GetSubscriberDevices(appID string, subscriberID string) ([]Device, error) {
	conn := stg.pool.Get()
	defer conn.Close()

	key := addPrefix("apps." + appID + ".subscribers." + subscriberID + ".devices")

	var devices map[string]string
	devices, err := redis.StringMap(conn.Do("HGETALL", key))
	if err != nil {
		return nil, err
	}

	var device Device
	var response []Device
	for _, deviceData := range devices {
		decoder := json.NewDecoder(strings.NewReader(deviceData))
		decoder.Decode(&device)
		response = append(response, device)
	}

	return response, nil
}

func (stg *Storage) AppExists(appID string) bool {
	conn := stg.pool.Get()
	defer conn.Close()

	status, err := redis.Int(conn.Do("SISMEMBER", addPrefix("apps"), appID))

	if status == 0 || err != nil {
		return false
	}

	return true
}

func (stg *Storage) CreateApp(appID string, appData string) {
	conn := stg.pool.Get()
	defer conn.Close()

	appsKey := addPrefix("apps")
	conn.Do("SADD", appsKey, appID)
	appKey := appsKey + "." + appID
	conn.Do("SET", appKey, appData)
}

func (stg *Storage) GetApp(appID string) (string, error) {
	conn := stg.pool.Get()
	defer conn.Close()

	appKey := addPrefix("apps") + "." + appID
	value, err := redis.String(conn.Do("GET", appKey))
	if err != nil {
		return "", err
	}
	return value, nil
}

func addPrefix(key string) string {
	return fmt.Sprintf("srv.push.%s", key)
}

func Init(conf *RedisConfig) *Storage {
	stg := new(Storage)
	pool := &redis.Pool{
		MaxIdle:     conf.MaxIdle,
		MaxActive:   conf.MaxActive,
		IdleTimeout: time.Duration(conf.IdleTimeout) * time.Second,
		Wait:        conf.Wait,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial(conf.Network, conf.Addr)
			if err != nil {
				return nil, err
			}

			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
	stg.pool = pool
	return stg
}
<|endoftext|>


<commit_before>package etcdstore

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/url"
	"strconv"
	"strings"

	etcdclient "github.com/coreos/etcd/client"
	engineapi "github.com/docker/docker/client"
	"github.com/projecteru2/core/types"
	"github.com/projecteru2/core/utils"
	log "github.com/sirupsen/logrus"
)

// AddNode save it to etcd
// storage path in etcd is `/pod/:podname/node/:nodename/info`
func (k *Krypton) AddNode(ctx context.Context, name, endpoint, podname, ca, cert, key string, cpu int, share, memory int64, labels map[string]string) (*types.Node, error) {
	if !strings.HasPrefix(endpoint, "tcp://") {
		return nil, fmt.Errorf("Endpoint must starts with tcp:// %q", endpoint)
	}

	_, err := k.GetPod(ctx, podname)
	if err != nil {
		return nil, err
	}

	nodeKey := fmt.Sprintf(nodeInfoKey, podname, name)
	if _, err := k.etcd.Get(ctx, nodeKey, nil); err == nil {
		return nil, fmt.Errorf("Node (%s, %s) already exists", podname, name)
	}

	// 如果有tls的证书需要保存就保存一下
	if ca != "" && cert != "" && key != "" {
		_, err = k.etcd.Set(ctx, fmt.Sprintf(nodeCaKey, podname, name), ca, nil)
		if err != nil {
			return nil, err
		}
		_, err = k.etcd.Set(ctx, fmt.Sprintf(nodeCertKey, podname, name), cert, nil)
		if err != nil {
			return nil, err
		}
		_, err = k.etcd.Set(ctx, fmt.Sprintf(nodeKeyKey, podname, name), key, nil)
		if err != nil {
			return nil, err
		}
	}

	// 尝试加载docker的客户端
	engine, err := k.makeDockerClient(ctx, podname, name, endpoint, true)
	if err != nil {
		k.deleteNode(ctx, podname, name, endpoint)
		return nil, err
	}

	// 判断这货是不是活着的
	info, err := engine.Info(ctx)
	if err != nil {
		k.deleteNode(ctx, podname, name, endpoint)
		return nil, err
	}

	ncpu := cpu
	memcap := memory
	if cpu == 0 && memory == 0 {
		ncpu = info.NCPU
		memcap = info.MemTotal - gigabyte
	}
	if share == 0 {
		share = k.config.Scheduler.ShareBase
	}
	cpumap := types.CPUMap{}
	for i := 0; i < ncpu; i++ {
		cpumap[strconv.Itoa(i)] = share
	}

	node := &types.Node{
		Name:      name,
		Endpoint:  endpoint,
		Podname:   podname,
		CPU:       cpumap,
		MemCap:    memcap,
		Available: true,
		Labels:    labels,
	}

	bytes, err := json.Marshal(node)
	if err != nil {
		k.deleteNode(ctx, podname, name, endpoint)
		return nil, err
	}

	_, err = k.etcd.Create(ctx, nodeKey, string(bytes))
	if err != nil {
		k.deleteNode(ctx, podname, name, endpoint)
		return nil, err
	}

	_, err = k.etcd.Create(ctx, fmt.Sprintf(nodePodKey, name), podname)
	if err != nil {
		k.deleteNode(ctx, podname, name, endpoint)
		return nil, err
	}

	return node, nil
}

// 因为是先写etcd的证书再拿client
// 所以可能出现实际上node创建失败但是却写好了证书的情况
// 所以需要删除这些留存的证书
// 至于结果是不是成功就无所谓了
func (k *Krypton) deleteNode(ctx context.Context, podname, nodename, endpoint string) {
	key := fmt.Sprintf(nodePrefixKey, podname, nodename)
	k.etcd.Delete(ctx, key, &etcdclient.DeleteOptions{Recursive: true})
	k.etcd.Delete(ctx, fmt.Sprintf(nodePodKey, nodename), &etcdclient.DeleteOptions{})
	u, err := url.Parse(endpoint)
	if err != nil {
		log.Errorf("[deleteNode] Bad endpoint: %s", endpoint)
		return
	}

	host, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		log.Errorf("[deleteNode] Bad addr: %s", u.Host)
		return
	}
	_cache.delete(host)
	log.Debugf("[deleteNode] Node (%s, %s, %s) deleted", podname, nodename, endpoint)
}

// DeleteNode delete a node
func (k *Krypton) DeleteNode(ctx context.Context, node *types.Node) {
	k.deleteNode(ctx, node.Podname, node.Name, node.Endpoint)
}

// GetNode get a node from etcd
// and construct it's docker client
// a node must belong to a pod
// and since node is not the smallest unit to user, to get a node we must specify the corresponding pod
// storage path in etcd is `/pod/:podname/node/:nodename/info`
func (k *Krypton) GetNode(ctx context.Context, podname, nodename string) (*types.Node, error) {
	key := fmt.Sprintf(nodeInfoKey, podname, nodename)
	resp, err := k.etcd.Get(ctx, key, nil)
	if err != nil {
		return nil, err
	}
	if resp.Node.Dir {
		return nil, fmt.Errorf("Node storage path %q in etcd is a directory", key)
	}

	node := &types.Node{}
	if err := json.Unmarshal([]byte(resp.Node.Value), node); err != nil {
		return nil, err
	}

	engine, err := k.makeDockerClient(ctx, podname, nodename, node.Endpoint, false)
	if err != nil {
		return nil, err
	}

	node.Engine = engine
	return node, nil
}

// GetNodeByName get node by name
func (k *Krypton) GetNodeByName(ctx context.Context, nodename string) (node *types.Node, err error) {
	key := fmt.Sprintf(nodePodKey, nodename)
	resp, err := k.etcd.Get(ctx, key, nil)
	if err != nil {
		return nil, err
	}
	podname := resp.Node.Value
	return k.GetNode(ctx, podname, nodename)
}

// GetNodesByPod get all nodes bound to pod
// here we use podname instead of pod instance
// storage path in etcd is `/pod/:podname/node`
func (k *Krypton) GetNodesByPod(ctx context.Context, podname string) (nodes []*types.Node, err error) {
	key := fmt.Sprintf(podNodesKey, podname)
	resp, err := k.etcd.Get(ctx, key, nil)
	if err != nil {
		return nodes, err
	}
	if !resp.Node.Dir {
		return nil, fmt.Errorf("Node storage path %q in etcd is not a directory", key)
	}

	for _, node := range resp.Node.Nodes {
		nodename := utils.Tail(node.Key)
		n, err := k.GetNode(ctx, podname, nodename)
		if err != nil {
			return nodes, err
		}
		nodes = append(nodes, n)
	}
	return nodes, err
}

// GetAllNodes get all nodes from etcd
// any error will break and return immediately
func (k *Krypton) GetAllNodes(ctx context.Context) (nodes []*types.Node, err error) {
	pods, err := k.GetAllPods(ctx)
	if err != nil {
		return nodes, err
	}

	for _, pod := range pods {
		ns, err := k.GetNodesByPod(ctx, pod.Name)
		if err != nil {
			return nodes, err
		}
		nodes = append(nodes, ns...)
	}
	return nodes, err
}

// UpdateNode update a node, save it to etcd
// storage path in etcd is `/pod/:podname/node/:nodename/info`
func (k *Krypton) UpdateNode(ctx context.Context, node *types.Node) error {
	lock, err := k.CreateLock(fmt.Sprintf("%s_%s", node.Podname, node.Name), k.config.LockTimeout)
	if err != nil {
		return err
	}

	if err := lock.Lock(ctx); err != nil {
		return err
	}
	defer lock.Unlock(ctx)

	key := fmt.Sprintf(nodeInfoKey, node.Podname, node.Name)
	bytes, err := json.Marshal(node)
	if err != nil {
		return err
	}

	_, err = k.etcd.Set(ctx, key, string(bytes), nil)
	return err
}

// UpdateNodeResource update cpu and mem on a node, either add or substract
// need to lock
func (k *Krypton) UpdateNodeResource(ctx context.Context, podname, nodename string, cpu types.CPUMap, mem int64, action string) error {
	lock, err := k.CreateLock(fmt.Sprintf("%s_%s", podname, nodename), k.config.LockTimeout)
	if err != nil {
		return err
	}

	if err := lock.Lock(ctx); err != nil {
		return err
	}
	defer lock.Unlock(ctx)

	nodeKey := fmt.Sprintf(nodeInfoKey, podname, nodename)
	resp, err := k.etcd.Get(ctx, nodeKey, nil)
	if err != nil {
		return err
	}
	if resp.Node.Dir {
		return fmt.Errorf("Node storage path %q in etcd is a directory", nodeKey)
	}

	node := &types.Node{}
	if err := json.Unmarshal([]byte(resp.Node.Value), node); err != nil {
		return err
	}

	if action == "add" || action == "+" {
		node.CPU.Add(cpu)
		node.MemCap += mem
	} else if action == "sub" || action == "-" {
		node.CPU.Sub(cpu)
		node.MemCap -= mem
	}

	bytes, err := json.Marshal(node)
	if err != nil {
		return err
	}

	log.Debugf("[UpdateNodeResource] new node info: %s", string(bytes))
	_, err = k.etcd.Set(ctx, nodeKey, string(bytes), nil)
	return err
}

func (k *Krypton) makeDockerClient(ctx context.Context, podname, nodename, endpoint string, force bool) (*engineapi.Client, error) {
	u, err := url.Parse(endpoint)
	if err != nil {
		return nil, err
	}

	host, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		return nil, err
	}

	// try get client, if nil, create a new one
	var client *engineapi.Client
	client = _cache.get(host)
	if client == nil || force {
		if k.config.Docker.CertPath == "" {
			client, err = makeRawClient(endpoint, k.config.Docker.APIVersion)
		} else {
			keyFormats := []string{nodeCaKey, nodeCertKey, nodeKeyKey}
			data := []string{}
			for i := 0; i < 3; i++ {
				resp, err := k.etcd.Get(ctx, fmt.Sprintf(keyFormats[i], podname, nodename), nil)
				if err != nil {
					return nil, err
				}
				data = append(data, resp.Node.Value)
			}
			caFile, err := ioutil.TempFile(k.config.Docker.CertPath, fmt.Sprintf("ca-%s", host))
			if err != nil {
				return nil, err
			}
			certFile, err := ioutil.TempFile(k.config.Docker.CertPath, fmt.Sprintf("cert-%s", host))
			if err != nil {
				return nil, err
			}
			keyFile, err := ioutil.TempFile(k.config.Docker.CertPath, fmt.Sprintf("key-%s", host))
			if err != nil {
				return nil, err
			}
			if err := dumpFromString(caFile, certFile, keyFile, data[0], data[1], data[2]); err != nil {
				return nil, err
			}
			client, err = makeRawClientWithTLS(caFile, certFile, keyFile, endpoint, k.config.Docker.APIVersion)
		}
		if err != nil {
			return nil, err
		}
		_cache.set(host, client)
	}
	return client, nil
}
<commit_msg>set default cpu and memory when add node<commit_after>package etcdstore

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/url"
	"strconv"
	"strings"

	etcdclient "github.com/coreos/etcd/client"
	engineapi "github.com/docker/docker/client"
	"github.com/projecteru2/core/types"
	"github.com/projecteru2/core/utils"
	log "github.com/sirupsen/logrus"
)

// AddNode save it to etcd
// storage path in etcd is `/pod/:podname/node/:nodename/info`
func (k *Krypton) AddNode(ctx context.Context, name, endpoint, podname, ca, cert, key string, cpu int, share, memory int64, labels map[string]string) (*types.Node, error) {
	if !strings.HasPrefix(endpoint, "tcp://") {
		return nil, fmt.Errorf("Endpoint must starts with tcp:// %q", endpoint)
	}

	_, err := k.GetPod(ctx, podname)
	if err != nil {
		return nil, err
	}

	nodeKey := fmt.Sprintf(nodeInfoKey, podname, name)
	if _, err := k.etcd.Get(ctx, nodeKey, nil); err == nil {
		return nil, fmt.Errorf("Node (%s, %s) already exists", podname, name)
	}

	// 如果有tls的证书需要保存就保存一下
	if ca != "" && cert != "" && key != "" {
		_, err = k.etcd.Set(ctx, fmt.Sprintf(nodeCaKey, podname, name), ca, nil)
		if err != nil {
			return nil, err
		}
		_, err = k.etcd.Set(ctx, fmt.Sprintf(nodeCertKey, podname, name), cert, nil)
		if err != nil {
			return nil, err
		}
		_, err = k.etcd.Set(ctx, fmt.Sprintf(nodeKeyKey, podname, name), key, nil)
		if err != nil {
			return nil, err
		}
	}

	// 尝试加载docker的客户端
	engine, err := k.makeDockerClient(ctx, podname, name, endpoint, true)
	if err != nil {
		k.deleteNode(ctx, podname, name, endpoint)
		return nil, err
	}

	// 判断这货是不是活着的
	info, err := engine.Info(ctx)
	if err != nil {
		k.deleteNode(ctx, podname, name, endpoint)
		return nil, err
	}

	ncpu := cpu
	memcap := memory
	if cpu == 0 {
		ncpu = info.NCPU
	}
	if memory == 0 {
		memcap = info.MemTotal - gigabyte
	}
	if share == 0 {
		share = k.config.Scheduler.ShareBase
	}
	cpumap := types.CPUMap{}
	for i := 0; i < ncpu; i++ {
		cpumap[strconv.Itoa(i)] = share
	}

	node := &types.Node{
		Name:      name,
		Endpoint:  endpoint,
		Podname:   podname,
		CPU:       cpumap,
		MemCap:    memcap,
		Available: true,
		Labels:    labels,
	}

	bytes, err := json.Marshal(node)
	if err != nil {
		k.deleteNode(ctx, podname, name, endpoint)
		return nil, err
	}

	_, err = k.etcd.Create(ctx, nodeKey, string(bytes))
	if err != nil {
		k.deleteNode(ctx, podname, name, endpoint)
		return nil, err
	}

	_, err = k.etcd.Create(ctx, fmt.Sprintf(nodePodKey, name), podname)
	if err != nil {
		k.deleteNode(ctx, podname, name, endpoint)
		return nil, err
	}

	return node, nil
}

// 因为是先写etcd的证书再拿client
// 所以可能出现实际上node创建失败但是却写好了证书的情况
// 所以需要删除这些留存的证书
// 至于结果是不是成功就无所谓了
func (k *Krypton) deleteNode(ctx context.Context, podname, nodename, endpoint string) {
	key := fmt.Sprintf(nodePrefixKey, podname, nodename)
	k.etcd.Delete(ctx, key, &etcdclient.DeleteOptions{Recursive: true})
	k.etcd.Delete(ctx, fmt.Sprintf(nodePodKey, nodename), &etcdclient.DeleteOptions{})
	u, err := url.Parse(endpoint)
	if err != nil {
		log.Errorf("[deleteNode] Bad endpoint: %s", endpoint)
		return
	}

	host, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		log.Errorf("[deleteNode] Bad addr: %s", u.Host)
		return
	}
	_cache.delete(host)
	log.Debugf("[deleteNode] Node (%s, %s, %s) deleted", podname, nodename, endpoint)
}

// DeleteNode delete a node
func (k *Krypton) DeleteNode(ctx context.Context, node *types.Node) {
	k.deleteNode(ctx, node.Podname, node.Name, node.Endpoint)
}

// GetNode get a node from etcd
// and construct it's docker client
// a node must belong to a pod
// and since node is not the smallest unit to user, to get a node we must specify the corresponding pod
// storage path in etcd is `/pod/:podname/node/:nodename/info`
func (k *Krypton) GetNode(ctx context.Context, podname, nodename string) (*types.Node, error) {
	key := fmt.Sprintf(nodeInfoKey, podname, nodename)
	resp, err := k.etcd.Get(ctx, key, nil)
	if err != nil {
		return nil, err
	}
	if resp.Node.Dir {
		return nil, fmt.Errorf("Node storage path %q in etcd is a directory", key)
	}

	node := &types.Node{}
	if err := json.Unmarshal([]byte(resp.Node.Value), node); err != nil {
		return nil, err
	}

	engine, err := k.makeDockerClient(ctx, podname, nodename, node.Endpoint, false)
	if err != nil {
		return nil, err
	}

	node.Engine = engine
	return node, nil
}

// GetNodeByName get node by name
func (k *Krypton) GetNodeByName(ctx context.Context, nodename string) (node *types.Node, err error) {
	key := fmt.Sprintf(nodePodKey, nodename)
	resp, err := k.etcd.Get(ctx, key, nil)
	if err != nil {
		return nil, err
	}
	podname := resp.Node.Value
	return k.GetNode(ctx, podname, nodename)
}

// GetNodesByPod get all nodes bound to pod
// here we use podname instead of pod instance
// storage path in etcd is `/pod/:podname/node`
func (k *Krypton) GetNodesByPod(ctx context.Context, podname string) (nodes []*types.Node, err error) {
	key := fmt.Sprintf(podNodesKey, podname)
	resp, err := k.etcd.Get(ctx, key, nil)
	if err != nil {
		return nodes, err
	}
	if !resp.Node.Dir {
		return nil, fmt.Errorf("Node storage path %q in etcd is not a directory", key)
	}

	for _, node := range resp.Node.Nodes {
		nodename := utils.Tail(node.Key)
		n, err := k.GetNode(ctx, podname, nodename)
		if err != nil {
			return nodes, err
		}
		nodes = append(nodes, n)
	}
	return nodes, err
}

// GetAllNodes get all nodes from etcd
// any error will break and return immediately
func (k *Krypton) GetAllNodes(ctx context.Context) (nodes []*types.Node, err error) {
	pods, err := k.GetAllPods(ctx)
	if err != nil {
		return nodes, err
	}

	for _, pod := range pods {
		ns, err := k.GetNodesByPod(ctx, pod.Name)
		if err != nil {
			return nodes, err
		}
		nodes = append(nodes, ns...)
	}
	return nodes, err
}

// UpdateNode update a node, save it to etcd
// storage path in etcd is `/pod/:podname/node/:nodename/info`
func (k *Krypton) UpdateNode(ctx context.Context, node *types.Node) error {
	lock, err := k.CreateLock(fmt.Sprintf("%s_%s", node.Podname, node.Name), k.config.LockTimeout)
	if err != nil {
		return err
	}

	if err := lock.Lock(ctx); err != nil {
		return err
	}
	defer lock.Unlock(ctx)

	key := fmt.Sprintf(nodeInfoKey, node.Podname, node.Name)
	bytes, err := json.Marshal(node)
	if err != nil {
		return err
	}

	_, err = k.etcd.Set(ctx, key, string(bytes), nil)
	return err
}

// UpdateNodeResource update cpu and mem on a node, either add or substract
// need to lock
func (k *Krypton) UpdateNodeResource(ctx context.Context, podname, nodename string, cpu types.CPUMap, mem int64, action string) error {
	lock, err := k.CreateLock(fmt.Sprintf("%s_%s", podname, nodename), k.config.LockTimeout)
	if err != nil {
		return err
	}

	if err := lock.Lock(ctx); err != nil {
		return err
	}
	defer lock.Unlock(ctx)

	nodeKey := fmt.Sprintf(nodeInfoKey, podname, nodename)
	resp, err := k.etcd.Get(ctx, nodeKey, nil)
	if err != nil {
		return err
	}
	if resp.Node.Dir {
		return fmt.Errorf("Node storage path %q in etcd is a directory", nodeKey)
	}

	node := &types.Node{}
	if err := json.Unmarshal([]byte(resp.Node.Value), node); err != nil {
		return err
	}

	if action == "add" || action == "+" {
		node.CPU.Add(cpu)
		node.MemCap += mem
	} else if action == "sub" || action == "-" {
		node.CPU.Sub(cpu)
		node.MemCap -= mem
	}

	bytes, err := json.Marshal(node)
	if err != nil {
		return err
	}

	log.Debugf("[UpdateNodeResource] new node info: %s", string(bytes))
	_, err = k.etcd.Set(ctx, nodeKey, string(bytes), nil)
	return err
}

func (k *Krypton) makeDockerClient(ctx context.Context, podname, nodename, endpoint string, force bool) (*engineapi.Client, error) {
	u, err := url.Parse(endpoint)
	if err != nil {
		return nil, err
	}

	host, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		return nil, err
	}

	// try get client, if nil, create a new one
	var client *engineapi.Client
	client = _cache.get(host)
	if client == nil || force {
		if k.config.Docker.CertPath == "" {
			client, err = makeRawClient(endpoint, k.config.Docker.APIVersion)
		} else {
			keyFormats := []string{nodeCaKey, nodeCertKey, nodeKeyKey}
			data := []string{}
			for i := 0; i < 3; i++ {
				resp, err := k.etcd.Get(ctx, fmt.Sprintf(keyFormats[i], podname, nodename), nil)
				if err != nil {
					return nil, err
				}
				data = append(data, resp.Node.Value)
			}
			caFile, err := ioutil.TempFile(k.config.Docker.CertPath, fmt.Sprintf("ca-%s", host))
			if err != nil {
				return nil, err
			}
			certFile, err := ioutil.TempFile(k.config.Docker.CertPath, fmt.Sprintf("cert-%s", host))
			if err != nil {
				return nil, err
			}
			keyFile, err := ioutil.TempFile(k.config.Docker.CertPath, fmt.Sprintf("key-%s", host))
			if err != nil {
				return nil, err
			}
			if err := dumpFromString(caFile, certFile, keyFile, data[0], data[1], data[2]); err != nil {
				return nil, err
			}
			client, err = makeRawClientWithTLS(caFile, certFile, keyFile, endpoint, k.config.Docker.APIVersion)
		}
		if err != nil {
			return nil, err
		}
		_cache.set(host, client)
	}
	return client, nil
}
<|endoftext|>