<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 kube // import "k8s.io/helm/pkg/kube"

import (
	"time"

	appsv1beta1 "k8s.io/api/apps/v1beta1"
	appsv1beta2 "k8s.io/api/apps/v1beta2"
	"k8s.io/api/core/v1"
	extensions "k8s.io/api/extensions/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/pkg/api/v1/helper"
	podutil "k8s.io/kubernetes/pkg/api/v1/pod"
	deploymentutil "k8s.io/kubernetes/pkg/controller/deployment/util"
)

// deployment holds associated replicaSets for a deployment
type deployment struct {
	replicaSets *extensions.ReplicaSet
	deployment  *extensions.Deployment
}

// waitForResources polls to get the current status of all pods, PVCs, and Services
// until all are ready or a timeout is reached
func (c *Client) waitForResources(timeout time.Duration, created Result) error {
	c.Log("beginning wait for %d resources with timeout of %v", len(created), timeout)

	kcs, err := c.KubernetesClientSet()
	if err != nil {
		return err
	}
	return wait.Poll(2*time.Second, timeout, func() (bool, error) {
		pods := []v1.Pod{}
		services := []v1.Service{}
		pvc := []v1.PersistentVolumeClaim{}
		deployments := []deployment{}
		for _, v := range created {
			obj, err := c.AsVersionedObject(v.Object)
			if err != nil && !runtime.IsNotRegisteredError(err) {
				return false, err
			}
			switch value := obj.(type) {
			case (*v1.ReplicationController):
				list, err := getPods(kcs, value.Namespace, value.Spec.Selector)
				if err != nil {
					return false, err
				}
				pods = append(pods, list...)
			case (*v1.Pod):
				pod, err := kcs.Core().Pods(value.Namespace).Get(value.Name, metav1.GetOptions{})
				if err != nil {
					return false, err
				}
				pods = append(pods, *pod)
			case (*extensions.Deployment):
				currentDeployment, err := kcs.Extensions().Deployments(value.Namespace).Get(value.Name, metav1.GetOptions{})
				if err != nil {
					return false, err
				}
				// Find RS associated with deployment
				newReplicaSet, err := deploymentutil.GetNewReplicaSet(currentDeployment, kcs.ExtensionsV1beta1())
				if err != nil || newReplicaSet == nil {
					return false, err
				}
				newDeployment := deployment{
					newReplicaSet,
					currentDeployment,
				}
				deployments = append(deployments, newDeployment)
			case (*extensions.DaemonSet):
				list, err := getPods(kcs, value.Namespace, value.Spec.Selector.MatchLabels)
				if err != nil {
					return false, err
				}
				pods = append(pods, list...)
			case (*appsv1beta1.StatefulSet):
				list, err := getPods(kcs, value.Namespace, value.Spec.Selector.MatchLabels)
				if err != nil {
					return false, err
				}
				pods = append(pods, list...)
			case (*appsv1beta2.StatefulSet):
				list, err := getPods(kcs, value.Namespace, value.Spec.Selector.MatchLabels)
				if err != nil {
					return false, err
				}
				pods = append(pods, list...)
			case (*extensions.ReplicaSet):
				list, err := getPods(kcs, value.Namespace, value.Spec.Selector.MatchLabels)
				if err != nil {
					return false, err
				}
				pods = append(pods, list...)
			case (*v1.PersistentVolumeClaim):
				claim, err := kcs.Core().PersistentVolumeClaims(value.Namespace).Get(value.Name, metav1.GetOptions{})
				if err != nil {
					return false, err
				}
				pvc = append(pvc, *claim)
			case (*v1.Service):
				svc, err := kcs.Core().Services(value.Namespace).Get(value.Name, metav1.GetOptions{})
				if err != nil {
					return false, err
				}
				services = append(services, *svc)
			}
		}
		isReady := c.podsReady(pods) && c.servicesReady(services) && c.volumesReady(pvc) && c.deploymentsReady(deployments)
		return isReady, nil
	})
}

func (c *Client) podsReady(pods []v1.Pod) bool {
	for _, pod := range pods {
		if !podutil.IsPodReady(&pod) {
			c.Log("Pod is not ready: %s/%s", pod.GetNamespace(), pod.GetName())
			return false
		}
	}
	return true
}

func (c *Client) servicesReady(svc []v1.Service) bool {
	for _, s := range svc {
		// ExternalName Services are external to cluster so helm shouldn't be checking to see if they're 'ready' (i.e. have an IP Set)
		if s.Spec.Type == v1.ServiceTypeExternalName {
			continue
		}

		// Make sure the service is not explicitly set to "None" before checking the IP
		if s.Spec.ClusterIP != v1.ClusterIPNone && !helper.IsServiceIPSet(&s) {
			c.Log("Service is not ready: %s/%s", s.GetNamespace(), s.GetName())
			return false
		}
		// This checks if the service has a LoadBalancer and that balancer has an Ingress defined
		if s.Spec.Type == v1.ServiceTypeLoadBalancer && s.Status.LoadBalancer.Ingress == nil {
			c.Log("Service is not ready: %s/%s", s.GetNamespace(), s.GetName())
			return false
		}
	}
	return true
}

func (c *Client) volumesReady(vols []v1.PersistentVolumeClaim) bool {
	for _, v := range vols {
		if v.Status.Phase != v1.ClaimBound {
			c.Log("PersistentVolumeClaim is not ready: %s/%s", v.GetNamespace(), v.GetName())
			return false
		}
	}
	return true
}

func (c *Client) deploymentsReady(deployments []deployment) bool {
	for _, v := range deployments {
		if !(v.replicaSets.Status.ReadyReplicas >= *v.deployment.Spec.Replicas-deploymentutil.MaxUnavailable(*v.deployment)) {
			c.Log("Deployment is not ready: %s/%s", v.deployment.GetNamespace(), v.deployment.GetName())
			return false
		}
	}
	return true
}

func getPods(client kubernetes.Interface, namespace string, selector map[string]string) ([]v1.Pod, error) {
	list, err := client.Core().Pods(namespace).List(metav1.ListOptions{
		FieldSelector: fields.Everything().String(),
		LabelSelector: labels.Set(selector).AsSelector().String(),
	})
	return list.Items, err
}
<commit_msg>Remove unnecessary parentheses in switch statements (#3175)<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 kube // import "k8s.io/helm/pkg/kube"

import (
	"time"

	appsv1beta1 "k8s.io/api/apps/v1beta1"
	appsv1beta2 "k8s.io/api/apps/v1beta2"
	"k8s.io/api/core/v1"
	extensions "k8s.io/api/extensions/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/pkg/api/v1/helper"
	podutil "k8s.io/kubernetes/pkg/api/v1/pod"
	deploymentutil "k8s.io/kubernetes/pkg/controller/deployment/util"
)

// deployment holds associated replicaSets for a deployment
type deployment struct {
	replicaSets *extensions.ReplicaSet
	deployment  *extensions.Deployment
}

// waitForResources polls to get the current status of all pods, PVCs, and Services
// until all are ready or a timeout is reached
func (c *Client) waitForResources(timeout time.Duration, created Result) error {
	c.Log("beginning wait for %d resources with timeout of %v", len(created), timeout)

	kcs, err := c.KubernetesClientSet()
	if err != nil {
		return err
	}
	return wait.Poll(2*time.Second, timeout, func() (bool, error) {
		pods := []v1.Pod{}
		services := []v1.Service{}
		pvc := []v1.PersistentVolumeClaim{}
		deployments := []deployment{}
		for _, v := range created {
			obj, err := c.AsVersionedObject(v.Object)
			if err != nil && !runtime.IsNotRegisteredError(err) {
				return false, err
			}
			switch value := obj.(type) {
			case *v1.ReplicationController:
				list, err := getPods(kcs, value.Namespace, value.Spec.Selector)
				if err != nil {
					return false, err
				}
				pods = append(pods, list...)
			case *v1.Pod:
				pod, err := kcs.Core().Pods(value.Namespace).Get(value.Name, metav1.GetOptions{})
				if err != nil {
					return false, err
				}
				pods = append(pods, *pod)
			case *extensions.Deployment:
				currentDeployment, err := kcs.Extensions().Deployments(value.Namespace).Get(value.Name, metav1.GetOptions{})
				if err != nil {
					return false, err
				}
				// Find RS associated with deployment
				newReplicaSet, err := deploymentutil.GetNewReplicaSet(currentDeployment, kcs.ExtensionsV1beta1())
				if err != nil || newReplicaSet == nil {
					return false, err
				}
				newDeployment := deployment{
					newReplicaSet,
					currentDeployment,
				}
				deployments = append(deployments, newDeployment)
			case *extensions.DaemonSet:
				list, err := getPods(kcs, value.Namespace, value.Spec.Selector.MatchLabels)
				if err != nil {
					return false, err
				}
				pods = append(pods, list...)
			case *appsv1beta1.StatefulSet:
				list, err := getPods(kcs, value.Namespace, value.Spec.Selector.MatchLabels)
				if err != nil {
					return false, err
				}
				pods = append(pods, list...)
			case *appsv1beta2.StatefulSet:
				list, err := getPods(kcs, value.Namespace, value.Spec.Selector.MatchLabels)
				if err != nil {
					return false, err
				}
				pods = append(pods, list...)
			case *extensions.ReplicaSet:
				list, err := getPods(kcs, value.Namespace, value.Spec.Selector.MatchLabels)
				if err != nil {
					return false, err
				}
				pods = append(pods, list...)
			case *v1.PersistentVolumeClaim:
				claim, err := kcs.Core().PersistentVolumeClaims(value.Namespace).Get(value.Name, metav1.GetOptions{})
				if err != nil {
					return false, err
				}
				pvc = append(pvc, *claim)
			case *v1.Service:
				svc, err := kcs.Core().Services(value.Namespace).Get(value.Name, metav1.GetOptions{})
				if err != nil {
					return false, err
				}
				services = append(services, *svc)
			}
		}
		isReady := c.podsReady(pods) && c.servicesReady(services) && c.volumesReady(pvc) && c.deploymentsReady(deployments)
		return isReady, nil
	})
}

func (c *Client) podsReady(pods []v1.Pod) bool {
	for _, pod := range pods {
		if !podutil.IsPodReady(&pod) {
			c.Log("Pod is not ready: %s/%s", pod.GetNamespace(), pod.GetName())
			return false
		}
	}
	return true
}

func (c *Client) servicesReady(svc []v1.Service) bool {
	for _, s := range svc {
		// ExternalName Services are external to cluster so helm shouldn't be checking to see if they're 'ready' (i.e. have an IP Set)
		if s.Spec.Type == v1.ServiceTypeExternalName {
			continue
		}

		// Make sure the service is not explicitly set to "None" before checking the IP
		if s.Spec.ClusterIP != v1.ClusterIPNone && !helper.IsServiceIPSet(&s) {
			c.Log("Service is not ready: %s/%s", s.GetNamespace(), s.GetName())
			return false
		}
		// This checks if the service has a LoadBalancer and that balancer has an Ingress defined
		if s.Spec.Type == v1.ServiceTypeLoadBalancer && s.Status.LoadBalancer.Ingress == nil {
			c.Log("Service is not ready: %s/%s", s.GetNamespace(), s.GetName())
			return false
		}
	}
	return true
}

func (c *Client) volumesReady(vols []v1.PersistentVolumeClaim) bool {
	for _, v := range vols {
		if v.Status.Phase != v1.ClaimBound {
			c.Log("PersistentVolumeClaim is not ready: %s/%s", v.GetNamespace(), v.GetName())
			return false
		}
	}
	return true
}

func (c *Client) deploymentsReady(deployments []deployment) bool {
	for _, v := range deployments {
		if !(v.replicaSets.Status.ReadyReplicas >= *v.deployment.Spec.Replicas-deploymentutil.MaxUnavailable(*v.deployment)) {
			c.Log("Deployment is not ready: %s/%s", v.deployment.GetNamespace(), v.deployment.GetName())
			return false
		}
	}
	return true
}

func getPods(client kubernetes.Interface, namespace string, selector map[string]string) ([]v1.Pod, error) {
	list, err := client.Core().Pods(namespace).List(metav1.ListOptions{
		FieldSelector: fields.Everything().String(),
		LabelSelector: labels.Set(selector).AsSelector().String(),
	})
	return list.Items, err
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 handlers define HTTP handlers.
package site

import (
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/google/triage-party/pkg/hubbub"
	"github.com/google/triage-party/pkg/triage"
	"github.com/google/triage-party/pkg/updater"

	"github.com/dustin/go-humanize"
	"github.com/google/go-github/v31/github"
	"gopkg.in/yaml.v2"

	"k8s.io/klog"
)

const VERSION = "2020-04-22.01"

var (
	nonWordRe = regexp.MustCompile(`\W`)
)

// Config is how external users interact with this package.
type Config struct {
	BaseDirectory string
	Name          string
	WarnAge       time.Duration
	Updater       *updater.Updater
	Party         *triage.Party
}

func New(c *Config) *Handlers {
	return &Handlers{
		baseDir:  c.BaseDirectory,
		updater:  c.Updater,
		party:    c.Party,
		siteName: c.Name,
		warnAge:  c.WarnAge,
	}
}

// Handlers is a mix of config and client interfaces to connect with.
type Handlers struct {
	baseDir  string
	updater  *updater.Updater
	party    *triage.Party
	siteName string
	warnAge  time.Duration
}

// Root redirects to leaderboard.
func (h *Handlers) Root() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		sts, err := h.party.ListCollections()
		if err != nil {
			klog.Errorf("collections: %v", err)
			return
		}
		http.Redirect(w, r, fmt.Sprintf("/s/%s", sts[0].ID), http.StatusSeeOther)
	}
}

// Page are values that are passed into the renderer
type Page struct {
	Version     string
	SiteName    string
	ID          string
	Title       string
	Description string
	Warning     string
	Total       int
	TotalShown  int
	Types       string
	UniqueItems []*hubbub.Conversation

	Player        int
	Players       int
	PlayerChoices []string
	Mode          int
	Index         int
	EmbedURL      string

	AverageResponseLatency time.Duration
	TotalPullRequests      int
	TotalIssues            int

	Collection  triage.Collection
	Collections []triage.Collection

	CollectionResult *triage.CollectionResult
	Stats            *triage.CollectionResult
	StatsID          string
}

// is this request an HTTP refresh?
func isRefresh(r *http.Request) bool {
	cc := r.Header["Cache-Control"]
	if len(cc) == 0 {
		return false
	}
	//	klog.Infof("cc=%s headers=%+v", cc, r.Header)
	return cc[0] == "max-age-0" || cc[0] == "no-cache"
}

// Board shows a stratgy board
func (h *Handlers) Collection() http.HandlerFunc {
	fmap := template.FuncMap{
		"toJS":          toJS,
		"toYAML":        toYAML,
		"toJSfunc":      toJSfunc,
		"toDays":        toDays,
		"HumanDuration": humanDuration,
		"HumanTime":     humanTime,
		"UnixNano":      unixNano,
		"Avatar":        avatar,
	}
	t := template.Must(template.New("collection").Funcs(fmap).ParseFiles(
		filepath.Join(h.baseDir, "collection.tmpl"),
		filepath.Join(h.baseDir, "base.tmpl"),
	))

	return func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		defer func() {
			klog.Infof("Collection request complete in %s", time.Since(start))
		}()
		id := strings.TrimPrefix(r.URL.Path, "/s/")
		playerChoices := []string{"Select a player"}
		players := getInt(r.URL, "players", 1)
		player := getInt(r.URL, "player", 0)
		mode := getInt(r.URL, "mode", 0)
		index := getInt(r.URL, "index", 1)

		for i := 0; i < players; i++ {
			playerChoices = append(playerChoices, fmt.Sprintf("Player %d", i+1))
		}

		klog.Infof("GET %s (%q): %v", r.URL.Path, id, r.Header)
		s, err := h.party.LookupCollection(id)
		if err != nil {
			http.Error(w, fmt.Sprintf("%q not found: old link or typo?", id), http.StatusNotFound)
			klog.Errorf("collection: %v", err)
			return
		}

		sts, err := h.party.ListCollections()
		if err != nil {
			klog.Errorf("collections: %v", err)
			http.Error(w, "list error", http.StatusInternalServerError)
			return
		}

		var result *triage.CollectionResult
		if isRefresh(r) {
			result = h.updater.ForceRefresh(r.Context(), id)
			klog.Infof("refresh %q result: %d items", id, len(result.RuleResults))
		} else {
			result = h.updater.Lookup(r.Context(), id, true)
			if result == nil {
				http.Error(w, fmt.Sprintf("%q no data", id), http.StatusNotFound)
				return
			}
			if result.RuleResults == nil {
				http.Error(w, fmt.Sprintf("%q no outcomes", id), http.StatusNotFound)
				return
			}

			klog.Infof("lookup %q result: %d items", id, len(result.RuleResults))
		}

		warning := ""
		if time.Since(result.Time) > h.warnAge {
			warning = fmt.Sprintf("Serving stale results (%s old) - refreshing results in background. Use Shift-Reload to force data to refresh at any time.", time.Since(result.Time))
		}

		total := 0
		for _, o := range result.RuleResults {
			total += len(o.Items)
		}

		unique := []*hubbub.Conversation{}
		seen := map[int]bool{}
		for _, o := range result.RuleResults {
			for _, i := range o.Items {
				if !seen[i.ID] {
					unique = append(unique, i)
					seen[i.ID] = true
				}
			}
		}

		if player > 0 && players > 1 {
			result = playerFilter(result, player, players)
		}

		uniqueFiltered := []*hubbub.Conversation{}
		seenFiltered := map[int]bool{}
		for _, o := range result.RuleResults {
			for _, i := range o.Items {
				if !seenFiltered[i.ID] {
					uniqueFiltered = append(uniqueFiltered, i)
					seenFiltered[i.ID] = true
				}
			}
		}

		embedURL := ""
		if mode == 1 {
			searchIndex := 0
			for _, o := range result.RuleResults {
				for _, i := range o.Items {
					searchIndex++
					if searchIndex == index {
						embedURL = i.URL
					}
				}
			}
		}
		p := &Page{
			ID:               s.ID,
			Version:          VERSION,
			SiteName:         h.siteName,
			Title:            s.Name,
			Collection:       s,
			Collections:      sts,
			Description:      s.Description,
			CollectionResult: result,
			Total:            len(unique),
			TotalShown:       len(uniqueFiltered),
			Types:            "Issues",
			PlayerChoices:    playerChoices,
			Player:           player,
			Players:          players,
			Mode:             mode,
			Index:            index,
			EmbedURL:         embedURL,
			Warning:          warning,
			UniqueItems:      uniqueFiltered,
		}

		for _, s := range sts {
			if s.UsedForStats {
				p.Stats = h.updater.Lookup(r.Context(), s.ID, false)
				p.StatsID = s.ID
			}
		}

		klog.V(2).Infof("page context: %+v", p)
		err = t.ExecuteTemplate(w, "base", p)
		if err != nil {
			klog.Errorf("tmpl: %v", err)
			return
		}
	}
}

// helper to get integers from a URL
func getInt(url *url.URL, key string, fallback int) int {
	vals := url.Query()[key]
	if len(vals) == 1 {
		i, err := strconv.ParseInt(vals[0], 10, 32)
		if err != nil {
			klog.Warningf("bad %s int value: %v", key, vals)
			return fallback
		}
		return int(i)
	}
	return fallback
}

func toYAML(v interface{}) string {
	s, err := yaml.Marshal(v)
	if err != nil {
		return fmt.Sprintf("yaml err: %v", err)
	}
	return string(s)
}

// Acknowledge JS sanitization issues (what data do we trust?)
func toJS(s string) template.JS {
	return template.JS(s)
}

// Acknowledge JS sanitization issues (what data do we trust?)
func toJSfunc(s string) template.JS {
	return template.JS(nonWordRe.ReplaceAllString(s, "_"))
}

func unixNano(t time.Time) int64 {
	return t.UnixNano()
}

func humanDuration(d time.Duration) string {
	return humanTime(time.Now().Add(-d))
}

func toDays(d time.Duration) string {
	return fmt.Sprintf("%0.1fd", d.Hours()/24)
}

func humanTime(t time.Time) string {
	if t.IsZero() {
		return ""
	}
	ds := humanize.Time(t)
	ds = strings.Replace(ds, " ago", "", 1)

	ds = strings.Replace(ds, " minutes", "min", 1)
	ds = strings.Replace(ds, " minute", "min", 1)

	ds = strings.Replace(ds, " hours", "h", 1)
	ds = strings.Replace(ds, " hour", "h", 1)

	ds = strings.Replace(ds, " days", "d", 1)
	ds = strings.Replace(ds, " day", "d", 1)

	ds = strings.Replace(ds, " months", "mo", 1)
	ds = strings.Replace(ds, " month", "mo", 1)

	ds = strings.Replace(ds, " years", "y", 1)
	ds = strings.Replace(ds, " year", "y", 1)

	ds = strings.Replace(ds, " weeks", "wk", 1)
	ds = strings.Replace(ds, " week", "wk", 1)

	return ds
}

func avatar(u *github.User) template.HTML {
	return template.HTML(fmt.Sprintf(`<a href="%s" title="%s"><img src="%s" width="20" height="20"></a>`, u.GetHTMLURL(), u.GetLogin(), u.GetAvatarURL()))
}

// playerFilter filters out results for a particular player
func playerFilter(result *triage.CollectionResult, player int, players int) *triage.CollectionResult {
	klog.Infof("Filtering for player %d of %d ...", player, players)
	os := []triage.RuleResult{}

	seen := map[string]*triage.Rule{}

	for _, o := range result.RuleResults {
		cs := []*hubbub.Conversation{}

		for _, i := range o.Items {
			if (i.ID % players) == (player - 1) {
				klog.Infof("%d belongs to player %d", i.ID, player)
				cs = append(cs, i)
			}
		}

		os = append(os, triage.SummarizeRuleResult(o.Rule, cs, seen))

	}
	return triage.SummarizeCollectionResult(os)
}
<commit_msg>Increment version number<commit_after>// Copyright 2020 Google LLC
//
// 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 handlers define HTTP handlers.
package site

import (
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/google/triage-party/pkg/hubbub"
	"github.com/google/triage-party/pkg/triage"
	"github.com/google/triage-party/pkg/updater"

	"github.com/dustin/go-humanize"
	"github.com/google/go-github/v31/github"
	"gopkg.in/yaml.v2"

	"k8s.io/klog"
)

const VERSION = "2020-04-27.00"

var (
	nonWordRe = regexp.MustCompile(`\W`)
)

// Config is how external users interact with this package.
type Config struct {
	BaseDirectory string
	Name          string
	WarnAge       time.Duration
	Updater       *updater.Updater
	Party         *triage.Party
}

func New(c *Config) *Handlers {
	return &Handlers{
		baseDir:  c.BaseDirectory,
		updater:  c.Updater,
		party:    c.Party,
		siteName: c.Name,
		warnAge:  c.WarnAge,
	}
}

// Handlers is a mix of config and client interfaces to connect with.
type Handlers struct {
	baseDir  string
	updater  *updater.Updater
	party    *triage.Party
	siteName string
	warnAge  time.Duration
}

// Root redirects to leaderboard.
func (h *Handlers) Root() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		sts, err := h.party.ListCollections()
		if err != nil {
			klog.Errorf("collections: %v", err)
			return
		}
		http.Redirect(w, r, fmt.Sprintf("/s/%s", sts[0].ID), http.StatusSeeOther)
	}
}

// Page are values that are passed into the renderer
type Page struct {
	Version     string
	SiteName    string
	ID          string
	Title       string
	Description string
	Warning     string
	Total       int
	TotalShown  int
	Types       string
	UniqueItems []*hubbub.Conversation

	Player        int
	Players       int
	PlayerChoices []string
	Mode          int
	Index         int
	EmbedURL      string

	AverageResponseLatency time.Duration
	TotalPullRequests      int
	TotalIssues            int

	Collection  triage.Collection
	Collections []triage.Collection

	CollectionResult *triage.CollectionResult
	Stats            *triage.CollectionResult
	StatsID          string
}

// is this request an HTTP refresh?
func isRefresh(r *http.Request) bool {
	cc := r.Header["Cache-Control"]
	if len(cc) == 0 {
		return false
	}
	//	klog.Infof("cc=%s headers=%+v", cc, r.Header)
	return cc[0] == "max-age-0" || cc[0] == "no-cache"
}

// Board shows a stratgy board
func (h *Handlers) Collection() http.HandlerFunc {
	fmap := template.FuncMap{
		"toJS":          toJS,
		"toYAML":        toYAML,
		"toJSfunc":      toJSfunc,
		"toDays":        toDays,
		"HumanDuration": humanDuration,
		"HumanTime":     humanTime,
		"UnixNano":      unixNano,
		"Avatar":        avatar,
	}
	t := template.Must(template.New("collection").Funcs(fmap).ParseFiles(
		filepath.Join(h.baseDir, "collection.tmpl"),
		filepath.Join(h.baseDir, "base.tmpl"),
	))

	return func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		defer func() {
			klog.Infof("Collection request complete in %s", time.Since(start))
		}()
		id := strings.TrimPrefix(r.URL.Path, "/s/")
		playerChoices := []string{"Select a player"}
		players := getInt(r.URL, "players", 1)
		player := getInt(r.URL, "player", 0)
		mode := getInt(r.URL, "mode", 0)
		index := getInt(r.URL, "index", 1)

		for i := 0; i < players; i++ {
			playerChoices = append(playerChoices, fmt.Sprintf("Player %d", i+1))
		}

		klog.Infof("GET %s (%q): %v", r.URL.Path, id, r.Header)
		s, err := h.party.LookupCollection(id)
		if err != nil {
			http.Error(w, fmt.Sprintf("%q not found: old link or typo?", id), http.StatusNotFound)
			klog.Errorf("collection: %v", err)
			return
		}

		sts, err := h.party.ListCollections()
		if err != nil {
			klog.Errorf("collections: %v", err)
			http.Error(w, "list error", http.StatusInternalServerError)
			return
		}

		var result *triage.CollectionResult
		if isRefresh(r) {
			result = h.updater.ForceRefresh(r.Context(), id)
			klog.Infof("refresh %q result: %d items", id, len(result.RuleResults))
		} else {
			result = h.updater.Lookup(r.Context(), id, true)
			if result == nil {
				http.Error(w, fmt.Sprintf("%q no data", id), http.StatusNotFound)
				return
			}
			if result.RuleResults == nil {
				http.Error(w, fmt.Sprintf("%q no outcomes", id), http.StatusNotFound)
				return
			}

			klog.Infof("lookup %q result: %d items", id, len(result.RuleResults))
		}

		warning := ""
		if time.Since(result.Time) > h.warnAge {
			warning = fmt.Sprintf("Serving stale results (%s old) - refreshing results in background. Use Shift-Reload to force data to refresh at any time.", time.Since(result.Time))
		}

		total := 0
		for _, o := range result.RuleResults {
			total += len(o.Items)
		}

		unique := []*hubbub.Conversation{}
		seen := map[int]bool{}
		for _, o := range result.RuleResults {
			for _, i := range o.Items {
				if !seen[i.ID] {
					unique = append(unique, i)
					seen[i.ID] = true
				}
			}
		}

		if player > 0 && players > 1 {
			result = playerFilter(result, player, players)
		}

		uniqueFiltered := []*hubbub.Conversation{}
		seenFiltered := map[int]bool{}
		for _, o := range result.RuleResults {
			for _, i := range o.Items {
				if !seenFiltered[i.ID] {
					uniqueFiltered = append(uniqueFiltered, i)
					seenFiltered[i.ID] = true
				}
			}
		}

		embedURL := ""
		if mode == 1 {
			searchIndex := 0
			for _, o := range result.RuleResults {
				for _, i := range o.Items {
					searchIndex++
					if searchIndex == index {
						embedURL = i.URL
					}
				}
			}
		}
		p := &Page{
			ID:               s.ID,
			Version:          VERSION,
			SiteName:         h.siteName,
			Title:            s.Name,
			Collection:       s,
			Collections:      sts,
			Description:      s.Description,
			CollectionResult: result,
			Total:            len(unique),
			TotalShown:       len(uniqueFiltered),
			Types:            "Issues",
			PlayerChoices:    playerChoices,
			Player:           player,
			Players:          players,
			Mode:             mode,
			Index:            index,
			EmbedURL:         embedURL,
			Warning:          warning,
			UniqueItems:      uniqueFiltered,
		}

		for _, s := range sts {
			if s.UsedForStats {
				p.Stats = h.updater.Lookup(r.Context(), s.ID, false)
				p.StatsID = s.ID
			}
		}

		klog.V(2).Infof("page context: %+v", p)
		err = t.ExecuteTemplate(w, "base", p)
		if err != nil {
			klog.Errorf("tmpl: %v", err)
			return
		}
	}
}

// helper to get integers from a URL
func getInt(url *url.URL, key string, fallback int) int {
	vals := url.Query()[key]
	if len(vals) == 1 {
		i, err := strconv.ParseInt(vals[0], 10, 32)
		if err != nil {
			klog.Warningf("bad %s int value: %v", key, vals)
			return fallback
		}
		return int(i)
	}
	return fallback
}

func toYAML(v interface{}) string {
	s, err := yaml.Marshal(v)
	if err != nil {
		return fmt.Sprintf("yaml err: %v", err)
	}
	return string(s)
}

// Acknowledge JS sanitization issues (what data do we trust?)
func toJS(s string) template.JS {
	return template.JS(s)
}

// Acknowledge JS sanitization issues (what data do we trust?)
func toJSfunc(s string) template.JS {
	return template.JS(nonWordRe.ReplaceAllString(s, "_"))
}

func unixNano(t time.Time) int64 {
	return t.UnixNano()
}

func humanDuration(d time.Duration) string {
	return humanTime(time.Now().Add(-d))
}

func toDays(d time.Duration) string {
	return fmt.Sprintf("%0.1fd", d.Hours()/24)
}

func humanTime(t time.Time) string {
	if t.IsZero() {
		return ""
	}
	ds := humanize.Time(t)
	ds = strings.Replace(ds, " ago", "", 1)

	ds = strings.Replace(ds, " minutes", "min", 1)
	ds = strings.Replace(ds, " minute", "min", 1)

	ds = strings.Replace(ds, " hours", "h", 1)
	ds = strings.Replace(ds, " hour", "h", 1)

	ds = strings.Replace(ds, " days", "d", 1)
	ds = strings.Replace(ds, " day", "d", 1)

	ds = strings.Replace(ds, " months", "mo", 1)
	ds = strings.Replace(ds, " month", "mo", 1)

	ds = strings.Replace(ds, " years", "y", 1)
	ds = strings.Replace(ds, " year", "y", 1)

	ds = strings.Replace(ds, " weeks", "wk", 1)
	ds = strings.Replace(ds, " week", "wk", 1)

	return ds
}

func avatar(u *github.User) template.HTML {
	return template.HTML(fmt.Sprintf(`<a href="%s" title="%s"><img src="%s" width="20" height="20"></a>`, u.GetHTMLURL(), u.GetLogin(), u.GetAvatarURL()))
}

// playerFilter filters out results for a particular player
func playerFilter(result *triage.CollectionResult, player int, players int) *triage.CollectionResult {
	klog.Infof("Filtering for player %d of %d ...", player, players)
	os := []triage.RuleResult{}

	seen := map[string]*triage.Rule{}

	for _, o := range result.RuleResults {
		cs := []*hubbub.Conversation{}

		for _, i := range o.Items {
			if (i.ID % players) == (player - 1) {
				klog.Infof("%d belongs to player %d", i.ID, player)
				cs = append(cs, i)
			}
		}

		os = append(os, triage.SummarizeRuleResult(o.Rule, cs, seen))

	}
	return triage.SummarizeCollectionResult(os)
}
<|endoftext|>


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

package maas

import (
	"net/http"

	"github.com/juju/errors"
	"github.com/juju/gomaasapi"
	jc "github.com/juju/testing/checkers"
	"github.com/juju/utils/set"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/feature"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/network"
	coretesting "github.com/juju/juju/testing"
)

type maas2EnvironSuite struct {
	baseProviderSuite
}

var _ = gc.Suite(&maas2EnvironSuite{})

func makeEnviron(c *gc.C) *maasEnviron {
	testAttrs := coretesting.Attrs{}
	for k, v := range maasEnvAttrs {
		testAttrs[k] = v
	}
	testAttrs["maas-server"] = "http://any-old-junk.invalid/"
	attrs := coretesting.FakeConfig().Merge(testAttrs)
	cfg, err := config.New(config.NoDefaults, attrs)
	c.Assert(err, jc.ErrorIsNil)
	env, err := NewEnviron(cfg)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(env, gc.NotNil)
	return env
}

func (suite *maas2EnvironSuite) SetUpTest(c *gc.C) {
	suite.baseProviderSuite.SetUpTest(c)
	suite.SetFeatureFlags(feature.MAAS2)
}

func (suite *maas2EnvironSuite) getEnvWithServer(c *gc.C) (*maasEnviron, error) {
	testServer := gomaasapi.NewSimpleServer()
	testServer.AddGetResponse("/api/2.0/version/", http.StatusOK, maas2VersionResponse)
	testServer.AddGetResponse("/api/2.0/users/?op=whoami", http.StatusOK, "{}")
	testServer.Start()
	suite.AddCleanup(func(*gc.C) { testServer.Close() })
	testAttrs := coretesting.Attrs{}
	for k, v := range maasEnvAttrs {
		testAttrs[k] = v
	}
	testAttrs["maas-server"] = testServer.Server.URL
	attrs := coretesting.FakeConfig().Merge(testAttrs)
	cfg, err := config.New(config.NoDefaults, attrs)
	c.Assert(err, jc.ErrorIsNil)
	return NewEnviron(cfg)
}

func (suite *maas2EnvironSuite) TestNewEnvironWithoutFeatureFlag(c *gc.C) {
	suite.SetFeatureFlags()
	_, err := suite.getEnvWithServer(c)
	c.Assert(err, jc.Satisfies, errors.IsNotSupported)
}

func (suite *maas2EnvironSuite) TestNewEnvironWithController(c *gc.C) {
	env, err := suite.getEnvWithServer(c)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(env, gc.NotNil)
}

func (suite *maas2EnvironSuite) TestSupportedArchitectures(c *gc.C) {
	controller := fakeController{
		bootResources: []gomaasapi.BootResource{
			&fakeBootResource{name: "wily", architecture: "amd64/blah"},
			&fakeBootResource{name: "wily", architecture: "amd64/something"},
			&fakeBootResource{name: "xenial", architecture: "arm/somethingelse"},
		},
	}
	suite.injectController(&controller)
	env := makeEnviron(c)
	result, err := env.SupportedArchitectures()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, jc.DeepEquals, []string{"amd64", "arm"})
}

func (suite *maas2EnvironSuite) TestSupportedArchitecturesError(c *gc.C) {
	suite.injectController(&fakeController{bootResourcesError: errors.New("Something terrible!")})
	env := makeEnviron(c)
	_, err := env.SupportedArchitectures()
	c.Assert(err, gc.ErrorMatches, "Something terrible!")
}

func (suite *maas2EnvironSuite) injectController(controller gomaasapi.Controller) {
	mockGetController := func(maasServer, apiKey string) (gomaasapi.Controller, error) {
		return controller, nil
	}
	suite.PatchValue(&GetMAAS2Controller, mockGetController)
}

func (suite *maas2EnvironSuite) makeEnvironWithMachines(c *gc.C, expectedSystemIDs []string, returnSystemIDs []string) *maasEnviron {
	var env *maasEnviron
	checkArgs := func(args gomaasapi.MachinesArgs) {
		c.Check(args.SystemIDs, jc.DeepEquals, expectedSystemIDs)
		c.Check(args.AgentName, gc.Equals, env.ecfg().maasAgentName())
	}
	machines := make([]gomaasapi.Machine, len(returnSystemIDs))
	for index, id := range returnSystemIDs {
		machines[index] = &fakeMachine{systemID: id}
	}
	suite.injectController(&fakeController{
		machines:          machines,
		machinesArgsCheck: checkArgs,
	})
	env = makeEnviron(c)
	return env
}

func (suite *maas2EnvironSuite) TestAllInstances(c *gc.C) {
	env := suite.makeEnvironWithMachines(
		c, []string{}, []string{"tuco", "tio", "gus"},
	)
	result, err := env.AllInstances()
	c.Assert(err, jc.ErrorIsNil)
	expectedMachines := set.NewStrings("tuco", "tio", "gus")
	actualMachines := set.NewStrings()
	for _, instance := range result {
		actualMachines.Add(string(instance.Id()))
	}
	c.Assert(actualMachines, jc.DeepEquals, expectedMachines)
}

func (suite *maas2EnvironSuite) TestAllInstancesError(c *gc.C) {
	suite.injectController(&fakeController{machinesError: errors.New("Something terrible!")})
	env := makeEnviron(c)
	_, err := env.AllInstances()
	c.Assert(err, gc.ErrorMatches, "Something terrible!")
}

func (suite *maas2EnvironSuite) TestInstances(c *gc.C) {
	env := suite.makeEnvironWithMachines(
		c, []string{"jake", "bonnibel"}, []string{"jake", "bonnibel"},
	)
	result, err := env.Instances([]instance.Id{"jake", "bonnibel"})
	c.Assert(err, jc.ErrorIsNil)
	expectedMachines := set.NewStrings("jake", "bonnibel")
	actualMachines := set.NewStrings()
	for _, machine := range result {
		actualMachines.Add(string(machine.Id()))
	}
	c.Assert(actualMachines, jc.DeepEquals, expectedMachines)
}

func (suite *maas2EnvironSuite) TestInstancesPartialResult(c *gc.C) {
	env := suite.makeEnvironWithMachines(
		c, []string{"jake", "bonnibel"}, []string{"tuco", "bonnibel"},
	)
	result, err := env.Instances([]instance.Id{"jake", "bonnibel"})
	c.Check(err, gc.Equals, environs.ErrPartialInstances)
	c.Assert(result, gc.HasLen, 2)
	c.Assert(result[0], gc.IsNil)
	c.Assert(result[1].Id(), gc.Equals, instance.Id("bonnibel"))
}

func (suite *maas2EnvironSuite) TestAvailabilityZones(c *gc.C) {
	suite.injectController(&fakeController{
		zones: []gomaasapi.Zone{
			&fakeZone{name: "mossack"},
			&fakeZone{name: "fonseca"},
		},
	})
	env := makeEnviron(c)
	result, err := env.AvailabilityZones()
	c.Assert(err, jc.ErrorIsNil)
	expectedZones := set.NewStrings("mossack", "fonseca")
	actualZones := set.NewStrings()
	for _, zone := range result {
		actualZones.Add(zone.Name())
	}
	c.Assert(actualZones, jc.DeepEquals, expectedZones)
}

func (suite *maas2EnvironSuite) TestAvailabilityZonesError(c *gc.C) {
	suite.injectController(&fakeController{
		zonesError: errors.New("a bad thing"),
	})
	env := makeEnviron(c)
	_, err := env.AvailabilityZones()
	c.Assert(err, gc.ErrorMatches, "a bad thing")
}

func (suite *maas2EnvironSuite) TestSpaces(c *gc.C) {
	suite.injectController(&fakeController{
		spaces: []gomaasapi.Space{
			fakeSpace{
				name: "pepper",
				id:   1234,
			},
			fakeSpace{
				name: "freckles",
				id:   4567,
				subnets: []gomaasapi.Subnet{
					fakeSubnet{id: 99, vlanVid: 66, cidr: "192.168.10.0/24"},
					fakeSubnet{id: 98, vlanVid: 67, cidr: "192.168.11.0/24"},
				},
			},
		},
	})
	env := makeEnviron(c)
	result, err := env.Spaces()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, gc.HasLen, 1)
	c.Assert(result[0].Name, gc.Equals, "freckles")
	c.Assert(result[0].ProviderId, gc.Equals, network.Id("4567"))
	subnets := result[0].Subnets
	c.Assert(subnets, gc.HasLen, 2)
	c.Assert(subnets[0].ProviderId, gc.Equals, network.Id("99"))
	c.Assert(subnets[0].VLANTag, gc.Equals, 66)
	c.Assert(subnets[0].CIDR, gc.Equals, "192.168.10.0/24")
	c.Assert(subnets[0].SpaceProviderId, gc.Equals, network.Id("4567"))
	c.Assert(subnets[1].ProviderId, gc.Equals, network.Id("98"))
	c.Assert(subnets[1].VLANTag, gc.Equals, 67)
	c.Assert(subnets[1].CIDR, gc.Equals, "192.168.11.0/24")
	c.Assert(subnets[1].SpaceProviderId, gc.Equals, network.Id("4567"))
}

func (suite *maas2EnvironSuite) TestSpacesError(c *gc.C) {
	suite.injectController(&fakeController{
		spacesError: errors.New("Joe Manginiello"),
	})
	env := makeEnviron(c)
	_, err := env.Spaces()
	c.Assert(err, gc.ErrorMatches, "Joe Manginiello")
}

func (suite *maas2EnvironSuite) TestStartInstanceError(c *gc.C) {
	suite.injectController(&fakeController{
		allocateMachineError: errors.New("Charles Babbage"),
	})
	env := makeEnviron(c)
	_, err := env.StartInstance(environs.StartInstanceParams{})
	c.Assert(err, gc.ErrorMatches, ".* cannot run instance: Charles Babbage")
}
<commit_msg>Initial test for StartInstance<commit_after>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package maas

import (
	"net/http"

	"github.com/juju/errors"
	"github.com/juju/gomaasapi"
	jc "github.com/juju/testing/checkers"
	"github.com/juju/utils/set"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/feature"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/network"
	coretesting "github.com/juju/juju/testing"
)

type maas2EnvironSuite struct {
	baseProviderSuite
}

var _ = gc.Suite(&maas2EnvironSuite{})

func makeEnviron(c *gc.C) *maasEnviron {
	testAttrs := coretesting.Attrs{}
	for k, v := range maasEnvAttrs {
		testAttrs[k] = v
	}
	testAttrs["maas-server"] = "http://any-old-junk.invalid/"
	attrs := coretesting.FakeConfig().Merge(testAttrs)
	cfg, err := config.New(config.NoDefaults, attrs)
	c.Assert(err, jc.ErrorIsNil)
	env, err := NewEnviron(cfg)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(env, gc.NotNil)
	return env
}

func (suite *maas2EnvironSuite) SetUpTest(c *gc.C) {
	suite.baseProviderSuite.SetUpTest(c)
	suite.SetFeatureFlags(feature.MAAS2)
}

func (suite *maas2EnvironSuite) getEnvWithServer(c *gc.C) (*maasEnviron, error) {
	testServer := gomaasapi.NewSimpleServer()
	testServer.AddGetResponse("/api/2.0/version/", http.StatusOK, maas2VersionResponse)
	testServer.AddGetResponse("/api/2.0/users/?op=whoami", http.StatusOK, "{}")
	testServer.Start()
	suite.AddCleanup(func(*gc.C) { testServer.Close() })
	testAttrs := coretesting.Attrs{}
	for k, v := range maasEnvAttrs {
		testAttrs[k] = v
	}
	testAttrs["maas-server"] = testServer.Server.URL
	attrs := coretesting.FakeConfig().Merge(testAttrs)
	cfg, err := config.New(config.NoDefaults, attrs)
	c.Assert(err, jc.ErrorIsNil)
	return NewEnviron(cfg)
}

func (suite *maas2EnvironSuite) TestNewEnvironWithoutFeatureFlag(c *gc.C) {
	suite.SetFeatureFlags()
	_, err := suite.getEnvWithServer(c)
	c.Assert(err, jc.Satisfies, errors.IsNotSupported)
}

func (suite *maas2EnvironSuite) TestNewEnvironWithController(c *gc.C) {
	env, err := suite.getEnvWithServer(c)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(env, gc.NotNil)
}

func (suite *maas2EnvironSuite) TestSupportedArchitectures(c *gc.C) {
	controller := fakeController{
		bootResources: []gomaasapi.BootResource{
			&fakeBootResource{name: "wily", architecture: "amd64/blah"},
			&fakeBootResource{name: "wily", architecture: "amd64/something"},
			&fakeBootResource{name: "xenial", architecture: "arm/somethingelse"},
		},
	}
	suite.injectController(&controller)
	env := makeEnviron(c)
	result, err := env.SupportedArchitectures()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, jc.DeepEquals, []string{"amd64", "arm"})
}

func (suite *maas2EnvironSuite) TestSupportedArchitecturesError(c *gc.C) {
	suite.injectController(&fakeController{bootResourcesError: errors.New("Something terrible!")})
	env := makeEnviron(c)
	_, err := env.SupportedArchitectures()
	c.Assert(err, gc.ErrorMatches, "Something terrible!")
}

func (suite *maas2EnvironSuite) injectController(controller gomaasapi.Controller) {
	mockGetController := func(maasServer, apiKey string) (gomaasapi.Controller, error) {
		return controller, nil
	}
	suite.PatchValue(&GetMAAS2Controller, mockGetController)
}

func (suite *maas2EnvironSuite) makeEnvironWithMachines(c *gc.C, expectedSystemIDs []string, returnSystemIDs []string) *maasEnviron {
	var env *maasEnviron
	checkArgs := func(args gomaasapi.MachinesArgs) {
		c.Check(args.SystemIDs, jc.DeepEquals, expectedSystemIDs)
		c.Check(args.AgentName, gc.Equals, env.ecfg().maasAgentName())
	}
	machines := make([]gomaasapi.Machine, len(returnSystemIDs))
	for index, id := range returnSystemIDs {
		machines[index] = &fakeMachine{systemID: id}
	}
	suite.injectController(&fakeController{
		machines:          machines,
		machinesArgsCheck: checkArgs,
	})
	env = makeEnviron(c)
	return env
}

func (suite *maas2EnvironSuite) TestAllInstances(c *gc.C) {
	env := suite.makeEnvironWithMachines(
		c, []string{}, []string{"tuco", "tio", "gus"},
	)
	result, err := env.AllInstances()
	c.Assert(err, jc.ErrorIsNil)
	expectedMachines := set.NewStrings("tuco", "tio", "gus")
	actualMachines := set.NewStrings()
	for _, instance := range result {
		actualMachines.Add(string(instance.Id()))
	}
	c.Assert(actualMachines, jc.DeepEquals, expectedMachines)
}

func (suite *maas2EnvironSuite) TestAllInstancesError(c *gc.C) {
	suite.injectController(&fakeController{machinesError: errors.New("Something terrible!")})
	env := makeEnviron(c)
	_, err := env.AllInstances()
	c.Assert(err, gc.ErrorMatches, "Something terrible!")
}

func (suite *maas2EnvironSuite) TestInstances(c *gc.C) {
	env := suite.makeEnvironWithMachines(
		c, []string{"jake", "bonnibel"}, []string{"jake", "bonnibel"},
	)
	result, err := env.Instances([]instance.Id{"jake", "bonnibel"})
	c.Assert(err, jc.ErrorIsNil)
	expectedMachines := set.NewStrings("jake", "bonnibel")
	actualMachines := set.NewStrings()
	for _, machine := range result {
		actualMachines.Add(string(machine.Id()))
	}
	c.Assert(actualMachines, jc.DeepEquals, expectedMachines)
}

func (suite *maas2EnvironSuite) TestInstancesPartialResult(c *gc.C) {
	env := suite.makeEnvironWithMachines(
		c, []string{"jake", "bonnibel"}, []string{"tuco", "bonnibel"},
	)
	result, err := env.Instances([]instance.Id{"jake", "bonnibel"})
	c.Check(err, gc.Equals, environs.ErrPartialInstances)
	c.Assert(result, gc.HasLen, 2)
	c.Assert(result[0], gc.IsNil)
	c.Assert(result[1].Id(), gc.Equals, instance.Id("bonnibel"))
}

func (suite *maas2EnvironSuite) TestAvailabilityZones(c *gc.C) {
	suite.injectController(&fakeController{
		zones: []gomaasapi.Zone{
			&fakeZone{name: "mossack"},
			&fakeZone{name: "fonseca"},
		},
	})
	env := makeEnviron(c)
	result, err := env.AvailabilityZones()
	c.Assert(err, jc.ErrorIsNil)
	expectedZones := set.NewStrings("mossack", "fonseca")
	actualZones := set.NewStrings()
	for _, zone := range result {
		actualZones.Add(zone.Name())
	}
	c.Assert(actualZones, jc.DeepEquals, expectedZones)
}

func (suite *maas2EnvironSuite) TestAvailabilityZonesError(c *gc.C) {
	suite.injectController(&fakeController{
		zonesError: errors.New("a bad thing"),
	})
	env := makeEnviron(c)
	_, err := env.AvailabilityZones()
	c.Assert(err, gc.ErrorMatches, "a bad thing")
}

func (suite *maas2EnvironSuite) TestSpaces(c *gc.C) {
	suite.injectController(&fakeController{
		spaces: []gomaasapi.Space{
			fakeSpace{
				name: "pepper",
				id:   1234,
			},
			fakeSpace{
				name: "freckles",
				id:   4567,
				subnets: []gomaasapi.Subnet{
					fakeSubnet{id: 99, vlanVid: 66, cidr: "192.168.10.0/24"},
					fakeSubnet{id: 98, vlanVid: 67, cidr: "192.168.11.0/24"},
				},
			},
		},
	})
	env := makeEnviron(c)
	result, err := env.Spaces()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, gc.HasLen, 1)
	c.Assert(result[0].Name, gc.Equals, "freckles")
	c.Assert(result[0].ProviderId, gc.Equals, network.Id("4567"))
	subnets := result[0].Subnets
	c.Assert(subnets, gc.HasLen, 2)
	c.Assert(subnets[0].ProviderId, gc.Equals, network.Id("99"))
	c.Assert(subnets[0].VLANTag, gc.Equals, 66)
	c.Assert(subnets[0].CIDR, gc.Equals, "192.168.10.0/24")
	c.Assert(subnets[0].SpaceProviderId, gc.Equals, network.Id("4567"))
	c.Assert(subnets[1].ProviderId, gc.Equals, network.Id("98"))
	c.Assert(subnets[1].VLANTag, gc.Equals, 67)
	c.Assert(subnets[1].CIDR, gc.Equals, "192.168.11.0/24")
	c.Assert(subnets[1].SpaceProviderId, gc.Equals, network.Id("4567"))
}

func (suite *maas2EnvironSuite) TestSpacesError(c *gc.C) {
	suite.injectController(&fakeController{
		spacesError: errors.New("Joe Manginiello"),
	})
	env := makeEnviron(c)
	_, err := env.Spaces()
	c.Assert(err, gc.ErrorMatches, "Joe Manginiello")
}

func (suite *maas2EnvironSuite) TestStartInstanceError(c *gc.C) {
	suite.injectController(&fakeController{
		allocateMachineError: errors.New("Charles Babbage"),
	})
	env := makeEnviron(c)
	_, err := env.StartInstance(environs.StartInstanceParams{})
	c.Assert(err, gc.ErrorMatches, ".* cannot run instance: Charles Babbage")
}

func (suite *maas2EnvironSuite) TestStartInstance(c *gc.C) {
	var env *maasEnviron
	suite.injectController(&fakeController{
		allocateMachineArgsCheck: func(args gomaasapi.AllocateMachineArgs) {
			c.Assert(args, jc.DeepEquals, gomaasapi.AllocateMachineArgs{
				AgentName: env.ecfg().maasAgentName()})
		},
		allocateMachine: &fakeMachine{systemID: "Bruce sterling"},
	})
	env = makeEnviron(c)
	result, err := env.StartInstance(environs.StartInstanceParams{})
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result.Instance.Id(), gc.Equals, instance.Id("Bruce Sterling"))
}
<|endoftext|>


<commit_before>package main

import "fmt"

// fibonacci is a function that returns
// a function that returns an int.
func fibonacci() func() int {
}

func main() {
    f := fibonacci()
    for i := 0; i < 10; i++ {
        fmt.Println(f())
    }
}
<commit_msg>Implement `fibonacci`.<commit_after>package main

import "fmt"

// fibonacci is a function that returns
// a function that returns an int.
func fibonacci() func() int {
	n, x, y := 0, 0, 1
	return func() int {
		if n == 0 {
			n++
		} else {
			x, y = y, x+y
		}
		return x
	}
}

func main() {
	f := fibonacci()
	for i := 0; i < 10; i++ {
		fmt.Println(f())
	}
}
<|endoftext|>


<commit_before>/************************
[钩子引擎 (version 0.3)]
@author:S.W.H
@E-mail:swh@admpub.com
@update:2014-01-18
************************/
package xweb

import (
	"errors"
	"fmt"
	"reflect"
)

var (
	ErrParamsNotAdapted             = errors.New("The number of params is not adapted.")
	XHook               *HookEngine = NewHookEngine(10)
)

type Hook []reflect.Value
type HookEngine struct {
	Hooks map[string]Hook
	Index map[string]uint
}

func (f *HookEngine) Bind(name string, fns ...interface{}) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = errors.New(name + " is not callable.")
		}
	}()
	if _, ok := f.Hooks[name]; !ok {
		f.Hooks[name] = make(Hook, 0)
	}
	if _, ok := f.Index[name]; !ok {
		f.Index[name] = 0
	}
	hln := uint(len(f.Hooks[name]))
	fln := f.Index[name] + 1 + uint(len(fns))
	if hln < fln {
		for _, fn := range fns {
			v := reflect.ValueOf(fn)
			f.Hooks[name] = append(f.Hooks[name], v)
			f.Index[name]++
		}
	} else {
		for _, fn := range fns {
			v := reflect.ValueOf(fn)
			f.Hooks[name][f.Index[name]] = v
			f.Index[name]++
		}
	}
	return
}

func (f *HookEngine) Call(name string, params ...interface{}) (result []reflect.Value, err error) {
	if _, ok := f.Hooks[name]; !ok {
		err = errors.New(name + " does not exist.")
		return
	}
	ln := len(params)
	in := make([]reflect.Value, ln)
	for k, param := range params {
		in[k] = reflect.ValueOf(param)
	}
	for _, v := range f.Hooks[name] {
		if v.IsValid() == false {
			continue
		}
		if ln != v.Type().NumIn() {
			continue
			err = ErrParamsNotAdapted
			return
		}
		result = v.Call(in)
		for _k, _v := range result {
			in[_k] = _v
		}
	}
	if len(result) == 0 {
		err = errors.New(name + " have nothing to do.")
	}
	return
}

func (f *HookEngine) Value(c []reflect.Value, index int) (r interface{}) {
	if len(c) >= index && c[index].CanInterface() {
		r = c[index].Interface()
	}
	return
}

func (f *HookEngine) String(c reflect.Value) string {
	return fmt.Sprintf("%s", c)
}

func NewHookEngine(size int) *HookEngine {
	h := &HookEngine{Hooks: make(map[string]Hook, size), Index: make(map[string]uint, size)}

	//func(mux *http.ServeMux) *http.ServeMux
	h.Hooks["MuxHandle"] = make(Hook, 0)

	//func(result *bool, serv *Server, w http.ResponseWriter, req *http.Request) *bool
	h.Hooks["BeforeProcess"] = make(Hook, 0)

	//func(result *bool, serv *Server, w http.ResponseWriter, req *http.Request) *bool
	h.Hooks["AfterProcess"] = make(Hook, 0)

	//func(content string, action *Action) string
	h.Hooks["BeforeRender"] = make(Hook, 0)

	//func(content []byte, action *Action) []byte
	h.Hooks["AfterRender"] = make(Hook, 0)
	return h
}
<commit_msg>update<commit_after>/************************
[钩子引擎 (version 0.3)]
@author:S.W.H
@E-mail:swh@admpub.com
@update:2014-01-18
************************/
package xweb

import (
	"errors"
	"fmt"
	"reflect"
	"sync"
)

var (
	ErrParamsNotAdapted             = errors.New("The number of params is not adapted.")
	XHook               *HookEngine = NewHookEngine(10)
)

type Hook []reflect.Value
type HookEngine struct {
	Hooks map[string]Hook
	Index map[string]uint
	lock  *sync.RWMutex
}

func (f *HookEngine) Bind(name string, fns ...interface{}) (err error) {
	f.lock.Lock()
	defer func() {
		f.lock.Unlock()
		if e := recover(); e != nil {
			err = errors.New(name + " is not callable.")
		}
	}()
	if _, ok := f.Hooks[name]; !ok {
		f.Hooks[name] = make(Hook, 0)
	}
	if _, ok := f.Index[name]; !ok {
		f.Index[name] = 0
	}
	hln := uint(len(f.Hooks[name]))
	fln := f.Index[name] + 1 + uint(len(fns))
	if hln < fln {
		for _, fn := range fns {
			v := reflect.ValueOf(fn)
			f.Hooks[name] = append(f.Hooks[name], v)
			f.Index[name]++
		}
	} else {
		for _, fn := range fns {
			v := reflect.ValueOf(fn)
			f.Hooks[name][f.Index[name]] = v
			f.Index[name]++
		}
	}
	return
}

func (f *HookEngine) Call(name string, params ...interface{}) (result []reflect.Value, err error) {
	f.lock.Lock()
	defer f.lock.Unlock()
	if _, ok := f.Hooks[name]; !ok {
		err = errors.New(name + " does not exist.")
		return
	}
	ln := len(params)
	in := make([]reflect.Value, ln)
	for k, param := range params {
		in[k] = reflect.ValueOf(param)
	}
	for _, v := range f.Hooks[name] {
		if v.IsValid() == false {
			continue
		}
		if ln != v.Type().NumIn() {
			continue
			err = ErrParamsNotAdapted
			return
		}
		result = v.Call(in)
		for _k, _v := range result {
			in[_k] = _v
		}
	}
	if len(result) == 0 {
		err = errors.New(name + " have nothing to do.")
	}
	return
}

func (f *HookEngine) Value(c []reflect.Value, index int) (r interface{}) {
	if len(c) >= index && c[index].CanInterface() {
		r = c[index].Interface()
	}
	return
}

func (f *HookEngine) String(c reflect.Value) string {
	return fmt.Sprintf("%s", c)
}

func NewHookEngine(size int) *HookEngine {
	h := &HookEngine{Hooks: make(map[string]Hook, size), Index: make(map[string]uint, size), lock: new(sync.RWMutex)}

	//func(mux *http.ServeMux) *http.ServeMux
	h.Hooks["MuxHandle"] = make(Hook, 0)

	//func(result *bool, serv *Server, w http.ResponseWriter, req *http.Request) *bool
	h.Hooks["BeforeProcess"] = make(Hook, 0)

	//func(result *bool, serv *Server, w http.ResponseWriter, req *http.Request) *bool
	h.Hooks["AfterProcess"] = make(Hook, 0)

	//func(content string, action *Action) string
	h.Hooks["BeforeRender"] = make(Hook, 0)

	//func(content []byte, action *Action) []byte
	h.Hooks["AfterRender"] = make(Hook, 0)
	return h
}
<|endoftext|>


<commit_before>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// See https://code.google.com/p/go/source/browse/CONTRIBUTORS
// Licensed under the same terms as Go itself:
// https://code.google.com/p/go/source/browse/LICENSE

// Package http2 implements the HTTP/2 protocol.
//
// This is a work in progress. This package is low-level and intended
// to be used directly by very few people. Most users will use it
// indirectly through integration with the net/http package. See
// ConfigureServer. That ConfigureServer call will likely be automatic
// or available via an empty import in the future.
//
// This package currently targets draft-14. See http://http2.github.io/
package http2

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"sync"
)

var VerboseLogs = false

const (
	// ClientPreface is the string that must be sent by new
	// connections from clients.
	ClientPreface = "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"

	// SETTINGS_MAX_FRAME_SIZE default
	// http://http2.github.io/http2-spec/#rfc.section.6.5.2
	initialMaxFrameSize = 16384

	// NextProtoTLS is the NPN/ALPN protocol negotiated during
	// HTTP/2's TLS setup.
	NextProtoTLS = "h2"

	// http://http2.github.io/http2-spec/#SettingValues
	initialHeaderTableSize = 4096

	initialWindowSize = 65535 // 6.9.2 Initial Flow Control Window Size

	defaultMaxReadFrameSize = 1 << 20
)

var (
	clientPreface = []byte(ClientPreface)
)

type streamState int

const (
	stateIdle streamState = iota
	stateOpen
	stateHalfClosedLocal
	stateHalfClosedRemote
	stateResvLocal
	stateResvRemote
	stateClosed
)

var stateName = [...]string{
	stateIdle:             "Idle",
	stateOpen:             "Open",
	stateHalfClosedLocal:  "HalfClosedLocal",
	stateHalfClosedRemote: "HalfClosedRemote",
	stateResvLocal:        "ResvLocal",
	stateResvRemote:       "ResvRemote",
	stateClosed:           "Closed",
}

func (st streamState) String() string {
	return stateName[st]
}

// Setting is a setting parameter: which setting it is, and its value.
type Setting struct {
	// ID is which setting is being set.
	// See http://http2.github.io/http2-spec/#SettingValues
	ID SettingID

	// Val is the value.
	Val uint32
}

func (s Setting) String() string {
	return fmt.Sprintf("[%v = %d]", s.ID, s.Val)
}

// Valid reports whether the setting is valid.
func (s Setting) Valid() error {
	// Limits and error codes from 6.5.2 Defined SETTINGS Parameters
	switch s.ID {
	case SettingEnablePush:
		if s.Val != 1 && s.Val != 0 {
			return ConnectionError(ErrCodeProtocol)
		}
	case SettingInitialWindowSize:
		if s.Val > 1<<31-1 {
			return ConnectionError(ErrCodeFlowControl)
		}
	case SettingMaxFrameSize:
		if s.Val < 16384 || s.Val > 1<<24-1 {
			return ConnectionError(ErrCodeProtocol)
		}
	}
	return nil
}

// A SettingID is an HTTP/2 setting as defined in
// http://http2.github.io/http2-spec/#iana-settings
type SettingID uint16

const (
	SettingHeaderTableSize      SettingID = 0x1
	SettingEnablePush           SettingID = 0x2
	SettingMaxConcurrentStreams SettingID = 0x3
	SettingInitialWindowSize    SettingID = 0x4
	SettingMaxFrameSize         SettingID = 0x5
	SettingMaxHeaderListSize    SettingID = 0x6
)

var settingName = map[SettingID]string{
	SettingHeaderTableSize:      "HEADER_TABLE_SIZE",
	SettingEnablePush:           "ENABLE_PUSH",
	SettingMaxConcurrentStreams: "MAX_CONCURRENT_STREAMS",
	SettingInitialWindowSize:    "INITIAL_WINDOW_SIZE",
	SettingMaxFrameSize:         "MAX_FRAME_SIZE",
	SettingMaxHeaderListSize:    "MAX_HEADER_LIST_SIZE",
}

func (s SettingID) String() string {
	if v, ok := settingName[s]; ok {
		return v
	}
	return fmt.Sprintf("UNKNOWN_SETTING_%d", uint16(s))
}

func validHeader(v string) bool {
	if len(v) == 0 {
		return false
	}
	for _, r := range v {
		// "Just as in HTTP/1.x, header field names are
		// strings of ASCII characters that are compared in a
		// case-insensitive fashion. However, header field
		// names MUST be converted to lowercase prior to their
		// encoding in HTTP/2. "
		if r >= 127 || ('A' <= r && r <= 'Z') {
			return false
		}
	}
	return true
}

var httpCodeStringCommon = map[int]string{} // n -> strconv.Itoa(n)

func init() {
	for i := 100; i <= 999; i++ {
		if v := http.StatusText(i); v != "" {
			httpCodeStringCommon[i] = strconv.Itoa(i)
		}
	}
}

func httpCodeString(code int) string {
	if s, ok := httpCodeStringCommon[code]; ok {
		return s
	}
	return strconv.Itoa(code)
}

// from pkg io
type stringWriter interface {
	WriteString(s string) (n int, err error)
}

// A gate lets two goroutines coordinate their activities.
type gate chan struct{}

func (g gate) Done() { g <- struct{}{} }
func (g gate) Wait() { <-g }

// A closeWaiter is like a sync.WaitGroup but only goes 1 to 0 (open to closed).
type closeWaiter chan struct{}

// Init makes a closeWaiter usable.
// It exists because so a closeWaiter value can be placed inside a
// larger struct and have the Mutex and Cond's memory in the same
// allocation.
func (cw *closeWaiter) Init() {
	*cw = make(chan struct{})
}

// Close marks the closeWaiter as closed and unblocks any waiters.
func (cw closeWaiter) Close() {
	close(cw)
}

// Wait waits for the closeWaiter to become closed.
func (cw closeWaiter) Wait() {
	<-cw
}

// bufferedWriter is a buffered writer that writes to w.
// Its buffered writer is lazily allocated as needed, to minimize
// idle memory usage with many connections.
type bufferedWriter struct {
	w  io.Writer     // immutable
	bw *bufio.Writer // non-nil when data is buffered
}

func newBufferedWriter(w io.Writer) *bufferedWriter {
	return &bufferedWriter{w: w}
}

var bufWriterPool = sync.Pool{
	New: func() interface{} {
		// TODO: pick something better? this is a bit under
		// (3 x typical 1500 byte MTU) at least.
		return bufio.NewWriterSize(nil, 4<<10)
	},
}

func (w *bufferedWriter) Write(p []byte) (n int, err error) {
	if w.bw == nil {
		bw := bufWriterPool.Get().(*bufio.Writer)
		bw.Reset(w.w)
		w.bw = bw
	}
	return w.bw.Write(p)
}

func (w *bufferedWriter) Flush() error {
	bw := w.bw
	if bw == nil {
		return nil
	}
	err := bw.Flush()
	bw.Reset(nil)
	bufWriterPool.Put(bw)
	w.bw = nil
	return err
}
<commit_msg>Remove draft-14 comment in package docs.<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// See https://code.google.com/p/go/source/browse/CONTRIBUTORS
// Licensed under the same terms as Go itself:
// https://code.google.com/p/go/source/browse/LICENSE

// Package http2 implements the HTTP/2 protocol.
//
// This is a work in progress. This package is low-level and intended
// to be used directly by very few people. Most users will use it
// indirectly through integration with the net/http package. See
// ConfigureServer. That ConfigureServer call will likely be automatic
// or available via an empty import in the future.
//
// See http://http2.github.io/
package http2

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"sync"
)

var VerboseLogs = false

const (
	// ClientPreface is the string that must be sent by new
	// connections from clients.
	ClientPreface = "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"

	// SETTINGS_MAX_FRAME_SIZE default
	// http://http2.github.io/http2-spec/#rfc.section.6.5.2
	initialMaxFrameSize = 16384

	// NextProtoTLS is the NPN/ALPN protocol negotiated during
	// HTTP/2's TLS setup.
	NextProtoTLS = "h2"

	// http://http2.github.io/http2-spec/#SettingValues
	initialHeaderTableSize = 4096

	initialWindowSize = 65535 // 6.9.2 Initial Flow Control Window Size

	defaultMaxReadFrameSize = 1 << 20
)

var (
	clientPreface = []byte(ClientPreface)
)

type streamState int

const (
	stateIdle streamState = iota
	stateOpen
	stateHalfClosedLocal
	stateHalfClosedRemote
	stateResvLocal
	stateResvRemote
	stateClosed
)

var stateName = [...]string{
	stateIdle:             "Idle",
	stateOpen:             "Open",
	stateHalfClosedLocal:  "HalfClosedLocal",
	stateHalfClosedRemote: "HalfClosedRemote",
	stateResvLocal:        "ResvLocal",
	stateResvRemote:       "ResvRemote",
	stateClosed:           "Closed",
}

func (st streamState) String() string {
	return stateName[st]
}

// Setting is a setting parameter: which setting it is, and its value.
type Setting struct {
	// ID is which setting is being set.
	// See http://http2.github.io/http2-spec/#SettingValues
	ID SettingID

	// Val is the value.
	Val uint32
}

func (s Setting) String() string {
	return fmt.Sprintf("[%v = %d]", s.ID, s.Val)
}

// Valid reports whether the setting is valid.
func (s Setting) Valid() error {
	// Limits and error codes from 6.5.2 Defined SETTINGS Parameters
	switch s.ID {
	case SettingEnablePush:
		if s.Val != 1 && s.Val != 0 {
			return ConnectionError(ErrCodeProtocol)
		}
	case SettingInitialWindowSize:
		if s.Val > 1<<31-1 {
			return ConnectionError(ErrCodeFlowControl)
		}
	case SettingMaxFrameSize:
		if s.Val < 16384 || s.Val > 1<<24-1 {
			return ConnectionError(ErrCodeProtocol)
		}
	}
	return nil
}

// A SettingID is an HTTP/2 setting as defined in
// http://http2.github.io/http2-spec/#iana-settings
type SettingID uint16

const (
	SettingHeaderTableSize      SettingID = 0x1
	SettingEnablePush           SettingID = 0x2
	SettingMaxConcurrentStreams SettingID = 0x3
	SettingInitialWindowSize    SettingID = 0x4
	SettingMaxFrameSize         SettingID = 0x5
	SettingMaxHeaderListSize    SettingID = 0x6
)

var settingName = map[SettingID]string{
	SettingHeaderTableSize:      "HEADER_TABLE_SIZE",
	SettingEnablePush:           "ENABLE_PUSH",
	SettingMaxConcurrentStreams: "MAX_CONCURRENT_STREAMS",
	SettingInitialWindowSize:    "INITIAL_WINDOW_SIZE",
	SettingMaxFrameSize:         "MAX_FRAME_SIZE",
	SettingMaxHeaderListSize:    "MAX_HEADER_LIST_SIZE",
}

func (s SettingID) String() string {
	if v, ok := settingName[s]; ok {
		return v
	}
	return fmt.Sprintf("UNKNOWN_SETTING_%d", uint16(s))
}

func validHeader(v string) bool {
	if len(v) == 0 {
		return false
	}
	for _, r := range v {
		// "Just as in HTTP/1.x, header field names are
		// strings of ASCII characters that are compared in a
		// case-insensitive fashion. However, header field
		// names MUST be converted to lowercase prior to their
		// encoding in HTTP/2. "
		if r >= 127 || ('A' <= r && r <= 'Z') {
			return false
		}
	}
	return true
}

var httpCodeStringCommon = map[int]string{} // n -> strconv.Itoa(n)

func init() {
	for i := 100; i <= 999; i++ {
		if v := http.StatusText(i); v != "" {
			httpCodeStringCommon[i] = strconv.Itoa(i)
		}
	}
}

func httpCodeString(code int) string {
	if s, ok := httpCodeStringCommon[code]; ok {
		return s
	}
	return strconv.Itoa(code)
}

// from pkg io
type stringWriter interface {
	WriteString(s string) (n int, err error)
}

// A gate lets two goroutines coordinate their activities.
type gate chan struct{}

func (g gate) Done() { g <- struct{}{} }
func (g gate) Wait() { <-g }

// A closeWaiter is like a sync.WaitGroup but only goes 1 to 0 (open to closed).
type closeWaiter chan struct{}

// Init makes a closeWaiter usable.
// It exists because so a closeWaiter value can be placed inside a
// larger struct and have the Mutex and Cond's memory in the same
// allocation.
func (cw *closeWaiter) Init() {
	*cw = make(chan struct{})
}

// Close marks the closeWaiter as closed and unblocks any waiters.
func (cw closeWaiter) Close() {
	close(cw)
}

// Wait waits for the closeWaiter to become closed.
func (cw closeWaiter) Wait() {
	<-cw
}

// bufferedWriter is a buffered writer that writes to w.
// Its buffered writer is lazily allocated as needed, to minimize
// idle memory usage with many connections.
type bufferedWriter struct {
	w  io.Writer     // immutable
	bw *bufio.Writer // non-nil when data is buffered
}

func newBufferedWriter(w io.Writer) *bufferedWriter {
	return &bufferedWriter{w: w}
}

var bufWriterPool = sync.Pool{
	New: func() interface{} {
		// TODO: pick something better? this is a bit under
		// (3 x typical 1500 byte MTU) at least.
		return bufio.NewWriterSize(nil, 4<<10)
	},
}

func (w *bufferedWriter) Write(p []byte) (n int, err error) {
	if w.bw == nil {
		bw := bufWriterPool.Get().(*bufio.Writer)
		bw.Reset(w.w)
		w.bw = bw
	}
	return w.bw.Write(p)
}

func (w *bufferedWriter) Flush() error {
	bw := w.bw
	if bw == nil {
		return nil
	}
	err := bw.Flush()
	bw.Reset(nil)
	bufWriterPool.Put(bw)
	w.bw = nil
	return err
}
<|endoftext|>


<commit_before>package goproxy

import (
	"bufio"
	"crypto/tls"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
)

type ConnectActionLiteral int

const (
	ConnectAccept = iota
	ConnectReject
	ConnectMitm
	ConnectHijack
	ConnectHTTPMitm
)

var (
	OkConnect       = &ConnectAction{Action: ConnectAccept}
	MitmConnect     = &ConnectAction{Action: ConnectMitm}
	HTTPMitmConnect = &ConnectAction{Action: ConnectHTTPMitm}
	RejectConnect   = &ConnectAction{Action: ConnectReject}
)

type ConnectAction struct {
	Action    ConnectActionLiteral
	TlsConfig *tls.Config
	Ca        *tls.Certificate
	Hijack    func(req *http.Request, client net.Conn, ctx *ProxyCtx)
}

func stripPort(s string) string {
	ix := strings.IndexRune(s, ':')
	if ix == -1 {
		return s
	}
	return s[:ix]
}

func (proxy *ProxyHttpServer) dial(network, addr string) (c net.Conn, err error) {
	if proxy.Tr.Dial != nil {
		return proxy.Tr.Dial(network, addr)
	}
	return net.Dial(network, addr)
}

func (proxy *ProxyHttpServer) connectDial(network, addr string) (c net.Conn, err error) {
	if proxy.ConnectDial == nil {
		return proxy.dial(network, addr)
	}
	return proxy.ConnectDial(network, addr)
}

func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request) {
	ctx := &ProxyCtx{Req: r, Session: atomic.AddInt64(&proxy.sess, 1), proxy: proxy}

	hij, ok := w.(http.Hijacker)
	if !ok {
		panic("httpserver does not support hijacking")
	}

	proxyClient, _, e := hij.Hijack()
	if e != nil {
		panic("Cannot hijack connection " + e.Error())
	}

	ctx.Logf("Running %d CONNECT handlers", len(proxy.httpsHandlers))
	todo, host := OkConnect, r.URL.Host
	ctx.Req = r
	for _, h := range proxy.httpsHandlers {
		newtodo, newhost := h.HandleConnect(host, ctx)
		if newtodo != nil {
			todo, host = newtodo, newhost
		}
		ctx.Logf("handler: %v %s", todo, host)
	}
	switch todo.Action {
	case ConnectAccept:
		if !hasPort.MatchString(host) {
			host += ":80"
		}
		targetSiteCon, err := proxy.connectDial("tcp", host)
		if err != nil {
			httpError(proxyClient, ctx, err)
			return
		}
		ctx.Logf("Accepting CONNECT to %s", host)
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		go copyAndClose(ctx, targetSiteCon, proxyClient)
		go copyAndClose(ctx, proxyClient, targetSiteCon)
	case ConnectHijack:
		ctx.Logf("Hijacking CONNECT to %s", host)
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		todo.Hijack(r, proxyClient, ctx)
	case ConnectHTTPMitm:
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		ctx.Logf("Assuming CONNECT is plain HTTP tunneling, mitm proxying it")
		targetSiteCon, err := proxy.connectDial("tcp", host)
		if err != nil {
			ctx.Warnf("Error dialing to %s: %s", host, err.Error())
			return
		}
		for {
			client := bufio.NewReader(proxyClient)
			remote := bufio.NewReader(targetSiteCon)
			req, err := http.ReadRequest(client)
			if err != nil && err != io.EOF {
				ctx.Warnf("cannot read request of MITM HTTP client: %+#v", err)
			}
			if err != nil {
				return
			}
			req, resp := proxy.filterRequest(req, ctx)
			if resp == nil {
				if err := req.Write(targetSiteCon); err != nil {
					httpError(proxyClient, ctx, err)
					return
				}
				resp, err = http.ReadResponse(remote, req)
				if err != nil {
					httpError(proxyClient, ctx, err)
					return
				}
			}
			resp = proxy.filterResponse(resp, ctx)
			if err := resp.Write(proxyClient); err != nil {
				httpError(proxyClient, ctx, err)
				return
			}
		}
	case ConnectMitm:
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		ctx.Logf("Assuming CONNECT is TLS, mitm proxying it")
		// this goes in a separate goroutine, so that the net/http server won't think we're
		// still handling the request even after hijacking the connection. Those HTTP CONNECT
		// request can take forever, and the server will be stuck when "closed".
		// TODO: Allow Server.Close() mechanism to shut down this connection as nicely as possible
		ca := todo.Ca
		if ca == nil {
			ca = &GoproxyCa
		}
		cert, err := signHost(*ca, []string{stripPort(host)})
		if err != nil {
			ctx.Warnf("Cannot sign host certificate with provided CA: %s", err)
			return
		}
		tlsConfig := tls.Config{}
		if todo.TlsConfig != nil {
			tlsConfig = *todo.TlsConfig
		} else {
			tlsConfig = *defaultTlsConfig
		}
		tlsConfig.Certificates = append(tlsConfig.Certificates, cert)
		go func() {
			//TODO: cache connections to the remote website
			rawClientTls := tls.Server(proxyClient, &tlsConfig)
			if err := rawClientTls.Handshake(); err != nil {
				ctx.Warnf("Cannot handshake client %v %v", r.Host, err)
				return
			}
			defer rawClientTls.Close()
			clientTlsReader := bufio.NewReader(rawClientTls)
			for !isEof(clientTlsReader) {
				req, err := http.ReadRequest(clientTlsReader)
				if err != nil && err != io.EOF {
					return
				}
				if err != nil {
					ctx.Warnf("Cannot read TLS request from mitm'd client %v %v", r.Host, err)
					return
				}
				ctx.Logf("req %v", r.Host)
				req.URL, err = url.Parse("https://" + r.Host + req.URL.String())
				req, resp := proxy.filterRequest(req, ctx)
				if resp == nil {
					if err != nil {
						ctx.Warnf("Illegal URL %s", "https://"+r.Host+req.URL.Path)
						return
					}
					removeProxyHeaders(ctx, req)
					resp, err = proxy.Tr.RoundTrip(req)
					if err != nil {
						ctx.Warnf("Cannot read TLS response from mitm'd server %v", err)
						return
					}
					ctx.Logf("resp %v", resp.Status)
				}
				resp = proxy.filterResponse(resp, ctx)
				text := resp.Status
				statusCode := strconv.Itoa(resp.StatusCode) + " "
				if strings.HasPrefix(text, statusCode) {
					text = text[len(statusCode):]
				}
				// always use 1.1 to support chunked encoding
				if _, err := io.WriteString(rawClientTls, "HTTP/1.1"+" "+statusCode+text+"\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response HTTP status from mitm'd client: %v", err)
					return
				}
				// Since we don't know the length of resp, return chunked encoded response
				// TODO: use a more reasonable scheme
				resp.Header.Del("Content-Length")
				resp.Header.Set("Transfer-Encoding", "chunked")
				if err := resp.Header.Write(rawClientTls); err != nil {
					ctx.Warnf("Cannot write TLS response header from mitm'd client: %v", err)
					return
				}
				if _, err = io.WriteString(rawClientTls, "\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response header end from mitm'd client: %v", err)
					return
				}
				chunked := newChunkedWriter(rawClientTls)
				if _, err := io.Copy(chunked, resp.Body); err != nil {
					ctx.Warnf("Cannot write TLS response body from mitm'd client: %v", err)
					return
				}
				if err := chunked.Close(); err != nil {
					ctx.Warnf("Cannot write TLS chunked EOF from mitm'd client: %v", err)
					return
				}
				if _, err = io.WriteString(rawClientTls, "\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response chunked trailer from mitm'd client: %v", err)
					return
				}
			}
			ctx.Logf("Exiting on EOF")
		}()
	case ConnectReject:
		if ctx.Resp != nil {
			if err := ctx.Resp.Write(proxyClient); err != nil {
				ctx.Warnf("Cannot write response that reject http CONNECT: %v", err)
			}
		}
		proxyClient.Close()
	}
}

func httpError(w io.WriteCloser, ctx *ProxyCtx, err error) {
	if _, err := io.WriteString(w, "HTTP/1.1 502 Bad Gateway\r\n\r\n"); err != nil {
		ctx.Warnf("Error responding to client: %s", err)
	}
	if err := w.Close(); err != nil {
		ctx.Warnf("Error closing client connection: %s", err)
	}
}

func copyAndClose(ctx *ProxyCtx, w net.Conn, r io.Reader) {
	connOk := true
	if _, err := io.Copy(w, r); err != nil {
		connOk = false
		ctx.Warnf("Error copying to client %s", err)
	}
	if err := w.Close(); err != nil && connOk {
		ctx.Warnf("Error closing %s", err)
	}
}

func dialerFromEnv(proxy *ProxyHttpServer) func(network, addr string) (net.Conn, error) {
	https_proxy := os.Getenv("HTTPS_PROXY")
	if https_proxy == "" {
		https_proxy = os.Getenv("https_proxy")
	}
	if https_proxy == "" {
		return nil
	}
	return proxy.NewConnectDialToProxy(https_proxy)
}

func (proxy *ProxyHttpServer) NewConnectDialToProxy(https_proxy string) func(network, addr string) (net.Conn, error) {
	u, err := url.Parse(https_proxy)
	if err != nil {
		return nil
	}
	if u.Scheme == "" || u.Scheme == "http" {
		if strings.IndexRune(u.Host, ':') == -1 {
			u.Host += ":80"
		}
		return func(network, addr string) (net.Conn, error) {
			connectReq := &http.Request{
				Method: "CONNECT",
				URL:    &url.URL{Opaque: addr},
				Host:   addr,
				Header: make(http.Header),
			}
			c, err := proxy.dial(network, u.Host)
			if err != nil {
				return nil, err
			}
			connectReq.Write(c)
			// Read response.
			// Okay to use and discard buffered reader here, because
			// TLS server will not speak until spoken to.
			br := bufio.NewReader(c)
			resp, err := http.ReadResponse(br, connectReq)
			if err != nil {
				c.Close()
				return nil, err
			}
			if resp.StatusCode != 200 {
				resp, _ := ioutil.ReadAll(resp.Body)
				c.Close()
				return nil, errors.New("proxy refused connection" + string(resp))
			}
			return c, nil
		}
	}
	if u.Scheme == "https" {
		if strings.IndexRune(u.Host, ':') == -1 {
			u.Host += ":443"
		}
		return func(network, addr string) (net.Conn, error) {
			c, err := proxy.dial(network, u.Host)
			if err != nil {
				return nil, err
			}
			c = tls.Client(c, proxy.Tr.TLSClientConfig)
			connectReq := &http.Request{
				Method: "CONNECT",
				URL:    &url.URL{Opaque: addr},
				Host:   addr,
				Header: make(http.Header),
			}
			connectReq.Write(c)
			// Read response.
			// Okay to use and discard buffered reader here, because
			// TLS server will not speak until spoken to.
			br := bufio.NewReader(c)
			resp, err := http.ReadResponse(br, connectReq)
			if err != nil {
				c.Close()
				return nil, err
			}
			if resp.StatusCode != 200 {
				body, _ := ioutil.ReadAll(io.LimitReader(resp.Body, 500))
				resp.Body.Close()
				c.Close()
				return nil, errors.New("proxy refused connection" + string(body))
			}
			return c, nil
		}
	}
	return nil
}
<commit_msg>https should use context roundtripper like http does<commit_after>package goproxy

import (
	"bufio"
	"crypto/tls"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
)

type ConnectActionLiteral int

const (
	ConnectAccept = iota
	ConnectReject
	ConnectMitm
	ConnectHijack
	ConnectHTTPMitm
)

var (
	OkConnect       = &ConnectAction{Action: ConnectAccept}
	MitmConnect     = &ConnectAction{Action: ConnectMitm}
	HTTPMitmConnect = &ConnectAction{Action: ConnectHTTPMitm}
	RejectConnect   = &ConnectAction{Action: ConnectReject}
)

type ConnectAction struct {
	Action    ConnectActionLiteral
	TlsConfig *tls.Config
	Ca        *tls.Certificate
	Hijack    func(req *http.Request, client net.Conn, ctx *ProxyCtx)
}

func stripPort(s string) string {
	ix := strings.IndexRune(s, ':')
	if ix == -1 {
		return s
	}
	return s[:ix]
}

func (proxy *ProxyHttpServer) dial(network, addr string) (c net.Conn, err error) {
	if proxy.Tr.Dial != nil {
		return proxy.Tr.Dial(network, addr)
	}
	return net.Dial(network, addr)
}

func (proxy *ProxyHttpServer) connectDial(network, addr string) (c net.Conn, err error) {
	if proxy.ConnectDial == nil {
		return proxy.dial(network, addr)
	}
	return proxy.ConnectDial(network, addr)
}

func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request) {
	ctx := &ProxyCtx{Req: r, Session: atomic.AddInt64(&proxy.sess, 1), proxy: proxy}

	hij, ok := w.(http.Hijacker)
	if !ok {
		panic("httpserver does not support hijacking")
	}

	proxyClient, _, e := hij.Hijack()
	if e != nil {
		panic("Cannot hijack connection " + e.Error())
	}

	ctx.Logf("Running %d CONNECT handlers", len(proxy.httpsHandlers))
	todo, host := OkConnect, r.URL.Host
	ctx.Req = r
	for _, h := range proxy.httpsHandlers {
		newtodo, newhost := h.HandleConnect(host, ctx)
		if newtodo != nil {
			todo, host = newtodo, newhost
		}
		ctx.Logf("handler: %v %s", todo, host)
	}
	switch todo.Action {
	case ConnectAccept:
		if !hasPort.MatchString(host) {
			host += ":80"
		}
		targetSiteCon, err := proxy.connectDial("tcp", host)
		if err != nil {
			httpError(proxyClient, ctx, err)
			return
		}
		ctx.Logf("Accepting CONNECT to %s", host)
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		go copyAndClose(ctx, targetSiteCon, proxyClient)
		go copyAndClose(ctx, proxyClient, targetSiteCon)
	case ConnectHijack:
		ctx.Logf("Hijacking CONNECT to %s", host)
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		todo.Hijack(r, proxyClient, ctx)
	case ConnectHTTPMitm:
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		ctx.Logf("Assuming CONNECT is plain HTTP tunneling, mitm proxying it")
		targetSiteCon, err := proxy.connectDial("tcp", host)
		if err != nil {
			ctx.Warnf("Error dialing to %s: %s", host, err.Error())
			return
		}
		for {
			client := bufio.NewReader(proxyClient)
			remote := bufio.NewReader(targetSiteCon)
			req, err := http.ReadRequest(client)
			if err != nil && err != io.EOF {
				ctx.Warnf("cannot read request of MITM HTTP client: %+#v", err)
			}
			if err != nil {
				return
			}
			req, resp := proxy.filterRequest(req, ctx)
			if resp == nil {
				if err := req.Write(targetSiteCon); err != nil {
					httpError(proxyClient, ctx, err)
					return
				}
				resp, err = http.ReadResponse(remote, req)
				if err != nil {
					httpError(proxyClient, ctx, err)
					return
				}
			}
			resp = proxy.filterResponse(resp, ctx)
			if err := resp.Write(proxyClient); err != nil {
				httpError(proxyClient, ctx, err)
				return
			}
		}
	case ConnectMitm:
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		ctx.Logf("Assuming CONNECT is TLS, mitm proxying it")
		// this goes in a separate goroutine, so that the net/http server won't think we're
		// still handling the request even after hijacking the connection. Those HTTP CONNECT
		// request can take forever, and the server will be stuck when "closed".
		// TODO: Allow Server.Close() mechanism to shut down this connection as nicely as possible
		ca := todo.Ca
		if ca == nil {
			ca = &GoproxyCa
		}
		cert, err := signHost(*ca, []string{stripPort(host)})
		if err != nil {
			ctx.Warnf("Cannot sign host certificate with provided CA: %s", err)
			return
		}
		tlsConfig := tls.Config{}
		if todo.TlsConfig != nil {
			tlsConfig = *todo.TlsConfig
		} else {
			tlsConfig = *defaultTlsConfig
		}
		tlsConfig.Certificates = append(tlsConfig.Certificates, cert)
		go func() {
			//TODO: cache connections to the remote website
			rawClientTls := tls.Server(proxyClient, &tlsConfig)
			if err := rawClientTls.Handshake(); err != nil {
				ctx.Warnf("Cannot handshake client %v %v", r.Host, err)
				return
			}
			defer rawClientTls.Close()
			clientTlsReader := bufio.NewReader(rawClientTls)
			for !isEof(clientTlsReader) {
				req, err := http.ReadRequest(clientTlsReader)
				if err != nil && err != io.EOF {
					return
				}
				if err != nil {
					ctx.Warnf("Cannot read TLS request from mitm'd client %v %v", r.Host, err)
					return
				}
				ctx.Logf("req %v", r.Host)
				req.URL, err = url.Parse("https://" + r.Host + req.URL.String())
				req, resp := proxy.filterRequest(req, ctx)
				if resp == nil {
					if err != nil {
						ctx.Warnf("Illegal URL %s", "https://"+r.Host+req.URL.Path)
						return
					}
					removeProxyHeaders(ctx, req)
					resp, err = ctx.RoundTrip(req)
					if err != nil {
						ctx.Warnf("Cannot read TLS response from mitm'd server %v", err)
						return
					}
					ctx.Logf("resp %v", resp.Status)
				}
				resp = proxy.filterResponse(resp, ctx)
				text := resp.Status
				statusCode := strconv.Itoa(resp.StatusCode) + " "
				if strings.HasPrefix(text, statusCode) {
					text = text[len(statusCode):]
				}
				// always use 1.1 to support chunked encoding
				if _, err := io.WriteString(rawClientTls, "HTTP/1.1"+" "+statusCode+text+"\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response HTTP status from mitm'd client: %v", err)
					return
				}
				// Since we don't know the length of resp, return chunked encoded response
				// TODO: use a more reasonable scheme
				resp.Header.Del("Content-Length")
				resp.Header.Set("Transfer-Encoding", "chunked")
				if err := resp.Header.Write(rawClientTls); err != nil {
					ctx.Warnf("Cannot write TLS response header from mitm'd client: %v", err)
					return
				}
				if _, err = io.WriteString(rawClientTls, "\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response header end from mitm'd client: %v", err)
					return
				}
				chunked := newChunkedWriter(rawClientTls)
				if _, err := io.Copy(chunked, resp.Body); err != nil {
					ctx.Warnf("Cannot write TLS response body from mitm'd client: %v", err)
					return
				}
				if err := chunked.Close(); err != nil {
					ctx.Warnf("Cannot write TLS chunked EOF from mitm'd client: %v", err)
					return
				}
				if _, err = io.WriteString(rawClientTls, "\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response chunked trailer from mitm'd client: %v", err)
					return
				}
			}
			ctx.Logf("Exiting on EOF")
		}()
	case ConnectReject:
		if ctx.Resp != nil {
			if err := ctx.Resp.Write(proxyClient); err != nil {
				ctx.Warnf("Cannot write response that reject http CONNECT: %v", err)
			}
		}
		proxyClient.Close()
	}
}

func httpError(w io.WriteCloser, ctx *ProxyCtx, err error) {
	if _, err := io.WriteString(w, "HTTP/1.1 502 Bad Gateway\r\n\r\n"); err != nil {
		ctx.Warnf("Error responding to client: %s", err)
	}
	if err := w.Close(); err != nil {
		ctx.Warnf("Error closing client connection: %s", err)
	}
}

func copyAndClose(ctx *ProxyCtx, w net.Conn, r io.Reader) {
	connOk := true
	if _, err := io.Copy(w, r); err != nil {
		connOk = false
		ctx.Warnf("Error copying to client %s", err)
	}
	if err := w.Close(); err != nil && connOk {
		ctx.Warnf("Error closing %s", err)
	}
}

func dialerFromEnv(proxy *ProxyHttpServer) func(network, addr string) (net.Conn, error) {
	https_proxy := os.Getenv("HTTPS_PROXY")
	if https_proxy == "" {
		https_proxy = os.Getenv("https_proxy")
	}
	if https_proxy == "" {
		return nil
	}
	return proxy.NewConnectDialToProxy(https_proxy)
}

func (proxy *ProxyHttpServer) NewConnectDialToProxy(https_proxy string) func(network, addr string) (net.Conn, error) {
	u, err := url.Parse(https_proxy)
	if err != nil {
		return nil
	}
	if u.Scheme == "" || u.Scheme == "http" {
		if strings.IndexRune(u.Host, ':') == -1 {
			u.Host += ":80"
		}
		return func(network, addr string) (net.Conn, error) {
			connectReq := &http.Request{
				Method: "CONNECT",
				URL:    &url.URL{Opaque: addr},
				Host:   addr,
				Header: make(http.Header),
			}
			c, err := proxy.dial(network, u.Host)
			if err != nil {
				return nil, err
			}
			connectReq.Write(c)
			// Read response.
			// Okay to use and discard buffered reader here, because
			// TLS server will not speak until spoken to.
			br := bufio.NewReader(c)
			resp, err := http.ReadResponse(br, connectReq)
			if err != nil {
				c.Close()
				return nil, err
			}
			if resp.StatusCode != 200 {
				resp, _ := ioutil.ReadAll(resp.Body)
				c.Close()
				return nil, errors.New("proxy refused connection" + string(resp))
			}
			return c, nil
		}
	}
	if u.Scheme == "https" {
		if strings.IndexRune(u.Host, ':') == -1 {
			u.Host += ":443"
		}
		return func(network, addr string) (net.Conn, error) {
			c, err := proxy.dial(network, u.Host)
			if err != nil {
				return nil, err
			}
			c = tls.Client(c, proxy.Tr.TLSClientConfig)
			connectReq := &http.Request{
				Method: "CONNECT",
				URL:    &url.URL{Opaque: addr},
				Host:   addr,
				Header: make(http.Header),
			}
			connectReq.Write(c)
			// Read response.
			// Okay to use and discard buffered reader here, because
			// TLS server will not speak until spoken to.
			br := bufio.NewReader(c)
			resp, err := http.ReadResponse(br, connectReq)
			if err != nil {
				c.Close()
				return nil, err
			}
			if resp.StatusCode != 200 {
				body, _ := ioutil.ReadAll(io.LimitReader(resp.Body, 500))
				resp.Body.Close()
				c.Close()
				return nil, errors.New("proxy refused connection" + string(body))
			}
			return c, nil
		}
	}
	return nil
}
<|endoftext|>