<commit_before>// Copyright 2016-2020 Authors of Cilium
//
// 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 watchers

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/cilium/cilium/pkg/controller"
	"github.com/cilium/cilium/pkg/endpoint"
	"github.com/cilium/cilium/pkg/k8s"
	cilium_v2 "github.com/cilium/cilium/pkg/k8s/apis/cilium.io/v2"
	k8sversion "github.com/cilium/cilium/pkg/k8s/version"
	pkgLabels "github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/option"

	go_version "github.com/blang/semver"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
)

const (
	// subsysEndpointSync is the value for logfields.LogSubsys
	subsysEndpointSync = "endpointsynchronizer"
)

// EndpointSynchronizer currently is an empty type, which wraps around syncing
// of CiliumEndpoint resources.
type EndpointSynchronizer struct{}

// RunK8sCiliumEndpointSync starts a controller that synchronizes the endpoint
// to the corresponding k8s CiliumEndpoint CRD. It is expected that each CEP
// has 1 controller that updates it, and a local copy is retained and only
// updates are pushed up.
// CiliumEndpoint objects have the same name as the pod they represent.
func (epSync *EndpointSynchronizer) RunK8sCiliumEndpointSync(e *endpoint.Endpoint, conf endpoint.EndpointStatusConfiguration) {
	var (
		endpointID     = e.ID
		controllerName = fmt.Sprintf("sync-to-k8s-ciliumendpoint (%v)", endpointID)
		scopedLog      = e.Logger(subsysEndpointSync).WithField("controller", controllerName)
	)

	if option.Config.DisableCiliumEndpointCRD {
		scopedLog.Warn("Not running controller. CEP CRD synchronization is disabled")
		return
	}

	if !k8s.IsEnabled() {
		scopedLog.Debug("Not starting controller because k8s is disabled")
		return
	}

	ciliumClient := k8s.CiliumClient().CiliumV2()

	// The health endpoint doesn't really exist in k8s and updates to it caused
	// arbitrary errors. Disable the controller for these endpoints.
	if isHealthEP := e.HasLabels(pkgLabels.LabelHealth); isHealthEP {
		scopedLog.Debug("Not starting unnecessary CEP controller for cilium-health endpoint")
		return
	}

	var (
		lastMdl  *cilium_v2.EndpointStatus
		localCEP *cilium_v2.CiliumEndpoint // the local copy of the CEP object. Reused.
		needInit = true                    // needInit indicates that we may need to create the CEP
	)

	// NOTE: The controller functions do NOT hold the endpoint locks
	e.UpdateController(controllerName,
		controller.ControllerParams{
			RunInterval: 10 * time.Second,
			DoFunc: func(ctx context.Context) (err error) {
				// Update logger as scopeLog might not have the podName when it
				// was created.
				scopedLog = e.Logger(subsysEndpointSync).WithField("controller", controllerName)

				if k8sversion.Version().Equals(go_version.Version{}) {
					return fmt.Errorf("Kubernetes apiserver is not available")
				}

				capabilities := k8sversion.Capabilities()

				// K8sPodName and K8sNamespace are not always available when an
				// endpoint is first created, so we collect them here.
				podName := e.GetK8sPodName()
				if podName == "" {
					scopedLog.Debug("Skipping CiliumEndpoint update because it has no k8s pod name")
					return nil
				}

				namespace := e.GetK8sNamespace()
				if namespace == "" {
					scopedLog.Debug("Skipping CiliumEndpoint update because it has no k8s namespace")
					return nil
				}

				// Serialize the endpoint into a model. It is compared with the one
				// from before, only updating on changes.
				mdl := e.GetCiliumEndpointStatus(conf)
				if mdl.DeepEqual(lastMdl) {
					scopedLog.Debug("Skipping CiliumEndpoint update because it has not changed")
					return nil
				}

				// Initialize the CEP by deleting the upstream instance and recreating
				// it. Deleting first allows for upgrade scenarios where the format has
				// changed but our k8s CEP code cannot read in the upstream value.
				if needInit {
					state := e.GetState()
					// Don't bother to create if the
					// endpoint is already disconnecting
					if state == endpoint.StateDisconnecting ||
						state == endpoint.StateDisconnected {
						return nil
					}

					scopedLog.Debug("Getting CEP during an initialization")
					localCEP, err = ciliumClient.CiliumEndpoints(namespace).Get(ctx, podName, meta_v1.GetOptions{})
					// It's only an error if it exists but something else happened
					switch {
					case k8serrors.IsNotFound(err):
						pod := e.GetPod()
						if pod == nil {
							scopedLog.Debug("Skipping CiliumEndpoint update because it has no k8s pod")
							return nil
						}

						// We can't create localCEP directly, it must come from the k8s
						// server via an API call.
						cep := &cilium_v2.CiliumEndpoint{
							ObjectMeta: meta_v1.ObjectMeta{
								Name: podName,
								OwnerReferences: []meta_v1.OwnerReference{
									{
										APIVersion: "v1",
										Kind:       "Pod",
										Name:       pod.GetObjectMeta().GetName(),
										UID:        types.UID(pod.ObjectMeta.UID),
									},
								},
							},
						}
						if mdl != nil {
							cep.Status = *mdl
						}
						localCEP, err = ciliumClient.CiliumEndpoints(namespace).Create(ctx, cep, meta_v1.CreateOptions{})
						if err != nil {
							scopedLog.WithError(err).Error("Cannot create CEP")
							return err
						}

						// We have successfully created the CEP and can return. Subsequent
						// runs will update using localCEP.
						needInit = false
						return nil
					case err != nil:
						scopedLog.WithError(err).Warn("Error getting CEP")
						return err
					default:

					}
				}

				// We have no localCEP copy. We need to fetch it for updates, below.
				// This is unexpected as there should be only 1 writer per CEP, this
				// controller, and the localCEP created on startup will be used.
				if localCEP == nil {
					localCEP, err = ciliumClient.CiliumEndpoints(namespace).Get(ctx, podName, meta_v1.GetOptions{})
					switch {
					// The CEP doesn't exist in k8s. This is unexpetected but may occur
					// if the endpoint was removed from k8s but not yet within the agent.
					// Mark the CEP for creation on the next controller iteration. This
					// may never occur if the controller is stopped on Endpoint delete.
					case err != nil && k8serrors.IsNotFound(err):
						needInit = true
						return err

					// We cannot read the upstream CEP. needInit will cause the next
					// iteration to delete and create the CEP. This is an unexpected
					// situation.
					case err != nil && k8serrors.IsInvalid(err):
						scopedLog.WithError(err).Warn("Invalid CEP during update")
						needInit = true
						return nil

					// A real error
					case err != nil:
						scopedLog.WithError(err).Error("Cannot get CEP during update")
						return err
					}
				}

				switch {
				case capabilities.Patch:
					// For json patch we don't need to perform a GET for endpoints

					// If it fails it means the test from the previous patch failed
					// so we can safely replace this node in the CNP status.
					replaceCEPStatus := []k8s.JSONPatch{
						{
							OP:    "replace",
							Path:  "/status",
							Value: mdl,
						},
					}
					var createStatusPatch []byte
					createStatusPatch, err = json.Marshal(replaceCEPStatus)
					if err != nil {
						return err
					}
					localCEP, err = ciliumClient.CiliumEndpoints(namespace).Patch(
						ctx, podName,
						types.JSONPatchType,
						createStatusPatch,
						meta_v1.PatchOptions{},
						"status")
				default:
					// We have an object to reuse. Update and push it up. In the case of an
					// update error, we retry in the next iteration of the controller using
					// the copy returned by Update.
					scopedLog.Debug("Updating CEP from local copy")
					mdl.DeepCopyInto(&localCEP.Status)
					switch {
					case capabilities.UpdateStatus:
						localCEP, err = ciliumClient.CiliumEndpoints(namespace).UpdateStatus(ctx, localCEP, meta_v1.UpdateOptions{})
					default:
						localCEP, err = ciliumClient.CiliumEndpoints(namespace).Update(ctx, localCEP, meta_v1.UpdateOptions{})
					}
				}

				// Handle Update errors or return successfully
				switch {
				// Return no error when we see a conflict. We want to retry without a
				// backoff and the Update* calls returned the current localCEP
				case err != nil && k8serrors.IsConflict(err):
					scopedLog.WithError(err).Warn("Cannot update CEP due to a revision conflict. The next controller execution will try again")
					needInit = true
					return nil

				// Ensure we re-init when we see a generic error. This will recrate the
				// CEP.
				case err != nil:
					scopedLog.WithError(err).Error("Cannot update CEP")
					needInit = true
					return err

				// A successful update means no more updates unless mdl changes
				default:
					lastMdl = mdl
					return nil
				}
			},
			StopFunc: func(ctx context.Context) error {
				podName := e.GetK8sPodName()
				if podName == "" {
					scopedLog.Debug("Skipping CiliumEndpoint deletion because it has no k8s pod name")
					return nil
				}
				namespace := e.GetK8sNamespace()
				if namespace == "" {
					scopedLog.Debug("Skipping CiliumEndpoint deletion because it has no k8s namespace")
					return nil
				}
				if err := ciliumClient.CiliumEndpoints(namespace).Delete(ctx, podName, meta_v1.DeleteOptions{}); err != nil {
					if !k8serrors.IsNotFound(err) {
						scopedLog.WithError(err).Warning("Unable to delete CEP")
					}
				}
				return nil
			},
		})
}
<commit_msg>k8s/watchers: set BlockOwnerDeletion to true in CEP<commit_after>// Copyright 2016-2020 Authors of Cilium
//
// 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 watchers

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/cilium/cilium/pkg/controller"
	"github.com/cilium/cilium/pkg/endpoint"
	"github.com/cilium/cilium/pkg/k8s"
	cilium_v2 "github.com/cilium/cilium/pkg/k8s/apis/cilium.io/v2"
	k8sversion "github.com/cilium/cilium/pkg/k8s/version"
	pkgLabels "github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/option"

	go_version "github.com/blang/semver"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
)

const (
	// subsysEndpointSync is the value for logfields.LogSubsys
	subsysEndpointSync = "endpointsynchronizer"
)

// EndpointSynchronizer currently is an empty type, which wraps around syncing
// of CiliumEndpoint resources.
type EndpointSynchronizer struct{}

// RunK8sCiliumEndpointSync starts a controller that synchronizes the endpoint
// to the corresponding k8s CiliumEndpoint CRD. It is expected that each CEP
// has 1 controller that updates it, and a local copy is retained and only
// updates are pushed up.
// CiliumEndpoint objects have the same name as the pod they represent.
func (epSync *EndpointSynchronizer) RunK8sCiliumEndpointSync(e *endpoint.Endpoint, conf endpoint.EndpointStatusConfiguration) {
	var (
		endpointID     = e.ID
		controllerName = fmt.Sprintf("sync-to-k8s-ciliumendpoint (%v)", endpointID)
		scopedLog      = e.Logger(subsysEndpointSync).WithField("controller", controllerName)
	)

	if option.Config.DisableCiliumEndpointCRD {
		scopedLog.Warn("Not running controller. CEP CRD synchronization is disabled")
		return
	}

	if !k8s.IsEnabled() {
		scopedLog.Debug("Not starting controller because k8s is disabled")
		return
	}

	ciliumClient := k8s.CiliumClient().CiliumV2()

	// The health endpoint doesn't really exist in k8s and updates to it caused
	// arbitrary errors. Disable the controller for these endpoints.
	if isHealthEP := e.HasLabels(pkgLabels.LabelHealth); isHealthEP {
		scopedLog.Debug("Not starting unnecessary CEP controller for cilium-health endpoint")
		return
	}

	var (
		lastMdl  *cilium_v2.EndpointStatus
		localCEP *cilium_v2.CiliumEndpoint // the local copy of the CEP object. Reused.
		needInit = true                    // needInit indicates that we may need to create the CEP
	)

	// NOTE: The controller functions do NOT hold the endpoint locks
	e.UpdateController(controllerName,
		controller.ControllerParams{
			RunInterval: 10 * time.Second,
			DoFunc: func(ctx context.Context) (err error) {
				// Update logger as scopeLog might not have the podName when it
				// was created.
				scopedLog = e.Logger(subsysEndpointSync).WithField("controller", controllerName)

				if k8sversion.Version().Equals(go_version.Version{}) {
					return fmt.Errorf("Kubernetes apiserver is not available")
				}

				capabilities := k8sversion.Capabilities()

				// K8sPodName and K8sNamespace are not always available when an
				// endpoint is first created, so we collect them here.
				podName := e.GetK8sPodName()
				if podName == "" {
					scopedLog.Debug("Skipping CiliumEndpoint update because it has no k8s pod name")
					return nil
				}

				namespace := e.GetK8sNamespace()
				if namespace == "" {
					scopedLog.Debug("Skipping CiliumEndpoint update because it has no k8s namespace")
					return nil
				}

				// Serialize the endpoint into a model. It is compared with the one
				// from before, only updating on changes.
				mdl := e.GetCiliumEndpointStatus(conf)
				if mdl.DeepEqual(lastMdl) {
					scopedLog.Debug("Skipping CiliumEndpoint update because it has not changed")
					return nil
				}

				// Initialize the CEP by deleting the upstream instance and recreating
				// it. Deleting first allows for upgrade scenarios where the format has
				// changed but our k8s CEP code cannot read in the upstream value.
				if needInit {
					state := e.GetState()
					// Don't bother to create if the
					// endpoint is already disconnecting
					if state == endpoint.StateDisconnecting ||
						state == endpoint.StateDisconnected {
						return nil
					}

					scopedLog.Debug("Getting CEP during an initialization")
					localCEP, err = ciliumClient.CiliumEndpoints(namespace).Get(ctx, podName, meta_v1.GetOptions{})
					// It's only an error if it exists but something else happened
					switch {
					case k8serrors.IsNotFound(err):
						pod := e.GetPod()
						if pod == nil {
							scopedLog.Debug("Skipping CiliumEndpoint update because it has no k8s pod")
							return nil
						}

						// We can't create localCEP directly, it must come from the k8s
						// server via an API call.
						cep := &cilium_v2.CiliumEndpoint{
							ObjectMeta: meta_v1.ObjectMeta{
								Name: podName,
								OwnerReferences: []meta_v1.OwnerReference{
									{
										APIVersion:         "v1",
										Kind:               "Pod",
										Name:               pod.GetObjectMeta().GetName(),
										UID:                pod.ObjectMeta.UID,
										BlockOwnerDeletion: func() *bool { a := true; return &a }(),
									},
								},
							},
						}
						if mdl != nil {
							cep.Status = *mdl
						}
						localCEP, err = ciliumClient.CiliumEndpoints(namespace).Create(ctx, cep, meta_v1.CreateOptions{})
						if err != nil {
							scopedLog.WithError(err).Error("Cannot create CEP")
							return err
						}

						// We have successfully created the CEP and can return. Subsequent
						// runs will update using localCEP.
						needInit = false
						return nil
					case err != nil:
						scopedLog.WithError(err).Warn("Error getting CEP")
						return err
					default:

					}
				}

				// We have no localCEP copy. We need to fetch it for updates, below.
				// This is unexpected as there should be only 1 writer per CEP, this
				// controller, and the localCEP created on startup will be used.
				if localCEP == nil {
					localCEP, err = ciliumClient.CiliumEndpoints(namespace).Get(ctx, podName, meta_v1.GetOptions{})
					switch {
					// The CEP doesn't exist in k8s. This is unexpetected but may occur
					// if the endpoint was removed from k8s but not yet within the agent.
					// Mark the CEP for creation on the next controller iteration. This
					// may never occur if the controller is stopped on Endpoint delete.
					case err != nil && k8serrors.IsNotFound(err):
						needInit = true
						return err

					// We cannot read the upstream CEP. needInit will cause the next
					// iteration to delete and create the CEP. This is an unexpected
					// situation.
					case err != nil && k8serrors.IsInvalid(err):
						scopedLog.WithError(err).Warn("Invalid CEP during update")
						needInit = true
						return nil

					// A real error
					case err != nil:
						scopedLog.WithError(err).Error("Cannot get CEP during update")
						return err
					}
				}

				switch {
				case capabilities.Patch:
					// For json patch we don't need to perform a GET for endpoints

					// If it fails it means the test from the previous patch failed
					// so we can safely replace this node in the CNP status.
					replaceCEPStatus := []k8s.JSONPatch{
						{
							OP:    "replace",
							Path:  "/status",
							Value: mdl,
						},
					}
					var createStatusPatch []byte
					createStatusPatch, err = json.Marshal(replaceCEPStatus)
					if err != nil {
						return err
					}
					localCEP, err = ciliumClient.CiliumEndpoints(namespace).Patch(
						ctx, podName,
						types.JSONPatchType,
						createStatusPatch,
						meta_v1.PatchOptions{},
						"status")
				default:
					// We have an object to reuse. Update and push it up. In the case of an
					// update error, we retry in the next iteration of the controller using
					// the copy returned by Update.
					scopedLog.Debug("Updating CEP from local copy")
					mdl.DeepCopyInto(&localCEP.Status)
					switch {
					case capabilities.UpdateStatus:
						localCEP, err = ciliumClient.CiliumEndpoints(namespace).UpdateStatus(ctx, localCEP, meta_v1.UpdateOptions{})
					default:
						localCEP, err = ciliumClient.CiliumEndpoints(namespace).Update(ctx, localCEP, meta_v1.UpdateOptions{})
					}
				}

				// Handle Update errors or return successfully
				switch {
				// Return no error when we see a conflict. We want to retry without a
				// backoff and the Update* calls returned the current localCEP
				case err != nil && k8serrors.IsConflict(err):
					scopedLog.WithError(err).Warn("Cannot update CEP due to a revision conflict. The next controller execution will try again")
					needInit = true
					return nil

				// Ensure we re-init when we see a generic error. This will recrate the
				// CEP.
				case err != nil:
					scopedLog.WithError(err).Error("Cannot update CEP")
					needInit = true
					return err

				// A successful update means no more updates unless mdl changes
				default:
					lastMdl = mdl
					return nil
				}
			},
			StopFunc: func(ctx context.Context) error {
				podName := e.GetK8sPodName()
				if podName == "" {
					scopedLog.Debug("Skipping CiliumEndpoint deletion because it has no k8s pod name")
					return nil
				}
				namespace := e.GetK8sNamespace()
				if namespace == "" {
					scopedLog.Debug("Skipping CiliumEndpoint deletion because it has no k8s namespace")
					return nil
				}
				if err := ciliumClient.CiliumEndpoints(namespace).Delete(ctx, podName, meta_v1.DeleteOptions{}); err != nil {
					if !k8serrors.IsNotFound(err) {
						scopedLog.WithError(err).Warning("Unable to delete CEP")
					}
				}
				return nil
			},
		})
}
<|endoftext|>


<commit_before>package notifiers

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/services/alerting"
	"github.com/grafana/grafana/pkg/setting"
)

func init() {
	alerting.RegisterNotifier(&alerting.NotifierPlugin{
		Type:        "slack",
		Name:        "Slack",
		Description: "Sends notifications to Slack via Slack Webhooks",
		Heading:     "Slack settings",
		Factory:     NewSlackNotifier,
		Options: []alerting.NotifierOption{
			{
				Label:        "Url",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Placeholder:  "Slack incoming webhook url",
				PropertyName: "url",
				Required:     true,
				Secure:       true,
			},
			{
				Label:        "Recipient",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Override default channel or user, use #channel-name, @username (has to be all lowercase, no whitespace), or user/channel Slack ID",
				PropertyName: "recipient",
			},
			{
				Label:        "Username",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Set the username for the bot's message",
				PropertyName: "username",
			},
			{
				Label:        "Icon emoji",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Provide an emoji to use as the icon for the bot's message. Overrides the icon URL.",
				PropertyName: "iconEmoji",
			},
			{
				Label:        "Icon URL",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Provide a URL to an image to use as the icon for the bot's message",
				PropertyName: "iconUrl",
			},
			{
				Label:        "Mention Users",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Mention one or more users (comma separated) when notifying in a channel, by ID (you can copy this from the user's Slack profile)",
				PropertyName: "mentionUsers",
			},
			{
				Label:        "Mention Groups",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Mention one or more groups (comma separated) when notifying in a channel (you can copy this from the group's Slack profile URL)",
				PropertyName: "mentionGroups",
			},
			{
				Label:   "Mention Channel",
				Element: alerting.ElementTypeSelect,
				SelectOptions: []alerting.SelectOption{
					{
						Value: "",
						Label: "Disabled",
					},
					{
						Value: "here",
						Label: "Every active channel member",
					},
					{
						Value: "channel",
						Label: "Every channel member",
					},
				},
				Description:  "Mention whole channel or just active members when notifying",
				PropertyName: "mentionChannel",
			},
			{
				Label:        "Token",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Provide a bot token to use the Slack file.upload API (starts with \"xoxb\"). Specify Recipient for this to work",
				PropertyName: "token",
				Secure:       true,
			},
		},
	})
}

var reRecipient *regexp.Regexp = regexp.MustCompile("^((@[a-z0-9][a-zA-Z0-9._-]*)|(#[^ .A-Z]{1,79})|([a-zA-Z0-9]+))$")

// NewSlackNotifier is the constructor for the Slack notifier
func NewSlackNotifier(model *models.AlertNotification) (alerting.Notifier, error) {
	url := model.DecryptedValue("url", model.Settings.Get("url").MustString())
	if url == "" {
		return nil, alerting.ValidationError{Reason: "Could not find url property in settings"}
	}

	recipient := strings.TrimSpace(model.Settings.Get("recipient").MustString())
	if recipient != "" && !reRecipient.MatchString(recipient) {
		return nil, alerting.ValidationError{Reason: fmt.Sprintf("Recipient on invalid format: %q", recipient)}
	}
	username := model.Settings.Get("username").MustString()
	iconEmoji := model.Settings.Get("icon_emoji").MustString()
	iconURL := model.Settings.Get("icon_url").MustString()
	mentionUsersStr := model.Settings.Get("mentionUsers").MustString()
	mentionGroupsStr := model.Settings.Get("mentionGroups").MustString()
	mentionChannel := model.Settings.Get("mentionChannel").MustString()
	token := model.DecryptedValue("token", model.Settings.Get("token").MustString())

	uploadImage := model.Settings.Get("uploadImage").MustBool(true)

	if mentionChannel != "" && mentionChannel != "here" && mentionChannel != "channel" {
		return nil, alerting.ValidationError{
			Reason: fmt.Sprintf("Invalid value for mentionChannel: %q", mentionChannel),
		}
	}
	mentionUsers := []string{}
	for _, u := range strings.Split(mentionUsersStr, ",") {
		u = strings.TrimSpace(u)
		if u != "" {
			mentionUsers = append(mentionUsers, u)
		}
	}
	mentionGroups := []string{}
	for _, g := range strings.Split(mentionGroupsStr, ",") {
		g = strings.TrimSpace(g)
		if g != "" {
			mentionGroups = append(mentionGroups, g)
		}
	}

	return &SlackNotifier{
		NotifierBase:   NewNotifierBase(model),
		URL:            url,
		Recipient:      recipient,
		Username:       username,
		IconEmoji:      iconEmoji,
		IconURL:        iconURL,
		MentionUsers:   mentionUsers,
		MentionGroups:  mentionGroups,
		MentionChannel: mentionChannel,
		Token:          token,
		Upload:         uploadImage,
		log:            log.New("alerting.notifier.slack"),
	}, nil
}

// SlackNotifier is responsible for sending
// alert notification to Slack.
type SlackNotifier struct {
	NotifierBase
	URL            string
	Recipient      string
	Username       string
	IconEmoji      string
	IconURL        string
	MentionUsers   []string
	MentionGroups  []string
	MentionChannel string
	Token          string
	Upload         bool
	log            log.Logger
}

// Notify send alert notification to Slack.
func (sn *SlackNotifier) Notify(evalContext *alerting.EvalContext) error {
	sn.log.Info("Executing slack notification", "ruleId", evalContext.Rule.ID, "notification", sn.Name)

	ruleURL, err := evalContext.GetRuleURL()
	if err != nil {
		sn.log.Error("Failed get rule link", "error", err)
		return err
	}

	fields := make([]map[string]interface{}, 0)
	fieldLimitCount := 4
	for index, evt := range evalContext.EvalMatches {
		fields = append(fields, map[string]interface{}{
			"title": evt.Metric,
			"value": evt.Value,
			"short": true,
		})
		if index > fieldLimitCount {
			break
		}
	}

	if evalContext.Error != nil {
		fields = append(fields, map[string]interface{}{
			"title": "Error message",
			"value": evalContext.Error.Error(),
			"short": false,
		})
	}

	mentionsBuilder := strings.Builder{}
	appendSpace := func() {
		if mentionsBuilder.Len() > 0 {
			mentionsBuilder.WriteString(" ")
		}
	}
	mentionChannel := strings.TrimSpace(sn.MentionChannel)
	if mentionChannel != "" {
		mentionsBuilder.WriteString(fmt.Sprintf("<!%s|%s>", mentionChannel, mentionChannel))
	}
	if len(sn.MentionGroups) > 0 {
		appendSpace()
		for _, g := range sn.MentionGroups {
			mentionsBuilder.WriteString(fmt.Sprintf("<!subteam^%s>", g))
		}
	}
	if len(sn.MentionUsers) > 0 {
		appendSpace()
		for _, u := range sn.MentionUsers {
			mentionsBuilder.WriteString(fmt.Sprintf("<@%s>", u))
		}
	}
	msg := ""
	if evalContext.Rule.State != models.AlertStateOK { // don't add message when going back to alert state ok.
		msg = evalContext.Rule.Message
	}
	imageURL := ""
	// default to file.upload API method if a token is provided
	if sn.Token == "" {
		imageURL = evalContext.ImagePublicURL
	}

	var blocks []map[string]interface{}
	if mentionsBuilder.Len() > 0 {
		blocks = []map[string]interface{}{
			{
				"type": "section",
				"text": map[string]interface{}{
					"type": "mrkdwn",
					"text": mentionsBuilder.String(),
				},
			},
		}
	}
	attachment := map[string]interface{}{
		"color":       evalContext.GetStateModel().Color,
		"title":       evalContext.GetNotificationTitle(),
		"title_link":  ruleURL,
		"text":        msg,
		"fallback":    evalContext.GetNotificationTitle(),
		"fields":      fields,
		"footer":      "Grafana v" + setting.BuildVersion,
		"footer_icon": "https://grafana.com/assets/img/fav32.png",
		"ts":          time.Now().Unix(),
	}
	if sn.NeedsImage() && imageURL != "" {
		attachment["image_url"] = imageURL
	}
	body := map[string]interface{}{
		"text": evalContext.GetNotificationTitle(),
		"attachments": []map[string]interface{}{
			attachment,
		},
		"parse": "full", // to linkify urls, users and channels in alert message.
	}
	if len(blocks) > 0 {
		body["blocks"] = blocks
	}

	// recipient override
	if sn.Recipient != "" {
		body["channel"] = sn.Recipient
	}
	if sn.Username != "" {
		body["username"] = sn.Username
	}
	if sn.IconEmoji != "" {
		body["icon_emoji"] = sn.IconEmoji
	}
	if sn.IconURL != "" {
		body["icon_url"] = sn.IconURL
	}
	data, err := json.Marshal(&body)
	if err != nil {
		return err
	}

	cmd := &models.SendWebhookSync{
		Url:        sn.URL,
		Body:       string(data),
		HttpMethod: http.MethodPost,
	}
	if sn.Token != "" {
		sn.log.Debug("Adding authorization header to HTTP request")
		cmd.HttpHeader = map[string]string{
			"Authorization": fmt.Sprintf("Bearer %s", sn.Token),
		}
	}
	if err := bus.DispatchCtx(evalContext.Ctx, cmd); err != nil {
		sn.log.Error("Failed to send slack notification", "error", err, "webhook", sn.Name)
		return err
	}
	if sn.Token != "" && sn.UploadImage {
		err = sn.slackFileUpload(evalContext, sn.log, "https://slack.com/api/files.upload", sn.Recipient, sn.Token)
		if err != nil {
			return err
		}
	}
	return nil
}

func (sn *SlackNotifier) slackFileUpload(evalContext *alerting.EvalContext, log log.Logger, url string, recipient string, token string) error {
	if evalContext.ImageOnDiskPath == "" {
		// nolint:gosec
		// We can ignore the gosec G304 warning on this one because `setting.HomePath` comes from Grafana's configuration file.
		evalContext.ImageOnDiskPath = filepath.Join(setting.HomePath, "public/img/mixed_styles.png")
	}
	log.Info("Uploading to slack via file.upload API")
	headers, uploadBody, err := sn.generateSlackBody(evalContext.ImageOnDiskPath, token, recipient)
	if err != nil {
		return err
	}
	cmd := &models.SendWebhookSync{Url: url, Body: uploadBody.String(), HttpHeader: headers, HttpMethod: "POST"}
	if err := bus.DispatchCtx(evalContext.Ctx, cmd); err != nil {
		log.Error("Failed to upload slack image", "error", err, "webhook", "file.upload")
		return err
	}
	return nil
}

func (sn *SlackNotifier) generateSlackBody(path string, token string, recipient string) (map[string]string, bytes.Buffer, error) {
	// Slack requires all POSTs to files.upload to present
	// an "application/x-www-form-urlencoded" encoded querystring
	// See https://api.slack.com/methods/files.upload
	var b bytes.Buffer
	w := multipart.NewWriter(&b)
	defer func() {
		if err := w.Close(); err != nil {
			// Shouldn't matter since we already close w explicitly on the non-error path
			sn.log.Warn("Failed to close multipart writer", "err", err)
		}
	}()

	// Add the generated image file
	// We can ignore the gosec G304 warning on this one because `imagePath` comes
	// from the alert `evalContext` that generates the images. `evalContext` in turn derives the root of the file
	// path from configuration variables.
	// nolint:gosec
	f, err := os.Open(path)
	if err != nil {
		return nil, b, err
	}
	defer func() {
		if err := f.Close(); err != nil {
			sn.log.Warn("Failed to close file", "path", path, "err", err)
		}
	}()
	fw, err := w.CreateFormFile("file", path)
	if err != nil {
		return nil, b, err
	}
	if _, err := io.Copy(fw, f); err != nil {
		return nil, b, err
	}
	// Add the authorization token
	if err := w.WriteField("token", token); err != nil {
		return nil, b, err
	}
	// Add the channel(s) to POST to
	if err := w.WriteField("channels", recipient); err != nil {
		return nil, b, err
	}
	if err := w.Close(); err != nil {
		return nil, b, fmt.Errorf("failed to close multipart writer: %w", err)
	}
	headers := map[string]string{
		"Content-Type":  w.FormDataContentType(),
		"Authorization": "auth_token=\"" + token + "\"",
	}
	return headers, b, nil
}
<commit_msg>Dashboard: Fix issue where Slack notifications won't link to users (#32861)<commit_after>package notifiers

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/services/alerting"
	"github.com/grafana/grafana/pkg/setting"
)

func init() {
	alerting.RegisterNotifier(&alerting.NotifierPlugin{
		Type:        "slack",
		Name:        "Slack",
		Description: "Sends notifications to Slack via Slack Webhooks",
		Heading:     "Slack settings",
		Factory:     NewSlackNotifier,
		Options: []alerting.NotifierOption{
			{
				Label:        "Url",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Placeholder:  "Slack incoming webhook url",
				PropertyName: "url",
				Required:     true,
				Secure:       true,
			},
			{
				Label:        "Recipient",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Override default channel or user, use #channel-name, @username (has to be all lowercase, no whitespace), or user/channel Slack ID",
				PropertyName: "recipient",
			},
			{
				Label:        "Username",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Set the username for the bot's message",
				PropertyName: "username",
			},
			{
				Label:        "Icon emoji",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Provide an emoji to use as the icon for the bot's message. Overrides the icon URL.",
				PropertyName: "iconEmoji",
			},
			{
				Label:        "Icon URL",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Provide a URL to an image to use as the icon for the bot's message",
				PropertyName: "iconUrl",
			},
			{
				Label:        "Mention Users",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Mention one or more users (comma separated) when notifying in a channel, by ID (you can copy this from the user's Slack profile)",
				PropertyName: "mentionUsers",
			},
			{
				Label:        "Mention Groups",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Mention one or more groups (comma separated) when notifying in a channel (you can copy this from the group's Slack profile URL)",
				PropertyName: "mentionGroups",
			},
			{
				Label:   "Mention Channel",
				Element: alerting.ElementTypeSelect,
				SelectOptions: []alerting.SelectOption{
					{
						Value: "",
						Label: "Disabled",
					},
					{
						Value: "here",
						Label: "Every active channel member",
					},
					{
						Value: "channel",
						Label: "Every channel member",
					},
				},
				Description:  "Mention whole channel or just active members when notifying",
				PropertyName: "mentionChannel",
			},
			{
				Label:        "Token",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Description:  "Provide a bot token to use the Slack file.upload API (starts with \"xoxb\"). Specify Recipient for this to work",
				PropertyName: "token",
				Secure:       true,
			},
		},
	})
}

var reRecipient *regexp.Regexp = regexp.MustCompile("^((@[a-z0-9][a-zA-Z0-9._-]*)|(#[^ .A-Z]{1,79})|([a-zA-Z0-9]+))$")

// NewSlackNotifier is the constructor for the Slack notifier
func NewSlackNotifier(model *models.AlertNotification) (alerting.Notifier, error) {
	url := model.DecryptedValue("url", model.Settings.Get("url").MustString())
	if url == "" {
		return nil, alerting.ValidationError{Reason: "Could not find url property in settings"}
	}

	recipient := strings.TrimSpace(model.Settings.Get("recipient").MustString())
	if recipient != "" && !reRecipient.MatchString(recipient) {
		return nil, alerting.ValidationError{Reason: fmt.Sprintf("Recipient on invalid format: %q", recipient)}
	}
	username := model.Settings.Get("username").MustString()
	iconEmoji := model.Settings.Get("icon_emoji").MustString()
	iconURL := model.Settings.Get("icon_url").MustString()
	mentionUsersStr := model.Settings.Get("mentionUsers").MustString()
	mentionGroupsStr := model.Settings.Get("mentionGroups").MustString()
	mentionChannel := model.Settings.Get("mentionChannel").MustString()
	token := model.DecryptedValue("token", model.Settings.Get("token").MustString())

	uploadImage := model.Settings.Get("uploadImage").MustBool(true)

	if mentionChannel != "" && mentionChannel != "here" && mentionChannel != "channel" {
		return nil, alerting.ValidationError{
			Reason: fmt.Sprintf("Invalid value for mentionChannel: %q", mentionChannel),
		}
	}
	mentionUsers := []string{}
	for _, u := range strings.Split(mentionUsersStr, ",") {
		u = strings.TrimSpace(u)
		if u != "" {
			mentionUsers = append(mentionUsers, u)
		}
	}
	mentionGroups := []string{}
	for _, g := range strings.Split(mentionGroupsStr, ",") {
		g = strings.TrimSpace(g)
		if g != "" {
			mentionGroups = append(mentionGroups, g)
		}
	}

	return &SlackNotifier{
		NotifierBase:   NewNotifierBase(model),
		URL:            url,
		Recipient:      recipient,
		Username:       username,
		IconEmoji:      iconEmoji,
		IconURL:        iconURL,
		MentionUsers:   mentionUsers,
		MentionGroups:  mentionGroups,
		MentionChannel: mentionChannel,
		Token:          token,
		Upload:         uploadImage,
		log:            log.New("alerting.notifier.slack"),
	}, nil
}

// SlackNotifier is responsible for sending
// alert notification to Slack.
type SlackNotifier struct {
	NotifierBase
	URL            string
	Recipient      string
	Username       string
	IconEmoji      string
	IconURL        string
	MentionUsers   []string
	MentionGroups  []string
	MentionChannel string
	Token          string
	Upload         bool
	log            log.Logger
}

// Notify send alert notification to Slack.
func (sn *SlackNotifier) Notify(evalContext *alerting.EvalContext) error {
	sn.log.Info("Executing slack notification", "ruleId", evalContext.Rule.ID, "notification", sn.Name)

	ruleURL, err := evalContext.GetRuleURL()
	if err != nil {
		sn.log.Error("Failed get rule link", "error", err)
		return err
	}

	fields := make([]map[string]interface{}, 0)
	fieldLimitCount := 4
	for index, evt := range evalContext.EvalMatches {
		fields = append(fields, map[string]interface{}{
			"title": evt.Metric,
			"value": evt.Value,
			"short": true,
		})
		if index > fieldLimitCount {
			break
		}
	}

	if evalContext.Error != nil {
		fields = append(fields, map[string]interface{}{
			"title": "Error message",
			"value": evalContext.Error.Error(),
			"short": false,
		})
	}

	mentionsBuilder := strings.Builder{}
	appendSpace := func() {
		if mentionsBuilder.Len() > 0 {
			mentionsBuilder.WriteString(" ")
		}
	}
	mentionChannel := strings.TrimSpace(sn.MentionChannel)
	if mentionChannel != "" {
		mentionsBuilder.WriteString(fmt.Sprintf("<!%s|%s>", mentionChannel, mentionChannel))
	}
	if len(sn.MentionGroups) > 0 {
		appendSpace()
		for _, g := range sn.MentionGroups {
			mentionsBuilder.WriteString(fmt.Sprintf("<!subteam^%s>", g))
		}
	}
	if len(sn.MentionUsers) > 0 {
		appendSpace()
		for _, u := range sn.MentionUsers {
			mentionsBuilder.WriteString(fmt.Sprintf("<@%s>", u))
		}
	}
	msg := ""
	if evalContext.Rule.State != models.AlertStateOK { // don't add message when going back to alert state ok.
		msg = evalContext.Rule.Message
	}
	imageURL := ""
	// default to file.upload API method if a token is provided
	if sn.Token == "" {
		imageURL = evalContext.ImagePublicURL
	}

	var blocks []map[string]interface{}
	if mentionsBuilder.Len() > 0 {
		blocks = []map[string]interface{}{
			{
				"type": "section",
				"text": map[string]interface{}{
					"type": "mrkdwn",
					"text": mentionsBuilder.String(),
				},
			},
		}
	}
	attachment := map[string]interface{}{
		"color":       evalContext.GetStateModel().Color,
		"title":       evalContext.GetNotificationTitle(),
		"title_link":  ruleURL,
		"text":        msg,
		"fallback":    evalContext.GetNotificationTitle(),
		"fields":      fields,
		"footer":      "Grafana v" + setting.BuildVersion,
		"footer_icon": "https://grafana.com/assets/img/fav32.png",
		"ts":          time.Now().Unix(),
	}
	if sn.NeedsImage() && imageURL != "" {
		attachment["image_url"] = imageURL
	}
	body := map[string]interface{}{
		"text": evalContext.GetNotificationTitle(),
		"attachments": []map[string]interface{}{
			attachment,
		},
	}
	if len(blocks) > 0 {
		body["blocks"] = blocks
	}

	// recipient override
	if sn.Recipient != "" {
		body["channel"] = sn.Recipient
	}
	if sn.Username != "" {
		body["username"] = sn.Username
	}
	if sn.IconEmoji != "" {
		body["icon_emoji"] = sn.IconEmoji
	}
	if sn.IconURL != "" {
		body["icon_url"] = sn.IconURL
	}
	data, err := json.Marshal(&body)
	if err != nil {
		return err
	}

	cmd := &models.SendWebhookSync{
		Url:        sn.URL,
		Body:       string(data),
		HttpMethod: http.MethodPost,
	}
	if sn.Token != "" {
		sn.log.Debug("Adding authorization header to HTTP request")
		cmd.HttpHeader = map[string]string{
			"Authorization": fmt.Sprintf("Bearer %s", sn.Token),
		}
	}
	if err := bus.DispatchCtx(evalContext.Ctx, cmd); err != nil {
		sn.log.Error("Failed to send slack notification", "error", err, "webhook", sn.Name)
		return err
	}
	if sn.Token != "" && sn.UploadImage {
		err = sn.slackFileUpload(evalContext, sn.log, "https://slack.com/api/files.upload", sn.Recipient, sn.Token)
		if err != nil {
			return err
		}
	}
	return nil
}

func (sn *SlackNotifier) slackFileUpload(evalContext *alerting.EvalContext, log log.Logger, url string, recipient string, token string) error {
	if evalContext.ImageOnDiskPath == "" {
		// nolint:gosec
		// We can ignore the gosec G304 warning on this one because `setting.HomePath` comes from Grafana's configuration file.
		evalContext.ImageOnDiskPath = filepath.Join(setting.HomePath, "public/img/mixed_styles.png")
	}
	log.Info("Uploading to slack via file.upload API")
	headers, uploadBody, err := sn.generateSlackBody(evalContext.ImageOnDiskPath, token, recipient)
	if err != nil {
		return err
	}
	cmd := &models.SendWebhookSync{Url: url, Body: uploadBody.String(), HttpHeader: headers, HttpMethod: "POST"}
	if err := bus.DispatchCtx(evalContext.Ctx, cmd); err != nil {
		log.Error("Failed to upload slack image", "error", err, "webhook", "file.upload")
		return err
	}
	return nil
}

func (sn *SlackNotifier) generateSlackBody(path string, token string, recipient string) (map[string]string, bytes.Buffer, error) {
	// Slack requires all POSTs to files.upload to present
	// an "application/x-www-form-urlencoded" encoded querystring
	// See https://api.slack.com/methods/files.upload
	var b bytes.Buffer
	w := multipart.NewWriter(&b)
	defer func() {
		if err := w.Close(); err != nil {
			// Shouldn't matter since we already close w explicitly on the non-error path
			sn.log.Warn("Failed to close multipart writer", "err", err)
		}
	}()

	// Add the generated image file
	// We can ignore the gosec G304 warning on this one because `imagePath` comes
	// from the alert `evalContext` that generates the images. `evalContext` in turn derives the root of the file
	// path from configuration variables.
	// nolint:gosec
	f, err := os.Open(path)
	if err != nil {
		return nil, b, err
	}
	defer func() {
		if err := f.Close(); err != nil {
			sn.log.Warn("Failed to close file", "path", path, "err", err)
		}
	}()
	fw, err := w.CreateFormFile("file", path)
	if err != nil {
		return nil, b, err
	}
	if _, err := io.Copy(fw, f); err != nil {
		return nil, b, err
	}
	// Add the authorization token
	if err := w.WriteField("token", token); err != nil {
		return nil, b, err
	}
	// Add the channel(s) to POST to
	if err := w.WriteField("channels", recipient); err != nil {
		return nil, b, err
	}
	if err := w.Close(); err != nil {
		return nil, b, fmt.Errorf("failed to close multipart writer: %w", err)
	}
	headers := map[string]string{
		"Content-Type":  w.FormDataContentType(),
		"Authorization": "auth_token=\"" + token + "\"",
	}
	return headers, b, nil
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2017 Intel Corporation
//
// 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 bat contains a number of helper functions that can be used to perform
// various operations on a ciao cluster such as creating an instance or retrieving
// a list of all the defined workloads, etc.  All of these helper functions are
// implemented by calling ciao-cli rather than by using ciao's REST APIs.  This
// package is mainly intended for use by BAT tests.  Manipulating the cluster
// via ciao-cli, rather than through the REST APIs, allows us to test a little
// bit more of ciao.
package bat

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"gopkg.in/yaml.v2"
)

// Source is provided to the disk structure to indicate whether the
// disk should be cloned from a volume or an image.
type Source struct {
	Type string `yaml:"service"`
	ID   string `yaml:"id"`
}

// Disk describes the storage for the workload definition.
type Disk struct {
	ID        *string `yaml:"volume_id"`
	Size      int     `yaml:"size"`
	Bootable  bool    `yaml:"bootable"`
	Source    *Source `yaml:"source"`
	Ephemeral bool    `yaml:"ephemeral"`
}

// DefaultResources indicate how many cpus and mem to allocate.
type DefaultResources struct {
	VCPUs int `yaml:"vcpus"`
	MemMB int `yaml:"mem_mb"`
}

// WorkloadOptions is used to generate a workload definition in yaml.
type WorkloadOptions struct {
	Description     string           `yaml:"description"`
	VMType          string           `yaml:"vm_type"`
	FWType          string           `yaml:"fw_type"`
	ImageName       string           `yaml:"image_name"`
	Defaults        DefaultResources `yaml:"defaults"`
	CloudConfigFile string           `yaml:"cloud_init"`
	Disks           []Disk           `yaml:"disks"`
}

// this function should output the yaml file for the workload definition
func createWorkloadDefinition(filename string, opt WorkloadOptions) error {
	// take the cloud config file and output it to a file.
	y, err := yaml.Marshal(opt)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(filename, y, 0644)
}

// this function should output a cloud init file for the workload definition
func createCloudConfig(filename string, config string) error {
	// take the cloud config file and output it to a file.
	// convert string to bytes.
	return ioutil.WriteFile(filename, []byte(config), 0644)
}

func createWorkload(ctx context.Context, tenant string, opt WorkloadOptions, config string, public bool) (ID string, err error) {
	cwd, err := os.Getwd()
	if err != nil {
		return "", err
	}

	// create temp dir for the yaml files
	dir, err := ioutil.TempDir("", "ciao_bat")
	if err != nil {
		return "", err
	}

	// you have to call workload create from the same working dir
	// as you have placed your yaml files, so change into that dir.
	err = os.Chdir(dir)
	if err != nil {
		os.RemoveAll(dir)
		return "", err
	}

	// change back to cwd and clean up temp dir on exit.
	defer func() {
		chdirErr := os.Chdir(cwd)
		if err == nil {
			err = chdirErr
		}

		os.RemoveAll(dir)
	}()

	// write out the cloud config.
	err = createCloudConfig("config.yaml", config)
	if err != nil {
		return "", err
	}

	opt.CloudConfigFile = "config.yaml"

	// write out the workload definition
	err = createWorkloadDefinition("workload.yaml", opt)
	if err != nil {
		return "", err
	}

	return CreateWorkloadFromFile(ctx, public, tenant, "workload.yaml")
}

func deleteWorkload(ctx context.Context, tenant string, workload string, public bool) (err error) {
	args := []string{"workload", "delete", "-workload", workload}
	if public {
		_, err = RunCIAOCLIAsAdmin(ctx, tenant, args)
	} else {
		_, err = RunCIAOCLI(ctx, tenant, args)
	}

	return err
}

// CreatePublicWorkload will call ciao-cli as admin to create a workload.
// It will first output the cloud init yaml file to the current working
// directory. Then it will output the workload definition to the current
// working directory. Finally it will call ciao-cli workload create -yaml
// to upload the workload definition. It will clean up all the files it
// created when it is done.
func CreatePublicWorkload(ctx context.Context, tenant string, opt WorkloadOptions, config string) (string, error) {
	return createWorkload(ctx, "", opt, config, true)
}

// CreateWorkloadFromFile will call ciao-cli workload create -yaml
// to create a workload from the specified file.  ciao-cli will be invoked
// using the admin user if the public flag is set to true.  The provided
// workload definition file will remain intact after the call returns.
// The id of the new workload is returned upon success.
func CreateWorkloadFromFile(ctx context.Context, public bool, tenant, wdPath string) (string, error) {
	args := []string{"workload", "create", "-yaml", wdPath}

	var data []byte
	var err error
	if public {
		data, err = RunCIAOCLIAsAdmin(ctx, tenant, args)
	} else {
		data, err = RunCIAOCLI(ctx, tenant, args)
	}

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

	// get the returned workload ID
	s := strings.SplitAfter(string(data), ":")
	if len(s) != 2 {
		return "", fmt.Errorf("Problem parsing workload ID")
	}

	return strings.TrimSpace(s[1]), nil
}

// DeletePublicWorkload will call ciao-cli as admin to delete a workload.
func DeletePublicWorkload(ctx context.Context, workload string) error {
	return deleteWorkload(ctx, "", workload, true)
}

// DeleteWorkload will call ciao-cli as a tenant to delete a workload.
func DeleteWorkload(ctx context.Context, tenant string, workload string) error {
	return deleteWorkload(ctx, tenant, workload, false)
}

// CreateWorkload will call ciao-cli to create a workload definition.
// It will first output the cloud init yaml file to the current working
// directory. Then it will output the workload definition to the current
// working directory. Finally it will call ciao-cli workload create -yaml
// to upload the workload definition. It will clean up all the files it
// created when it is done.
func CreateWorkload(ctx context.Context, tenant string, opt WorkloadOptions, config string) (string, error) {
	return createWorkload(ctx, tenant, opt, config, false)
}

// GetAllWorkloads retrieves a list of all workloads in the cluster by calling
// ciao-cli workload list.  An error will be returned if the following
// environment variables are not set; CIAO_IDENTITY,  CIAO_CONTROLLER,
// CIAO_USERNAME, CIAO_PASSWORD.
func GetAllWorkloads(ctx context.Context, tenant string) ([]Workload, error) {
	var workloads []Workload

	args := []string{"workload", "list", "-f", "{{tojson .}}"}
	err := RunCIAOCLIJS(ctx, tenant, args, &workloads)
	if err != nil {
		return nil, err
	}

	return workloads, nil
}

// GetWorkloadByName will return a specific workload referenced by name.
// An error will be returned if either no workloads exist in the cluster,
// or if the specific workload does not exist. It inherits all error
// conditions of GetAllWorkloads.
func GetWorkloadByName(ctx context.Context, tenant string, name string) (Workload, error) {
	wls, err := GetAllWorkloads(ctx, tenant)
	if err != nil {
		return Workload{}, err
	}

	if len(wls) == 0 {
		return Workload{}, fmt.Errorf("No workloads defined for tenant %s", tenant)
	}

	for _, w := range wls {
		if w.Name == name {
			return w, nil
		}
	}

	return Workload{}, fmt.Errorf("No matching workload for %s", name)
}

// GetWorkloadByID will return a specific workload referenced by name.
// An error will be returned if either no workloads exist in the cluster,
// or if the specific workload does not exist. It inherits all error
// conditions of GetAllWorkloads.
func GetWorkloadByID(ctx context.Context, tenant string, ID string) (Workload, error) {
	wls, err := GetAllWorkloads(ctx, tenant)
	if err != nil {
		return Workload{}, err
	}

	if len(wls) == 0 {
		return Workload{}, fmt.Errorf("No workloads defined for tenant %s", tenant)
	}

	for _, w := range wls {
		if w.ID == ID {
			return w, nil
		}
	}

	return Workload{}, fmt.Errorf("No matching workload for %s", ID)
}
<commit_msg>bat: Use absolute paths in bat.CreateWorkload()<commit_after>//
// Copyright (c) 2017 Intel Corporation
//
// 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 bat contains a number of helper functions that can be used to perform
// various operations on a ciao cluster such as creating an instance or retrieving
// a list of all the defined workloads, etc.  All of these helper functions are
// implemented by calling ciao-cli rather than by using ciao's REST APIs.  This
// package is mainly intended for use by BAT tests.  Manipulating the cluster
// via ciao-cli, rather than through the REST APIs, allows us to test a little
// bit more of ciao.
package bat

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

	"gopkg.in/yaml.v2"
)

// Source is provided to the disk structure to indicate whether the
// disk should be cloned from a volume or an image.
type Source struct {
	Type string `yaml:"service"`
	ID   string `yaml:"id"`
}

// Disk describes the storage for the workload definition.
type Disk struct {
	ID        *string `yaml:"volume_id"`
	Size      int     `yaml:"size"`
	Bootable  bool    `yaml:"bootable"`
	Source    *Source `yaml:"source"`
	Ephemeral bool    `yaml:"ephemeral"`
}

// DefaultResources indicate how many cpus and mem to allocate.
type DefaultResources struct {
	VCPUs int `yaml:"vcpus"`
	MemMB int `yaml:"mem_mb"`
}

// WorkloadOptions is used to generate a workload definition in yaml.
type WorkloadOptions struct {
	Description     string           `yaml:"description"`
	VMType          string           `yaml:"vm_type"`
	FWType          string           `yaml:"fw_type"`
	ImageName       string           `yaml:"image_name"`
	Defaults        DefaultResources `yaml:"defaults"`
	CloudConfigFile string           `yaml:"cloud_init"`
	Disks           []Disk           `yaml:"disks"`
}

// this function should output the yaml file for the workload definition
func createWorkloadDefinition(filename string, opt WorkloadOptions) error {
	// take the cloud config file and output it to a file.
	y, err := yaml.Marshal(opt)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(filename, y, 0644)
}

// this function should output a cloud init file for the workload definition
func createCloudConfig(filename string, config string) error {
	// take the cloud config file and output it to a file.
	// convert string to bytes.
	return ioutil.WriteFile(filename, []byte(config), 0644)
}

func createWorkload(ctx context.Context, tenant string, opt WorkloadOptions, config string, public bool) (ID string, err error) {
	// create temp dir for the yaml files
	dir, err := ioutil.TempDir("", "ciao_bat")
	if err != nil {
		return "", err
	}
	defer func() {
		os.RemoveAll(dir)
	}()

	opt.CloudConfigFile = path.Join(dir, "config.yaml")

	// write out the cloud config.
	err = createCloudConfig(opt.CloudConfigFile, config)
	if err != nil {
		return "", err
	}

	// write out the workload definition
	fp := path.Join(dir, "workload.yaml")
	err = createWorkloadDefinition(fp, opt)
	if err != nil {
		return "", err
	}

	return CreateWorkloadFromFile(ctx, public, tenant, fp)
}

func deleteWorkload(ctx context.Context, tenant string, workload string, public bool) (err error) {
	args := []string{"workload", "delete", "-workload", workload}
	if public {
		_, err = RunCIAOCLIAsAdmin(ctx, tenant, args)
	} else {
		_, err = RunCIAOCLI(ctx, tenant, args)
	}

	return err
}

// CreatePublicWorkload will call ciao-cli as admin to create a workload.
// It will first output the cloud init yaml file to the current working
// directory. Then it will output the workload definition to the current
// working directory. Finally it will call ciao-cli workload create -yaml
// to upload the workload definition. It will clean up all the files it
// created when it is done.
func CreatePublicWorkload(ctx context.Context, tenant string, opt WorkloadOptions, config string) (string, error) {
	return createWorkload(ctx, "", opt, config, true)
}

// CreateWorkloadFromFile will call ciao-cli workload create -yaml
// to create a workload from the specified file.  ciao-cli will be invoked
// using the admin user if the public flag is set to true.  The provided
// workload definition file will remain intact after the call returns.
// The id of the new workload is returned upon success.
func CreateWorkloadFromFile(ctx context.Context, public bool, tenant, wdPath string) (string, error) {
	args := []string{"workload", "create", "-yaml", wdPath}

	var data []byte
	var err error
	if public {
		data, err = RunCIAOCLIAsAdmin(ctx, tenant, args)
	} else {
		data, err = RunCIAOCLI(ctx, tenant, args)
	}

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

	// get the returned workload ID
	s := strings.SplitAfter(string(data), ":")
	if len(s) != 2 {
		return "", fmt.Errorf("Problem parsing workload ID")
	}

	return strings.TrimSpace(s[1]), nil
}

// DeletePublicWorkload will call ciao-cli as admin to delete a workload.
func DeletePublicWorkload(ctx context.Context, workload string) error {
	return deleteWorkload(ctx, "", workload, true)
}

// DeleteWorkload will call ciao-cli as a tenant to delete a workload.
func DeleteWorkload(ctx context.Context, tenant string, workload string) error {
	return deleteWorkload(ctx, tenant, workload, false)
}

// CreateWorkload will call ciao-cli to create a workload definition.
// It will first output the cloud init yaml file to the current working
// directory. Then it will output the workload definition to the current
// working directory. Finally it will call ciao-cli workload create -yaml
// to upload the workload definition. It will clean up all the files it
// created when it is done.
func CreateWorkload(ctx context.Context, tenant string, opt WorkloadOptions, config string) (string, error) {
	return createWorkload(ctx, tenant, opt, config, false)
}

// GetAllWorkloads retrieves a list of all workloads in the cluster by calling
// ciao-cli workload list.  An error will be returned if the following
// environment variables are not set; CIAO_IDENTITY,  CIAO_CONTROLLER,
// CIAO_USERNAME, CIAO_PASSWORD.
func GetAllWorkloads(ctx context.Context, tenant string) ([]Workload, error) {
	var workloads []Workload

	args := []string{"workload", "list", "-f", "{{tojson .}}"}
	err := RunCIAOCLIJS(ctx, tenant, args, &workloads)
	if err != nil {
		return nil, err
	}

	return workloads, nil
}

// GetWorkloadByName will return a specific workload referenced by name.
// An error will be returned if either no workloads exist in the cluster,
// or if the specific workload does not exist. It inherits all error
// conditions of GetAllWorkloads.
func GetWorkloadByName(ctx context.Context, tenant string, name string) (Workload, error) {
	wls, err := GetAllWorkloads(ctx, tenant)
	if err != nil {
		return Workload{}, err
	}

	if len(wls) == 0 {
		return Workload{}, fmt.Errorf("No workloads defined for tenant %s", tenant)
	}

	for _, w := range wls {
		if w.Name == name {
			return w, nil
		}
	}

	return Workload{}, fmt.Errorf("No matching workload for %s", name)
}

// GetWorkloadByID will return a specific workload referenced by name.
// An error will be returned if either no workloads exist in the cluster,
// or if the specific workload does not exist. It inherits all error
// conditions of GetAllWorkloads.
func GetWorkloadByID(ctx context.Context, tenant string, ID string) (Workload, error) {
	wls, err := GetAllWorkloads(ctx, tenant)
	if err != nil {
		return Workload{}, err
	}

	if len(wls) == 0 {
		return Workload{}, fmt.Errorf("No workloads defined for tenant %s", tenant)
	}

	for _, w := range wls {
		if w.ID == ID {
			return w, nil
		}
	}

	return Workload{}, fmt.Errorf("No matching workload for %s", ID)
}
<|endoftext|>


<commit_before>package qshell

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/astaxie/beego/logs"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"qiniu/api.v6/auth/digest"
	rio "qiniu/api.v6/resumable/io"
	"qiniu/api.v6/rs"
	"qiniu/rpc"
	"strconv"
	"strings"
	"time"
)

//range get and chunk upload

const (
	RETRY_MAX_TIMES = 5
	RETRY_INTERVAL  = time.Second * 1
	HTTP_TIMEOUT    = time.Second * 10
)

type PutRet struct {
	Key      string `json:"key"`
	Hash     string `json:"hash"`
	MimeType string `json:"mimeType"`
	Fsize    int64  `json:"fsize"`
}

type SyncProgress struct {
	BlkCtxs   []rio.BlkputRet `json:"blk_ctxs"`
	Offset    int64           `json:"offset"`
	TotalSize int64           `json:"total_size"`
}

func Sync(mac *digest.Mac, srcResUrl, bucket, key, upHostIp string) (putRet PutRet, err error) {
	if exists, cErr := checkExists(mac, bucket, key); cErr != nil {
		err = cErr
		return
	} else if exists {
		err = errors.New("File with same key already exists in bucket")
		return
	}

	syncProgress := SyncProgress{}
	//create sync id
	syncId := Md5Hex(fmt.Sprintf("%s:%s:%s", srcResUrl, bucket, key))

	//local storage path
	storePath := filepath.Join(QShellRootPath, ".qshell", "sync")
	if mkdirErr := os.MkdirAll(storePath, 0775); mkdirErr != nil {
		logs.Error("Failed to mkdir `%s` due to `%s`", storePath, mkdirErr)
		return
	}

	progressFile := filepath.Join(storePath, fmt.Sprintf("%s.progress", syncId))
	if statInfo, statErr := os.Stat(progressFile); statErr == nil {
		//check file last modified time, if older than one week, ignore
		if statInfo.ModTime().Add(time.Hour * 24 * 5).After(time.Now()) {
			//try read old progress
			progressFh, openErr := os.Open(progressFile)
			if openErr == nil {
				decoder := json.NewDecoder(progressFh)
				decoder.Decode(&syncProgress)
				progressFh.Close()
			}
		}
	}

	//check offset valid or not
	if syncProgress.Offset%BLOCK_SIZE != 0 {
		logs.Informational("Invalid offset from progress file,", syncProgress.Offset)
		syncProgress.Offset = 0
		syncProgress.TotalSize = 0
		syncProgress.BlkCtxs = make([]rio.BlkputRet, 0)
	}

	//check offset and blk ctxs
	if syncProgress.Offset != 0 && syncProgress.BlkCtxs != nil {
		if int(syncProgress.Offset/BLOCK_SIZE) != len(syncProgress.BlkCtxs) {
			logs.Informational("Invalid offset and block contexts")
			syncProgress.Offset = 0
			syncProgress.TotalSize = 0
			syncProgress.BlkCtxs = make([]rio.BlkputRet, 0)
		}
	}

	//check blk ctxs, when no progress found
	if syncProgress.Offset == 0 || syncProgress.BlkCtxs == nil {
		syncProgress.Offset = 0
		syncProgress.TotalSize = 0
		syncProgress.BlkCtxs = make([]rio.BlkputRet, 0)
	}

	//get total size
	totalSize, hErr := getRemoteFileLength(srcResUrl)
	if hErr != nil {
		err = hErr
		return
	}

	if totalSize != syncProgress.TotalSize {
		if syncProgress.TotalSize != 0 {
			logs.Warning("Remote file length changed, progress file out of date")
		}
		syncProgress.Offset = 0
		syncProgress.TotalSize = totalSize
		syncProgress.BlkCtxs = make([]rio.BlkputRet, 0)
	}

	//get total block count
	totalBlkCnt := 0
	if totalSize%BLOCK_SIZE == 0 {
		totalBlkCnt = int(totalSize / BLOCK_SIZE)
	} else {
		totalBlkCnt = int(totalSize/BLOCK_SIZE) + 1
	}

	//init the range offset
	rangeStartOffset := syncProgress.Offset
	fromBlkIndex := int(rangeStartOffset / BLOCK_SIZE)

	lastBlock := false

	//create upload token
	policy := rs.PutPolicy{Scope: bucket}
	//token is valid for one year
	policy.Expires = 3600 * 24 * 365
	policy.ReturnBody = `{"key":"$(key)","hash":"$(etag)","fsize":$(fsize),"mimeType":"$(mimeType)"}`
	uptoken := policy.Token(mac)
	putClient := rio.NewClient(uptoken, upHostIp)

	//range get and mkblk upload
	for blkIndex := fromBlkIndex; blkIndex < totalBlkCnt; blkIndex++ {
		if blkIndex == totalBlkCnt-1 {
			lastBlock = true
		}

		syncPercent := fmt.Sprintf("%.2f", float64(blkIndex+1)*100.0/float64(totalBlkCnt))
		logs.Informational("Syncing block %d [%s%%] ...", blkIndex, syncPercent)
		blkCtx, pErr := rangeMkblkPipe(srcResUrl, rangeStartOffset, BLOCK_SIZE, lastBlock, putClient)
		if pErr != nil {
			logs.Error(pErr.Error())
			time.Sleep(RETRY_INTERVAL)

			for retryTimes := 1; retryTimes <= RETRY_MAX_TIMES; retryTimes++ {
				logs.Informational(fmt.Sprintf("Retrying %d time range & mkblk block [%d]", retryTimes, blkIndex))
				blkCtx, pErr = rangeMkblkPipe(srcResUrl, rangeStartOffset, BLOCK_SIZE, lastBlock, putClient)
				if pErr != nil {
					logs.Error(pErr)
					//wait a interval and retry
					time.Sleep(RETRY_INTERVAL)
					continue
				} else {
					break
				}
			}
		}

		if pErr != nil {
			err = errors.New("Max retry reached and range & mkblk still failed, check your network")
			return
		}

		//advance range offset
		rangeStartOffset += BLOCK_SIZE

		syncProgress.BlkCtxs = append(syncProgress.BlkCtxs, blkCtx)
		syncProgress.Offset = rangeStartOffset

		rErr := recordProgress(progressFile, syncProgress)
		if rErr != nil {
			logs.Informational(rErr.Error())
		}
	}

	//make file
	putExtra := rio.PutExtra{
		Progresses: syncProgress.BlkCtxs,
	}
	mkErr := rio.Mkfile(putClient, nil, &putRet, key, true, totalSize, &putExtra)
	if mkErr != nil {
		err = fmt.Errorf("Mkfile error, %s", mkErr.Error())
		return
	}

	//delete progress file
	os.Remove(progressFile)

	return
}

func rangeMkblkPipe(srcResUrl string, rangeStartOffset int64, rangeBlockSize int64, lastBlock bool,
	putClient rpc.Client) (putRet rio.BlkputRet, err error) {
	//range get
	dReq, dReqErr := http.NewRequest("GET", srcResUrl, nil)
	if dReqErr != nil {
		err = fmt.Errorf("New request error, %s", dReqErr.Error())
		return
	}

	//set range header
	rangeEndOffset := rangeStartOffset + rangeBlockSize - 1
	dReq.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", rangeStartOffset, rangeEndOffset))

	//get resp
	client := http.DefaultClient
	client.Timeout = time.Duration(HTTP_TIMEOUT)
	dResp, dRespErr := client.Do(dReq)
	if dRespErr != nil {
		err = fmt.Errorf("Get response error, %s", dRespErr.Error())
		return
	}
	defer dResp.Body.Close()

	//status error
	if dResp.StatusCode/100 != 2 {
		err = fmt.Errorf("Get resource error, %s", dResp.Status)
		return
	}

	//if not support range, go back and err
	if dResp.Header.Get("Accept-Ranges") == "" {
		err = errors.New("Remote server not support range")
		return
	}

	//parse content-range
	contentRange := dResp.Header.Get("Content-Range")
	rangeSize, _ := parseContentRange(contentRange)

	//check ranged block size
	if !lastBlock && rangeSize != rangeBlockSize {
		err = errors.New("Block read error, only the last range block can has bytes less than <RangeBlockSize>")
		return
	}

	//read content
	buffer := bytes.NewBuffer(nil)
	cpCnt, cpErr := io.Copy(buffer, dResp.Body)
	if cpErr != nil || cpCnt != rangeSize {
		err = errors.New("Read range block response error, not fully read")
		return
	}

	//mkblk
	blkPutRet := rio.BlkputRet{}
	blockSize := int(rangeSize)
	blockDataReader := bytes.NewReader(buffer.Bytes())
	blockDataSize := buffer.Len()

	mkErr := rio.Mkblock(putClient, nil, &blkPutRet, blockSize, blockDataReader, blockDataSize)
	if mkErr != nil {
		err = fmt.Errorf("Mkblk error, %s", mkErr.Error())
		return
	}

	putRet = blkPutRet

	return
}

func recordProgress(progressFile string, syncProgress SyncProgress) (err error) {
	fh, openErr := os.Create(progressFile)
	if openErr != nil {
		err = fmt.Errorf("Open progress file %s error, %s", progressFile, openErr.Error())
		return
	}
	defer fh.Close()

	jsonBytes, mErr := json.Marshal(&syncProgress)
	if mErr != nil {
		err = fmt.Errorf("Marshal sync progress error, %s", mErr.Error())
		return
	}

	_, wErr := fh.Write(jsonBytes)
	if wErr != nil {
		err = fmt.Errorf("Write sync progress error, %s", wErr.Error())
	}

	return
}

func getRemoteFileLength(srcResUrl string) (totalSize int64, err error) {
	resp, respErr := http.Head(srcResUrl)
	if respErr != nil {
		err = fmt.Errorf("New head request failed, %s", respErr.Error())
		return
	}
	defer resp.Body.Close()

	contentLength := resp.Header.Get("Content-Length")
	if contentLength == "" {
		err = errors.New("Head request with no Content-Length found error")
		return
	}

	totalSize, _ = strconv.ParseInt(contentLength, 10, 64)

	return
}

func checkExists(mac *digest.Mac, bucket, key string) (exists bool, err error) {
	client := rs.NewMac(mac)
	entry, sErr := client.Stat(nil, bucket, key)
	if sErr != nil {
		if v, ok := sErr.(*rpc.ErrorInfo); !ok {
			err = fmt.Errorf("Check file exists error, %s", sErr.Error())
			return
		} else {
			if v.Code != 612 {
				err = fmt.Errorf("Check file exists error, %s", v.Err)
				return
			} else {
				exists = false
				return
			}
		}
	}

	if entry.Hash != "" {
		exists = true
	}

	return
}

//Content-Range: bytes 25538640-25538647/25538648
func parseContentRange(contentRange string) (rangeSize, totalSize int64) {
	contentRangeItems := strings.Split(contentRange, " ")
	sizeItems := strings.Split(contentRangeItems[1], "/")

	rangePartItems := strings.Split(sizeItems[0], "-")
	totalSize, _ = strconv.ParseInt(sizeItems[1], 10, 64)

	fromOffset, _ := strconv.ParseInt(rangePartItems[0], 10, 64)
	toOffset, _ := strconv.ParseInt(rangePartItems[1], 10, 64)

	rangeSize = toOffset - fromOffset + 1

	return
}
<commit_msg>fmt output<commit_after>package qshell

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/astaxie/beego/logs"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"qiniu/api.v6/auth/digest"
	rio "qiniu/api.v6/resumable/io"
	"qiniu/api.v6/rs"
	"qiniu/rpc"
	"strconv"
	"strings"
	"time"
)

//range get and chunk upload

const (
	RETRY_MAX_TIMES = 5
	RETRY_INTERVAL  = time.Second * 1
	HTTP_TIMEOUT    = time.Second * 10
)

type PutRet struct {
	Key      string `json:"key"`
	Hash     string `json:"hash"`
	MimeType string `json:"mimeType"`
	Fsize    int64  `json:"fsize"`
}

type SyncProgress struct {
	BlkCtxs   []rio.BlkputRet `json:"blk_ctxs"`
	Offset    int64           `json:"offset"`
	TotalSize int64           `json:"total_size"`
}

func Sync(mac *digest.Mac, srcResUrl, bucket, key, upHostIp string) (putRet PutRet, err error) {
	if exists, cErr := checkExists(mac, bucket, key); cErr != nil {
		err = cErr
		return
	} else if exists {
		err = errors.New("File with same key already exists in bucket")
		return
	}

	syncProgress := SyncProgress{}
	//create sync id
	syncId := Md5Hex(fmt.Sprintf("%s:%s:%s", srcResUrl, bucket, key))

	//local storage path
	storePath := filepath.Join(QShellRootPath, ".qshell", "sync")
	if mkdirErr := os.MkdirAll(storePath, 0775); mkdirErr != nil {
		logs.Error("Failed to mkdir ", storePath, " due to ", mkdirErr)
		return
	}

	progressFile := filepath.Join(storePath, fmt.Sprintf("%s.progress", syncId))
	if statInfo, statErr := os.Stat(progressFile); statErr == nil {
		//check file last modified time, if older than one week, ignore
		if statInfo.ModTime().Add(time.Hour * 24 * 5).After(time.Now()) {
			//try read old progress
			progressFh, openErr := os.Open(progressFile)
			if openErr == nil {
				decoder := json.NewDecoder(progressFh)
				decoder.Decode(&syncProgress)
				progressFh.Close()
			}
		}
	}

	//check offset valid or not
	if syncProgress.Offset%BLOCK_SIZE != 0 {
		logs.Informational("Invalid offset from progress file,", syncProgress.Offset)
		syncProgress.Offset = 0
		syncProgress.TotalSize = 0
		syncProgress.BlkCtxs = make([]rio.BlkputRet, 0)
	}

	//check offset and blk ctxs
	if syncProgress.Offset != 0 && syncProgress.BlkCtxs != nil {
		if int(syncProgress.Offset/BLOCK_SIZE) != len(syncProgress.BlkCtxs) {
			logs.Informational("Invalid offset and block contexts")
			syncProgress.Offset = 0
			syncProgress.TotalSize = 0
			syncProgress.BlkCtxs = make([]rio.BlkputRet, 0)
		}
	}

	//check blk ctxs, when no progress found
	if syncProgress.Offset == 0 || syncProgress.BlkCtxs == nil {
		syncProgress.Offset = 0
		syncProgress.TotalSize = 0
		syncProgress.BlkCtxs = make([]rio.BlkputRet, 0)
	}

	//get total size
	totalSize, hErr := getRemoteFileLength(srcResUrl)
	if hErr != nil {
		err = hErr
		return
	}

	if totalSize != syncProgress.TotalSize {
		if syncProgress.TotalSize != 0 {
			logs.Warning("Remote file length changed, progress file out of date")
		}
		syncProgress.Offset = 0
		syncProgress.TotalSize = totalSize
		syncProgress.BlkCtxs = make([]rio.BlkputRet, 0)
	}

	//get total block count
	totalBlkCnt := 0
	if totalSize%BLOCK_SIZE == 0 {
		totalBlkCnt = int(totalSize / BLOCK_SIZE)
	} else {
		totalBlkCnt = int(totalSize/BLOCK_SIZE) + 1
	}

	//init the range offset
	rangeStartOffset := syncProgress.Offset
	fromBlkIndex := int(rangeStartOffset / BLOCK_SIZE)

	lastBlock := false

	//create upload token
	policy := rs.PutPolicy{Scope: bucket}
	//token is valid for one year
	policy.Expires = 3600 * 24 * 365
	policy.ReturnBody = `{"key":"$(key)","hash":"$(etag)","fsize":$(fsize),"mimeType":"$(mimeType)"}`
	uptoken := policy.Token(mac)
	putClient := rio.NewClient(uptoken, upHostIp)

	//range get and mkblk upload
	for blkIndex := fromBlkIndex; blkIndex < totalBlkCnt; blkIndex++ {
		if blkIndex == totalBlkCnt-1 {
			lastBlock = true
		}

		syncPercent := fmt.Sprintf("%.2f", float64(blkIndex+1)*100.0/float64(totalBlkCnt))
		logs.Info(fmt.Sprintf("Syncing block %d [%s%%] ...", blkIndex, syncPercent))
		blkCtx, pErr := rangeMkblkPipe(srcResUrl, rangeStartOffset, BLOCK_SIZE, lastBlock, putClient)
		if pErr != nil {
			logs.Error(pErr.Error())
			time.Sleep(RETRY_INTERVAL)

			for retryTimes := 1; retryTimes <= RETRY_MAX_TIMES; retryTimes++ {
				logs.Informational(fmt.Sprintf("Retrying %d time range & mkblk block [%d]", retryTimes, blkIndex))
				blkCtx, pErr = rangeMkblkPipe(srcResUrl, rangeStartOffset, BLOCK_SIZE, lastBlock, putClient)
				if pErr != nil {
					logs.Error(pErr)
					//wait a interval and retry
					time.Sleep(RETRY_INTERVAL)
					continue
				} else {
					break
				}
			}
		}

		if pErr != nil {
			err = errors.New("Max retry reached and range & mkblk still failed, check your network")
			return
		}

		//advance range offset
		rangeStartOffset += BLOCK_SIZE

		syncProgress.BlkCtxs = append(syncProgress.BlkCtxs, blkCtx)
		syncProgress.Offset = rangeStartOffset

		rErr := recordProgress(progressFile, syncProgress)
		if rErr != nil {
			logs.Informational(rErr.Error())
		}
	}

	//make file
	putExtra := rio.PutExtra{
		Progresses: syncProgress.BlkCtxs,
	}
	mkErr := rio.Mkfile(putClient, nil, &putRet, key, true, totalSize, &putExtra)
	if mkErr != nil {
		err = fmt.Errorf("Mkfile error, %s", mkErr.Error())
		return
	}

	//delete progress file
	os.Remove(progressFile)

	return
}

func rangeMkblkPipe(srcResUrl string, rangeStartOffset int64, rangeBlockSize int64, lastBlock bool,
	putClient rpc.Client) (putRet rio.BlkputRet, err error) {
	//range get
	dReq, dReqErr := http.NewRequest("GET", srcResUrl, nil)
	if dReqErr != nil {
		err = fmt.Errorf("New request error, %s", dReqErr.Error())
		return
	}

	//set range header
	rangeEndOffset := rangeStartOffset + rangeBlockSize - 1
	dReq.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", rangeStartOffset, rangeEndOffset))

	//get resp
	client := http.DefaultClient
	client.Timeout = time.Duration(HTTP_TIMEOUT)
	dResp, dRespErr := client.Do(dReq)
	if dRespErr != nil {
		err = fmt.Errorf("Get response error, %s", dRespErr.Error())
		return
	}
	defer dResp.Body.Close()

	//status error
	if dResp.StatusCode/100 != 2 {
		err = fmt.Errorf("Get resource error, %s", dResp.Status)
		return
	}

	//if not support range, go back and err
	if dResp.Header.Get("Accept-Ranges") == "" {
		err = errors.New("Remote server not support range")
		return
	}

	//parse content-range
	contentRange := dResp.Header.Get("Content-Range")
	rangeSize, _ := parseContentRange(contentRange)

	//check ranged block size
	if !lastBlock && rangeSize != rangeBlockSize {
		err = errors.New("Block read error, only the last range block can has bytes less than <RangeBlockSize>")
		return
	}

	//read content
	buffer := bytes.NewBuffer(nil)
	cpCnt, cpErr := io.Copy(buffer, dResp.Body)
	if cpErr != nil || cpCnt != rangeSize {
		err = errors.New("Read range block response error, not fully read")
		return
	}

	//mkblk
	blkPutRet := rio.BlkputRet{}
	blockSize := int(rangeSize)
	blockDataReader := bytes.NewReader(buffer.Bytes())
	blockDataSize := buffer.Len()

	mkErr := rio.Mkblock(putClient, nil, &blkPutRet, blockSize, blockDataReader, blockDataSize)
	if mkErr != nil {
		err = fmt.Errorf("Mkblk error, %s", mkErr.Error())
		return
	}

	putRet = blkPutRet

	return
}

func recordProgress(progressFile string, syncProgress SyncProgress) (err error) {
	fh, openErr := os.Create(progressFile)
	if openErr != nil {
		err = fmt.Errorf("Open progress file %s error, %s", progressFile, openErr.Error())
		return
	}
	defer fh.Close()

	jsonBytes, mErr := json.Marshal(&syncProgress)
	if mErr != nil {
		err = fmt.Errorf("Marshal sync progress error, %s", mErr.Error())
		return
	}

	_, wErr := fh.Write(jsonBytes)
	if wErr != nil {
		err = fmt.Errorf("Write sync progress error, %s", wErr.Error())
	}

	return
}

func getRemoteFileLength(srcResUrl string) (totalSize int64, err error) {
	resp, respErr := http.Head(srcResUrl)
	if respErr != nil {
		err = fmt.Errorf("New head request failed, %s", respErr.Error())
		return
	}
	defer resp.Body.Close()

	contentLength := resp.Header.Get("Content-Length")
	if contentLength == "" {
		err = errors.New("Head request with no Content-Length found error")
		return
	}

	totalSize, _ = strconv.ParseInt(contentLength, 10, 64)

	return
}

func checkExists(mac *digest.Mac, bucket, key string) (exists bool, err error) {
	client := rs.NewMac(mac)
	entry, sErr := client.Stat(nil, bucket, key)
	if sErr != nil {
		if v, ok := sErr.(*rpc.ErrorInfo); !ok {
			err = fmt.Errorf("Check file exists error, %s", sErr.Error())
			return
		} else {
			if v.Code != 612 {
				err = fmt.Errorf("Check file exists error, %s", v.Err)
				return
			} else {
				exists = false
				return
			}
		}
	}

	if entry.Hash != "" {
		exists = true
	}

	return
}

//Content-Range: bytes 25538640-25538647/25538648
func parseContentRange(contentRange string) (rangeSize, totalSize int64) {
	contentRangeItems := strings.Split(contentRange, " ")
	sizeItems := strings.Split(contentRangeItems[1], "/")

	rangePartItems := strings.Split(sizeItems[0], "-")
	totalSize, _ = strconv.ParseInt(sizeItems[1], 10, 64)

	fromOffset, _ := strconv.ParseInt(rangePartItems[0], 10, 64)
	toOffset, _ := strconv.ParseInt(rangePartItems[1], 10, 64)

	rangeSize = toOffset - fromOffset + 1

	return
}
<|endoftext|>


<commit_before>package nodejs

import (
	"errors"
	"fmt"
	"net"
	"os/exec"
	"regexp"
	"runtime"
	"strings"

	"github.com/PuerkitoBio/goquery"

	"github.com/markelog/eclectica/variables"
)

var (
	VersionsLink = "https://nodejs.org/dist"

	bins = []string{"node", "npm"}
)

type Node struct {
	Version string
}

func (node Node) Install() error {
	return nil
}

func (node Node) PostInstall() error {
	return nil
}

func (node Node) Environment() (string, error) {
	return "", nil
}

func (node Node) Info() (map[string]string, error) {
	result := make(map[string]string)
	sourcesUrl := fmt.Sprintf("%s/v%s", VersionsLink, node.Version)

	result["filename"] = fmt.Sprintf("node-v%s-%s-x64", node.Version, runtime.GOOS)
	result["url"] = fmt.Sprintf("%s/%s.tar.gz", sourcesUrl, result["filename"])

	return result, nil
}

func (node Node) Bins() []string {
	return bins
}

func (node Node) Current() string {
	bin := variables.GetBin("node")
	out, _ := exec.Command(bin, "--version").Output()

	version := strings.TrimSpace(string(out))

	return strings.Replace(version, "v", "", 1)
}

func (node Node) ListRemote() ([]string, error) {
	doc, err := goquery.NewDocument(VersionsLink)

	if err != nil {
		if _, ok := err.(net.Error); ok {
			return nil, errors.New("Can't establish connection")
		}

		return nil, err
	}

	tmp := []string{}
	result := []string{}

	version := regexp.MustCompile("v\\d+\\.\\d+\\.\\d+$")
	remove := regexp.MustCompile("0\\.[0-7]")

	links := doc.Find("a")

	for i := range links.Nodes {
		href, _ := links.Eq(i).Attr("href")

		href = strings.Replace(href, "/", "", 1)
		if version.MatchString(href) {
			href = strings.Replace(href, "v", "", 1)
			tmp = append(tmp, href)
		}
	}

	// Remove < 0.8 versions
	for _, element := range tmp {
		if remove.MatchString(element) == false {
			result = append(result, element)
		}
	}

	return result, nil
}
<commit_msg>Move around regexp patterns<commit_after>package nodejs

import (
	"errors"
	"fmt"
	"net"
	"os/exec"
	"regexp"
	"runtime"
	"strings"

	"github.com/PuerkitoBio/goquery"

	"github.com/markelog/eclectica/variables"
)

var (
	VersionsLink   = "https://nodejs.org/dist"
	versionPattern = "v\\d+\\.\\d+\\.\\d+$"
	removePattern  = "0\\.[0-7]"

	bins = []string{"node", "npm"}
)

type Node struct {
	Version string
}

func (node Node) Install() error {
	return nil
}

func (node Node) PostInstall() error {
	return nil
}

func (node Node) Environment() (string, error) {
	return "", nil
}

func (node Node) Info() (map[string]string, error) {
	result := make(map[string]string)
	sourcesUrl := fmt.Sprintf("%s/v%s", VersionsLink, node.Version)

	result["filename"] = fmt.Sprintf("node-v%s-%s-x64", node.Version, runtime.GOOS)
	result["url"] = fmt.Sprintf("%s/%s.tar.gz", sourcesUrl, result["filename"])

	return result, nil
}

func (node Node) Bins() []string {
	return bins
}

func (node Node) Current() string {
	bin := variables.GetBin("node")
	out, _ := exec.Command(bin, "--version").Output()

	version := strings.TrimSpace(string(out))

	return strings.Replace(version, "v", "", 1)
}

func (node Node) ListRemote() ([]string, error) {
	doc, err := goquery.NewDocument(VersionsLink)

	if err != nil {
		if _, ok := err.(net.Error); ok {
			return nil, errors.New("Can't establish connection")
		}

		return nil, err
	}

	tmp := []string{}
	result := []string{}

	version := regexp.MustCompile(versionPattern)
	remove := regexp.MustCompile(removePattern)

	links := doc.Find("a")

	for i := range links.Nodes {
		href, _ := links.Eq(i).Attr("href")

		href = strings.Replace(href, "/", "", 1)
		if version.MatchString(href) {
			href = strings.Replace(href, "v", "", 1)
			tmp = append(tmp, href)
		}
	}

	// Remove < 0.8 versions
	for _, element := range tmp {
		if remove.MatchString(element) == false {
			result = append(result, element)
		}
	}

	return result, nil
}
<|endoftext|>


<commit_before>package octorrent

import (
	"bufio"
	"bytes"
	"math"
	"reflect"
	"strconv"
	"testing"
)

func getDecoder(d []byte) decoder {
	buf := bytes.NewBuffer(d)

	return decoder{
		*bufio.NewReader(buf),
	}
}

func TestReadInt(t *testing.T) {
	i, o := "412e51", int64(412)
	dec := getDecoder([]byte(i))

	num, err := dec.readInt()
	if err != nil {
		t.Errorf("readInt returned error: %v", err)
	}

	if num != o {
		t.Errorf("readInt did return %i, but should return %i", num, o)
	}
}

func TestReadIntUntil(t *testing.T) {
	i, o := "142", int64(14)
	dec := getDecoder([]byte(i))

	num, err := dec.readIntUntil('2')
	if err != nil {
		t.Errorf("readIntUntil returned error: %v", err)
	}

	if num != o {
		t.Errorf("readIntUntil did return %i, but should return %i", num, o)
	}

	// Check broken reader
	dec = getDecoder([]byte(""))
	num, err = dec.readIntUntil('t')
	if err == nil {
		t.Errorf("readIntUntil with empty reader returned no error, but should")
	}

	dec = getDecoder([]byte("test"))
	num, err = dec.readIntUntil('e')
	if err == nil {
		t.Errorf("readIntUntil with no digits returned no error, but should")
	}

	// Check unsigned int
	var ui uint64
	ui = math.MaxInt64 + 1
	i = "i" + strconv.FormatUint(ui, 10) + "e"
	dec = getDecoder([]byte(i))
	num, err = dec.readIntUntil('e')
	if err == nil {
		t.Errorf("readIntUntil with unsigned int returned no error, but should")
	}
}

func TestReadString(t *testing.T) {
	i, o := "4:testfoobar", "test"
	dec := getDecoder([]byte(i))
	if str, err := dec.readString(); err != nil {
		t.Errorf("readString returned error: %v", err)
	} else if str != o {
		t.Errorf("readString returned %s, but should return %s", str, o)
	}

	// Test broken reader
	dec = getDecoder([]byte(""))
	if _, err := dec.readString(); err == nil {
		t.Errorf("readSstring with empty buffer returned no error, but should")
	}

	// Test wrong format
	dec = getDecoder([]byte("4:foo"))
	if _, err := dec.readString(); err == nil {
		t.Errorf("readString with wrong formatted string returned no error, but should")
	}
}

func TestReadValue(t *testing.T) {
	dec := getDecoder([]byte(""))
	if _, err := dec.readValue(); err == nil {
		t.Errorf("readValue with empty buffer returned no error, but should")
	}

	dec = getDecoder([]byte("ifooe"))
	if _, err := dec.readValue(); err == nil {
		t.Errorf("readValue with wrongly formatted buffer returned no error, but should")
	}
}

func TestIsEnd(t *testing.T) {
	i, o := "4:test", false
	dec := getDecoder([]byte(i))
	end, err := dec.isEnd()
	if err != nil {
		t.Errorf("isEnd returned error: %v", err)
	}
	if end != o {
		t.Errorf("isEnd returned %v, but should return %v", end, o)
	}

	i, o = "e4:test", true
	dec = getDecoder([]byte(i))
	end, err = dec.isEnd()
	if err != nil {
		t.Errorf("isEnd returned error: %v", err)
	}
	if end != o {
		t.Errorf("isEnd returned %v, but should return %v", end, o)
	}

	dec = getDecoder([]byte(""))
	if _, err := dec.isEnd(); err == nil {
		t.Errorf("isEnd with empty buffer returned no error, but should")
	}
}

func TestReadList(t *testing.T) {
	const i = "4:testi4ee"
	o := []interface{}{"test", int64(4)}
	dec := getDecoder([]byte(i))

	list, err := dec.readList()
	if err != nil {
		t.Errorf("readList returned error: %v", err)
	}

	if !reflect.DeepEqual(list, o) {
		t.Errorf("readList returned %v, but should return %v", list, o)
	}

	dec = getDecoder([]byte(""))
	if _, err := dec.readList(); err == nil {
		t.Errorf("readList with empty buffer returned no error, but should")
	}

	dec = getDecoder([]byte("4:test"))
	if _, err := dec.readList(); err == nil {
		t.Errorf("readList with no 'e'nd returned no error, but should")
	}
}

func TestReadDictionary(t *testing.T) {
	const i = "4:testi1337e3:foo3:bare"
	o := map[string]interface{}{
		"test": int64(1337),
		"foo":  "bar",
	}
	dec := getDecoder([]byte(i))

	dict, err := dec.readDictionary()
	if err != nil {
		t.Errorf("readDictionary returned error: %v", err)
	}

	if !reflect.DeepEqual(dict, o) {
		t.Errorf("readDictionary returned %v, but should %v", dict, o)
	}

	dec = getDecoder([]byte(""))
	if _, err := dec.readDictionary(); err == nil {
		t.Errorf("readDict with empty buffer returned no error, but should")
	}

	dec = getDecoder([]byte(":"))
	if _, err := dec.readDictionary(); err == nil {
		t.Errorf("readDict with invalid key returned no error, but should")
	}

	dec = getDecoder([]byte("3:fooe"))
	if _, err := dec.readDictionary(); err == nil {
		t.Errorf("readDict with no value returned no error, but should")
	}

	dec = getDecoder([]byte("3:foo"))
	if _, err := dec.readDictionary(); err == nil {
		t.Errorf("readDict with only one key and no value returned no error, but should")
	}

	dec = getDecoder([]byte("3:foo3:bar"))
	if _, err := dec.readDictionary(); err == nil {
		t.Errorf("readDict with no 'e'nd returned no error, but should")
	}
}

func TestDecode(t *testing.T) {
	const i = "d4:testli4ei3ee3:foo3:bare"
	o := map[string]interface{}{
		"test": []interface{}{
			int64(4),
			int64(3),
		},
		"foo": "bar",
	}
	buf := bytes.NewBufferString(i)

	res, err := Decode(buf)
	if err != nil {
		t.Errorf("Decode returned error: %v", err)
	}

	if !reflect.DeepEqual(res, o) {
		t.Errorf("Decode returned %v, but should return %v", res, o)
	}

	buf = bytes.NewBufferString("")
	if _, err := Decode(buf); err == nil {
		t.Errorf("Decode with empty buffer returned no error, but should")
	}

	buf = bytes.NewBufferString("3:foo")
	if _, err := Decode(buf); err == nil {
		t.Errorf("Decode with other value than dictionary returned no error, but should")
	}
}
<commit_msg>Format if statements inline for better readability<commit_after>package octorrent

import (
	"bufio"
	"bytes"
	"math"
	"reflect"
	"strconv"
	"testing"
)

func getDecoder(d []byte) decoder {
	buf := bytes.NewBuffer(d)

	return decoder{
		*bufio.NewReader(buf),
	}
}

func TestReadInt(t *testing.T) {
	i, o := "412e51", int64(412)
	dec := getDecoder([]byte(i))

	if num, err := dec.readInt(); err != nil {
		t.Errorf("readInt returned error: %v", err)
	} else if num != o {
		t.Errorf("readInt did return %i, but should return %i", num, o)
	}
}

func TestReadIntUntil(t *testing.T) {
	i, o := "142", int64(14)
	dec := getDecoder([]byte(i))

	if num, err := dec.readIntUntil('2'); err != nil {
		t.Errorf("readIntUntil returned error: %v", err)
	} else if num != o {
		t.Errorf("readIntUntil did return %i, but should return %i", num, o)
	}

	// Check broken reader
	dec = getDecoder([]byte(""))
	if _, err := dec.readIntUntil('t'); err == nil {
		t.Errorf("readIntUntil with empty reader returned no error, but should")
	}

	dec = getDecoder([]byte("test"))
	if _, err := dec.readIntUntil('e'); err == nil {
		t.Errorf("readIntUntil with no digits returned no error, but should")
	}

	// Check unsigned int
	var ui uint64
	ui = math.MaxInt64 + 1
	i = "i" + strconv.FormatUint(ui, 10) + "e"
	dec = getDecoder([]byte(i))
	if _, err := dec.readIntUntil('e'); err == nil {
		t.Errorf("readIntUntil with unsigned int returned no error, but should")
	}
}

func TestReadString(t *testing.T) {
	i, o := "4:testfoobar", "test"
	dec := getDecoder([]byte(i))
	if str, err := dec.readString(); err != nil {
		t.Errorf("readString returned error: %v", err)
	} else if str != o {
		t.Errorf("readString returned %s, but should return %s", str, o)
	}

	// Test broken reader
	dec = getDecoder([]byte(""))
	if _, err := dec.readString(); err == nil {
		t.Errorf("readSstring with empty buffer returned no error, but should")
	}

	// Test wrong format
	dec = getDecoder([]byte("4:foo"))
	if _, err := dec.readString(); err == nil {
		t.Errorf("readString with wrong formatted string returned no error, but should")
	}
}

func TestReadValue(t *testing.T) {
	dec := getDecoder([]byte(""))
	if _, err := dec.readValue(); err == nil {
		t.Errorf("readValue with empty buffer returned no error, but should")
	}

	dec = getDecoder([]byte("ifooe"))
	if _, err := dec.readValue(); err == nil {
		t.Errorf("readValue with wrongly formatted buffer returned no error, but should")
	}
}

func TestIsEnd(t *testing.T) {
	i, o := "4:test", false
	dec := getDecoder([]byte(i))
	if end, err := dec.isEnd(); err != nil {
		t.Errorf("isEnd returned error: %v", err)
	} else if end != o {
		t.Errorf("isEnd returned %v, but should return %v", end, o)
	}

	i, o = "e4:test", true
	dec = getDecoder([]byte(i))
	if end, err := dec.isEnd(); err != nil {
		t.Errorf("isEnd returned error: %v", err)
	} else if end != o {
		t.Errorf("isEnd returned %v, but should return %v", end, o)
	}

	dec = getDecoder([]byte(""))
	if _, err := dec.isEnd(); err == nil {
		t.Errorf("isEnd with empty buffer returned no error, but should")
	}
}

func TestReadList(t *testing.T) {
	i := "4:testi4ee"
	o := []interface{}{"test", int64(4)}
	dec := getDecoder([]byte(i))

	if list, err := dec.readList(); err != nil {
		t.Errorf("readList returned error: %v", err)
	} else if !reflect.DeepEqual(list, o) {
		t.Errorf("readList returned %v, but should return %v", list, o)
	}

	dec = getDecoder([]byte(""))
	if _, err := dec.readList(); err == nil {
		t.Errorf("readList with empty buffer returned no error, but should")
	}

	dec = getDecoder([]byte("4:test"))
	if _, err := dec.readList(); err == nil {
		t.Errorf("readList with no 'e'nd returned no error, but should")
	}
}

func TestReadDictionary(t *testing.T) {
	i := "4:testi1337e3:foo3:bare"
	o := map[string]interface{}{
		"test": int64(1337),
		"foo":  "bar",
	}
	dec := getDecoder([]byte(i))

	if dict, err := dec.readDictionary(); err != nil {
		t.Errorf("readDictionary returned error: %v", err)
	} else if !reflect.DeepEqual(dict, o) {
		t.Errorf("readDictionary returned %v, but should %v", dict, o)
	}

	dec = getDecoder([]byte(""))
	if _, err := dec.readDictionary(); err == nil {
		t.Errorf("readDict with empty buffer returned no error, but should")
	}

	dec = getDecoder([]byte(":"))
	if _, err := dec.readDictionary(); err == nil {
		t.Errorf("readDict with invalid key returned no error, but should")
	}

	dec = getDecoder([]byte("3:fooe"))
	if _, err := dec.readDictionary(); err == nil {
		t.Errorf("readDict with no value returned no error, but should")
	}

	dec = getDecoder([]byte("3:foo"))
	if _, err := dec.readDictionary(); err == nil {
		t.Errorf("readDict with only one key and no value returned no error, but should")
	}

	dec = getDecoder([]byte("3:foo3:bar"))
	if _, err := dec.readDictionary(); err == nil {
		t.Errorf("readDict with no 'e'nd returned no error, but should")
	}
}

func TestDecode(t *testing.T) {
	const i = "d4:testli4ei3ee3:foo3:bare"
	o := map[string]interface{}{
		"test": []interface{}{
			int64(4),
			int64(3),
		},
		"foo": "bar",
	}
	buf := bytes.NewBufferString(i)

	if res, err := Decode(buf); err != nil {
		t.Errorf("Decode returned error: %v", err)
	} else if !reflect.DeepEqual(res, o) {
		t.Errorf("Decode returned %v, but should return %v", res, o)
	}

	buf = bytes.NewBufferString("")
	if _, err := Decode(buf); err == nil {
		t.Errorf("Decode with empty buffer returned no error, but should")
	}

	buf = bytes.NewBufferString("3:foo")
	if _, err := Decode(buf); err == nil {
		t.Errorf("Decode with other value than dictionary returned no error, but should")
	}
}
<|endoftext|>