<commit_before>/*
Copyright 2018 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

    https://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 webhook

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"path/filepath"
	"reflect"
	"time"

	"github.com/golang/glog"

	admissionv1beta1 "k8s.io/api/admission/v1beta1"
	"k8s.io/api/admissionregistration/v1beta1"
	apiv1 "k8s.io/api/core/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"

	crinformers "github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/client/informers/externalversions"
	crdlisters "github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/client/listers/sparkoperator.k8s.io/v1beta1"
	"github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/config"
	"github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/util"
)

const (
	webhookName    = "webhook.sparkoperator.k8s.io"
	serverCertFile = "server-cert.pem"
	serverKeyFile  = "server-key.pem"
	caCertFile     = "ca-cert.pem"
)

var podResource = metav1.GroupVersionResource{
	Group:    corev1.SchemeGroupVersion.Group,
	Version:  corev1.SchemeGroupVersion.Version,
	Resource: "pods",
}

// WebHook encapsulates things needed to run the webhook.
type WebHook struct {
	clientset         kubernetes.Interface
	lister            crdlisters.SparkApplicationLister
	server            *http.Server
	cert              *certBundle
	serviceRef        *v1beta1.ServiceReference
	sparkJobNamespace string
}

// New creates a new WebHook instance.
func New(
	clientset kubernetes.Interface,
	informerFactory crinformers.SharedInformerFactory,
	certDir string,
	webhookServiceNamespace string,
	webhookServiceName string,
	webhookPort int,
	jobNamespace string) (*WebHook, error) {
	cert := &certBundle{
		serverCertFile: filepath.Join(certDir, serverCertFile),
		serverKeyFile:  filepath.Join(certDir, serverKeyFile),
		caCertFile:     filepath.Join(certDir, caCertFile),
	}
	path := "/webhook"
	serviceRef := &v1beta1.ServiceReference{
		Namespace: webhookServiceNamespace,
		Name:      webhookServiceName,
		Path:      &path,
	}
	hook := &WebHook{
		clientset:         clientset,
		lister:            informerFactory.Sparkoperator().V1beta1().SparkApplications().Lister(),
		cert:              cert,
		serviceRef:        serviceRef,
		sparkJobNamespace: jobNamespace,
	}

	mux := http.NewServeMux()
	mux.HandleFunc(path, hook.serve)
	tlsConfig, err := configServerTLS(cert)
	if err != nil {
		return nil, err
	}
	hook.server = &http.Server{
		Addr:      fmt.Sprintf(":%d", webhookPort),
		Handler:   mux,
		TLSConfig: tlsConfig,
	}

	return hook, nil
}

// Start starts the admission webhook server and registers itself to the API server.
func (wh *WebHook) Start(webhookConfigName string) error {
	go func() {
		glog.Info("Starting the Spark pod admission webhook server")
		if err := wh.server.ListenAndServeTLS("", ""); err != nil && err != http.ErrServerClosed {
			glog.Errorf("error while serving the Spark pod admission webhook: %v\n", err)
		}
	}()

	return wh.selfRegistration(webhookConfigName)
}

// Stop deregisters itself with the API server and stops the admission webhook server.
func (wh *WebHook) Stop(webhookConfigName string) error {
	if err := wh.selfDeregistration(webhookConfigName); err != nil {
		return err
	}
	glog.Infof("Webhook %s deregistered", webhookConfigName)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	glog.Info("Stopping the Spark pod admission webhook server")
	return wh.server.Shutdown(ctx)
}

func (wh *WebHook) serve(w http.ResponseWriter, r *http.Request) {
	glog.V(2).Info("Serving admission request")
	var body []byte
	if r.Body != nil {
		data, err := ioutil.ReadAll(r.Body)
		if err != nil {
			glog.Errorf("failed to read the request body")
			http.Error(w, "failed to read the request body", http.StatusInternalServerError)
			return
		}
		body = data
	}

	if len(body) == 0 {
		glog.Errorf("empty request body")
		http.Error(w, "empty request body", http.StatusBadRequest)
		return
	}

	contentType := r.Header.Get("Content-Type")
	if contentType != "application/json" {
		glog.Errorf("Content-Type=%s, expected application/json", contentType)
		http.Error(w, "invalid Content-Type, expected `application/json`", http.StatusUnsupportedMediaType)
		return
	}

	var reviewResponse *admissionv1beta1.AdmissionResponse
	review := &admissionv1beta1.AdmissionReview{}
	deserializer := codecs.UniversalDeserializer()
	if _, _, err := deserializer.Decode(body, nil, review); err != nil {
		glog.Error(err)
		reviewResponse = toAdmissionResponse(err)
	} else {
		reviewResponse = mutatePods(review, wh.lister, wh.sparkJobNamespace)
	}

	response := admissionv1beta1.AdmissionReview{}
	if reviewResponse != nil {
		response.Response = reviewResponse
		if review.Request != nil {
			response.Response.UID = review.Request.UID
		}
	}

	resp, err := json.Marshal(response)
	if err != nil {
		glog.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if _, err := w.Write(resp); err != nil {
		glog.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func (wh *WebHook) selfRegistration(webhookConfigName string) error {
	client := wh.clientset.AdmissionregistrationV1beta1().MutatingWebhookConfigurations()
	existing, getErr := client.Get(webhookConfigName, metav1.GetOptions{})
	if getErr != nil && !errors.IsNotFound(getErr) {
		return getErr
	}

	ignorePolicy := v1beta1.Ignore
	caCert, err := readCertFile(wh.cert.caCertFile)
	if err != nil {
		return err
	}
	webhook := v1beta1.Webhook{
		Name: webhookName,
		Rules: []v1beta1.RuleWithOperations{
			{
				Operations: []v1beta1.OperationType{v1beta1.Create, v1beta1.Update},
				Rule: v1beta1.Rule{
					APIGroups:   []string{""},
					APIVersions: []string{"v1"},
					Resources:   []string{"pods"},
				},
			},
		},
		ClientConfig: v1beta1.WebhookClientConfig{
			Service:  wh.serviceRef,
			CABundle: caCert,
		},
		FailurePolicy: &ignorePolicy,
	}
	webhooks := []v1beta1.Webhook{webhook}

	if getErr == nil && existing != nil {
		// Update case.
		glog.Info("Updating existing MutatingWebhookConfiguration for the Spark pod admission webhook")
		if !reflect.DeepEqual(webhooks, existing.Webhooks) {
			existing.Webhooks = webhooks
			if _, err := client.Update(existing); err != nil {
				return err
			}
		}
	} else {
		// Create case.
		glog.Info("Creating a MutatingWebhookConfiguration for the Spark pod admission webhook")
		webhookConfig := &v1beta1.MutatingWebhookConfiguration{
			ObjectMeta: metav1.ObjectMeta{
				Name: webhookConfigName,
			},
			Webhooks: webhooks,
		}
		if _, err := client.Create(webhookConfig); err != nil {
			return err
		}
	}

	return nil
}

func (wh *WebHook) selfDeregistration(webhookConfigName string) error {
	client := wh.clientset.AdmissionregistrationV1beta1().MutatingWebhookConfigurations()
	return client.Delete(webhookConfigName, metav1.NewDeleteOptions(0))
}

func mutatePods(
	review *admissionv1beta1.AdmissionReview,
	lister crdlisters.SparkApplicationLister,
	sparkJobNs string) *admissionv1beta1.AdmissionResponse {
	if review.Request.Resource != podResource {
		glog.Errorf("expected resource to be %s, got %s", podResource, review.Request.Resource)
		return nil
	}

	raw := review.Request.Object.Raw
	pod := &corev1.Pod{}
	if err := json.Unmarshal(raw, pod); err != nil {
		glog.Errorf("failed to unmarshal a Pod from the raw data in the admission request: %v", err)
		return toAdmissionResponse(err)
	}

	response := &admissionv1beta1.AdmissionResponse{Allowed: true}

	if !isSparkPod(pod) || !inSparkJobNamespace(review.Request.Namespace, sparkJobNs) {
		glog.V(2).Infof("Pod %s in namespace %s is not subject to mutation", pod.GetObjectMeta().GetName(), review.Request.Namespace)
		return response
	}

	// Try getting the SparkApplication name from the annotation for that.
	appName := pod.Labels[config.SparkAppNameLabel]
	if appName == "" {
		return response
	}
	app, err := lister.SparkApplications(review.Request.Namespace).Get(appName)
	if err != nil {
		glog.Errorf("failed to get SparkApplication %s/%s: %v", review.Request.Namespace, appName, err)
		return toAdmissionResponse(err)
	}

	patchOps := patchSparkPod(pod, app)
	if len(patchOps) > 0 {
		glog.V(2).Infof("Pod %s in namespace %s is subject to mutation", pod.GetObjectMeta().GetName(), review.Request.Namespace)
		patchBytes, err := json.Marshal(patchOps)
		if err != nil {
			glog.Errorf("failed to marshal patch operations %v: %v", patchOps, err)
			return toAdmissionResponse(err)
		}
		response.Patch = patchBytes
		patchType := admissionv1beta1.PatchTypeJSONPatch
		response.PatchType = &patchType
	}

	return response
}

func toAdmissionResponse(err error) *admissionv1beta1.AdmissionResponse {
	return &admissionv1beta1.AdmissionResponse{
		Allowed: true,
		Result: &metav1.Status{
			Message: err.Error(),
			Code:    http.StatusInternalServerError,
		},
	}
}

func inSparkJobNamespace(podNs string, sparkJobNamespace string) bool {
	if sparkJobNamespace == apiv1.NamespaceAll {
		return true
	}
	return podNs == sparkJobNamespace
}

func isSparkPod(pod *corev1.Pod) bool {
	return util.IsLaunchedBySparkOperator(pod) && (util.IsDriverPod(pod) || util.IsExecutorPod(pod))
}
<commit_msg>Only use the webhook on pod creations<commit_after>/*
Copyright 2018 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

    https://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 webhook

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"path/filepath"
	"reflect"
	"time"

	"github.com/golang/glog"

	admissionv1beta1 "k8s.io/api/admission/v1beta1"
	"k8s.io/api/admissionregistration/v1beta1"
	apiv1 "k8s.io/api/core/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"

	crinformers "github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/client/informers/externalversions"
	crdlisters "github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/client/listers/sparkoperator.k8s.io/v1beta1"
	"github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/config"
	"github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/util"
)

const (
	webhookName    = "webhook.sparkoperator.k8s.io"
	serverCertFile = "server-cert.pem"
	serverKeyFile  = "server-key.pem"
	caCertFile     = "ca-cert.pem"
)

var podResource = metav1.GroupVersionResource{
	Group:    corev1.SchemeGroupVersion.Group,
	Version:  corev1.SchemeGroupVersion.Version,
	Resource: "pods",
}

// WebHook encapsulates things needed to run the webhook.
type WebHook struct {
	clientset         kubernetes.Interface
	lister            crdlisters.SparkApplicationLister
	server            *http.Server
	cert              *certBundle
	serviceRef        *v1beta1.ServiceReference
	sparkJobNamespace string
}

// New creates a new WebHook instance.
func New(
	clientset kubernetes.Interface,
	informerFactory crinformers.SharedInformerFactory,
	certDir string,
	webhookServiceNamespace string,
	webhookServiceName string,
	webhookPort int,
	jobNamespace string) (*WebHook, error) {
	cert := &certBundle{
		serverCertFile: filepath.Join(certDir, serverCertFile),
		serverKeyFile:  filepath.Join(certDir, serverKeyFile),
		caCertFile:     filepath.Join(certDir, caCertFile),
	}
	path := "/webhook"
	serviceRef := &v1beta1.ServiceReference{
		Namespace: webhookServiceNamespace,
		Name:      webhookServiceName,
		Path:      &path,
	}
	hook := &WebHook{
		clientset:         clientset,
		lister:            informerFactory.Sparkoperator().V1beta1().SparkApplications().Lister(),
		cert:              cert,
		serviceRef:        serviceRef,
		sparkJobNamespace: jobNamespace,
	}

	mux := http.NewServeMux()
	mux.HandleFunc(path, hook.serve)
	tlsConfig, err := configServerTLS(cert)
	if err != nil {
		return nil, err
	}
	hook.server = &http.Server{
		Addr:      fmt.Sprintf(":%d", webhookPort),
		Handler:   mux,
		TLSConfig: tlsConfig,
	}

	return hook, nil
}

// Start starts the admission webhook server and registers itself to the API server.
func (wh *WebHook) Start(webhookConfigName string) error {
	go func() {
		glog.Info("Starting the Spark pod admission webhook server")
		if err := wh.server.ListenAndServeTLS("", ""); err != nil && err != http.ErrServerClosed {
			glog.Errorf("error while serving the Spark pod admission webhook: %v\n", err)
		}
	}()

	return wh.selfRegistration(webhookConfigName)
}

// Stop deregisters itself with the API server and stops the admission webhook server.
func (wh *WebHook) Stop(webhookConfigName string) error {
	if err := wh.selfDeregistration(webhookConfigName); err != nil {
		return err
	}
	glog.Infof("Webhook %s deregistered", webhookConfigName)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	glog.Info("Stopping the Spark pod admission webhook server")
	return wh.server.Shutdown(ctx)
}

func (wh *WebHook) serve(w http.ResponseWriter, r *http.Request) {
	glog.V(2).Info("Serving admission request")
	var body []byte
	if r.Body != nil {
		data, err := ioutil.ReadAll(r.Body)
		if err != nil {
			glog.Errorf("failed to read the request body")
			http.Error(w, "failed to read the request body", http.StatusInternalServerError)
			return
		}
		body = data
	}

	if len(body) == 0 {
		glog.Errorf("empty request body")
		http.Error(w, "empty request body", http.StatusBadRequest)
		return
	}

	contentType := r.Header.Get("Content-Type")
	if contentType != "application/json" {
		glog.Errorf("Content-Type=%s, expected application/json", contentType)
		http.Error(w, "invalid Content-Type, expected `application/json`", http.StatusUnsupportedMediaType)
		return
	}

	var reviewResponse *admissionv1beta1.AdmissionResponse
	review := &admissionv1beta1.AdmissionReview{}
	deserializer := codecs.UniversalDeserializer()
	if _, _, err := deserializer.Decode(body, nil, review); err != nil {
		glog.Error(err)
		reviewResponse = toAdmissionResponse(err)
	} else {
		reviewResponse = mutatePods(review, wh.lister, wh.sparkJobNamespace)
	}

	response := admissionv1beta1.AdmissionReview{}
	if reviewResponse != nil {
		response.Response = reviewResponse
		if review.Request != nil {
			response.Response.UID = review.Request.UID
		}
	}

	resp, err := json.Marshal(response)
	if err != nil {
		glog.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if _, err := w.Write(resp); err != nil {
		glog.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func (wh *WebHook) selfRegistration(webhookConfigName string) error {
	client := wh.clientset.AdmissionregistrationV1beta1().MutatingWebhookConfigurations()
	existing, getErr := client.Get(webhookConfigName, metav1.GetOptions{})
	if getErr != nil && !errors.IsNotFound(getErr) {
		return getErr
	}

	ignorePolicy := v1beta1.Ignore
	caCert, err := readCertFile(wh.cert.caCertFile)
	if err != nil {
		return err
	}
	webhook := v1beta1.Webhook{
		Name: webhookName,
		Rules: []v1beta1.RuleWithOperations{
			{
				Operations: []v1beta1.OperationType{v1beta1.Create},
				Rule: v1beta1.Rule{
					APIGroups:   []string{""},
					APIVersions: []string{"v1"},
					Resources:   []string{"pods"},
				},
			},
		},
		ClientConfig: v1beta1.WebhookClientConfig{
			Service:  wh.serviceRef,
			CABundle: caCert,
		},
		FailurePolicy: &ignorePolicy,
	}
	webhooks := []v1beta1.Webhook{webhook}

	if getErr == nil && existing != nil {
		// Update case.
		glog.Info("Updating existing MutatingWebhookConfiguration for the Spark pod admission webhook")
		if !reflect.DeepEqual(webhooks, existing.Webhooks) {
			existing.Webhooks = webhooks
			if _, err := client.Update(existing); err != nil {
				return err
			}
		}
	} else {
		// Create case.
		glog.Info("Creating a MutatingWebhookConfiguration for the Spark pod admission webhook")
		webhookConfig := &v1beta1.MutatingWebhookConfiguration{
			ObjectMeta: metav1.ObjectMeta{
				Name: webhookConfigName,
			},
			Webhooks: webhooks,
		}
		if _, err := client.Create(webhookConfig); err != nil {
			return err
		}
	}

	return nil
}

func (wh *WebHook) selfDeregistration(webhookConfigName string) error {
	client := wh.clientset.AdmissionregistrationV1beta1().MutatingWebhookConfigurations()
	return client.Delete(webhookConfigName, metav1.NewDeleteOptions(0))
}

func mutatePods(
	review *admissionv1beta1.AdmissionReview,
	lister crdlisters.SparkApplicationLister,
	sparkJobNs string) *admissionv1beta1.AdmissionResponse {
	if review.Request.Resource != podResource {
		glog.Errorf("expected resource to be %s, got %s", podResource, review.Request.Resource)
		return nil
	}

	raw := review.Request.Object.Raw
	pod := &corev1.Pod{}
	if err := json.Unmarshal(raw, pod); err != nil {
		glog.Errorf("failed to unmarshal a Pod from the raw data in the admission request: %v", err)
		return toAdmissionResponse(err)
	}

	response := &admissionv1beta1.AdmissionResponse{Allowed: true}

	if !isSparkPod(pod) || !inSparkJobNamespace(review.Request.Namespace, sparkJobNs) {
		glog.V(2).Infof("Pod %s in namespace %s is not subject to mutation", pod.GetObjectMeta().GetName(), review.Request.Namespace)
		return response
	}

	// Try getting the SparkApplication name from the annotation for that.
	appName := pod.Labels[config.SparkAppNameLabel]
	if appName == "" {
		return response
	}
	app, err := lister.SparkApplications(review.Request.Namespace).Get(appName)
	if err != nil {
		glog.Errorf("failed to get SparkApplication %s/%s: %v", review.Request.Namespace, appName, err)
		return toAdmissionResponse(err)
	}

	patchOps := patchSparkPod(pod, app)
	if len(patchOps) > 0 {
		glog.V(2).Infof("Pod %s in namespace %s is subject to mutation", pod.GetObjectMeta().GetName(), review.Request.Namespace)
		patchBytes, err := json.Marshal(patchOps)
		if err != nil {
			glog.Errorf("failed to marshal patch operations %v: %v", patchOps, err)
			return toAdmissionResponse(err)
		}
		response.Patch = patchBytes
		patchType := admissionv1beta1.PatchTypeJSONPatch
		response.PatchType = &patchType
	}

	return response
}

func toAdmissionResponse(err error) *admissionv1beta1.AdmissionResponse {
	return &admissionv1beta1.AdmissionResponse{
		Allowed: true,
		Result: &metav1.Status{
			Message: err.Error(),
			Code:    http.StatusInternalServerError,
		},
	}
}

func inSparkJobNamespace(podNs string, sparkJobNamespace string) bool {
	if sparkJobNamespace == apiv1.NamespaceAll {
		return true
	}
	return podNs == sparkJobNamespace
}

func isSparkPod(pod *corev1.Pod) bool {
	return util.IsLaunchedBySparkOperator(pod) && (util.IsDriverPod(pod) || util.IsExecutorPod(pod))
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/stathat/go"
	"log"
	"runtime"
	"strings"
	"time"
)

func (zs *Zones) statHatPoster() {

	if len(Config.StatHat.ApiKey) == 0 {
		return
	}

	stathatGroups := append(serverGroups, "total", serverId)
	suffix := strings.Join(stathatGroups, ",")

	lastCounts := map[string]int64{}
	lastEdnsCounts := map[string]int64{}

	for name, zone := range *zs {
		if zone.Logging.StatHat == true {
			lastCounts[name] = zone.Metrics.Queries.Count()
			lastEdnsCounts[name] = zone.Metrics.EdnsQueries.Count()
		}
	}

	for {
		time.Sleep(60 * time.Second)

		for name, zone := range *zs {
			if zone.Logging != nil && zone.Logging.StatHat == true {

				apiKey := zone.Logging.StatHatAPI
				if len(apiKey) == 0 {
					apiKey = Config.StatHat.ApiKey
				}
				if len(apiKey) == 0 {
					continue
				}

				count := zone.Metrics.Queries.Count()
				newCount := count - lastCounts[name]
				lastCounts[name] = count
				stathat.PostEZCount("zone "+name+" queries~"+suffix, Config.StatHat.ApiKey, int(newCount))

				ednsCount := zone.Metrics.EdnsQueries.Count()
				newEdnsCount := ednsCount - lastEdnsCounts[name]
				lastEdnsCounts[name] = ednsCount
				stathat.PostEZCount("zone "+name+" edns queries~"+suffix, Config.StatHat.ApiKey, int(newEdnsCount))

			}
		}
	}
}

func statHatPoster() {

	lastQueryCount := qCounter.Count()
	stathatGroups := append(serverGroups, "total", serverId)
	suffix := strings.Join(stathatGroups, ",")
	// stathat.Verbose = true

	for {
		time.Sleep(60 * time.Second)

		if !Config.Flags.HasStatHat {
			log.Println("No stathat configuration")
			continue
		}

		log.Println("Posting to stathat")

		current := qCounter.Count()
		newQueries := current - lastQueryCount
		lastQueryCount = current

		stathat.PostEZCount("queries~"+suffix, Config.StatHat.ApiKey, int(newQueries))
		stathat.PostEZValue("goroutines "+serverId, Config.StatHat.ApiKey, float64(runtime.NumGoroutine()))

	}
}
<commit_msg>Don't send backlogged zone counts to stathat when support is enabled<commit_after>package main

import (
	"github.com/stathat/go"
	"log"
	"runtime"
	"strings"
	"time"
)

func (zs *Zones) statHatPoster() {

	if len(Config.StatHat.ApiKey) == 0 {
		return
	}

	stathatGroups := append(serverGroups, "total", serverId)
	suffix := strings.Join(stathatGroups, ",")

	lastCounts := map[string]int64{}
	lastEdnsCounts := map[string]int64{}

	for name, zone := range *zs {
		if zone.Logging.StatHat == true {
			lastCounts[name] = zone.Metrics.Queries.Count()
			lastEdnsCounts[name] = zone.Metrics.EdnsQueries.Count()
		}
	}

	for {
		time.Sleep(60 * time.Second)

		for name, zone := range *zs {

			count := zone.Metrics.Queries.Count()
			newCount := count - lastCounts[name]
			lastCounts[name] = count

			if zone.Logging != nil && zone.Logging.StatHat == true {

				apiKey := zone.Logging.StatHatAPI
				if len(apiKey) == 0 {
					apiKey = Config.StatHat.ApiKey
				}
				if len(apiKey) == 0 {
					continue
				}
				stathat.PostEZCount("zone "+name+" queries~"+suffix, Config.StatHat.ApiKey, int(newCount))

				ednsCount := zone.Metrics.EdnsQueries.Count()
				newEdnsCount := ednsCount - lastEdnsCounts[name]
				lastEdnsCounts[name] = ednsCount
				stathat.PostEZCount("zone "+name+" edns queries~"+suffix, Config.StatHat.ApiKey, int(newEdnsCount))

			}
		}
	}
}

func statHatPoster() {

	lastQueryCount := qCounter.Count()
	stathatGroups := append(serverGroups, "total", serverId)
	suffix := strings.Join(stathatGroups, ",")
	// stathat.Verbose = true

	for {
		time.Sleep(60 * time.Second)

		if !Config.Flags.HasStatHat {
			log.Println("No stathat configuration")
			continue
		}

		log.Println("Posting to stathat")

		current := qCounter.Count()
		newQueries := current - lastQueryCount
		lastQueryCount = current

		stathat.PostEZCount("queries~"+suffix, Config.StatHat.ApiKey, int(newQueries))
		stathat.PostEZValue("goroutines "+serverId, Config.StatHat.ApiKey, float64(runtime.NumGoroutine()))

	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package templates

import (
	"bytes"
	"fmt"
	"strings"
	"text/template"
	"unicode"

	"k8s.io/kubectl/pkg/util/term"

	"github.com/spf13/cobra"
	flag "github.com/spf13/pflag"
)

type FlagExposer interface {
	ExposeFlags(cmd *cobra.Command, flags ...string) FlagExposer
}

func ActsAsRootCommand(cmd *cobra.Command, filters []string, groups ...CommandGroup) FlagExposer {
	if cmd == nil {
		panic("nil root command")
	}
	templater := &templater{
		RootCmd:       cmd,
		UsageTemplate: MainUsageTemplate(),
		HelpTemplate:  MainHelpTemplate(),
		CommandGroups: groups,
		Filtered:      filters,
	}
	cmd.SetUsageFunc(templater.UsageFunc())
	cmd.SetHelpFunc(templater.HelpFunc())
	return templater
}

func UseOptionsTemplates(cmd *cobra.Command) {
	templater := &templater{
		UsageTemplate: OptionsUsageTemplate(),
		HelpTemplate:  OptionsHelpTemplate(),
	}
	cmd.SetUsageFunc(templater.UsageFunc())
	cmd.SetHelpFunc(templater.HelpFunc())
}

type templater struct {
	UsageTemplate string
	HelpTemplate  string
	RootCmd       *cobra.Command
	CommandGroups
	Filtered []string
}

func (templater *templater) ExposeFlags(cmd *cobra.Command, flags ...string) FlagExposer {
	cmd.SetUsageFunc(templater.UsageFunc(flags...))
	return templater
}

func (templater *templater) HelpFunc() func(*cobra.Command, []string) {
	return func(c *cobra.Command, s []string) {
		t := template.New("help")
		t.Funcs(templater.templateFuncs())
		template.Must(t.Parse(templater.HelpTemplate))
		out := term.NewResponsiveWriter(c.OutOrStdout())
		err := t.Execute(out, c)
		if err != nil {
			c.Println(err)
		}
	}
}

func (templater *templater) UsageFunc(exposedFlags ...string) func(*cobra.Command) error {
	return func(c *cobra.Command) error {
		t := template.New("usage")
		t.Funcs(templater.templateFuncs(exposedFlags...))
		template.Must(t.Parse(templater.UsageTemplate))
		out := term.NewResponsiveWriter(c.OutOrStderr())
		return t.Execute(out, c)
	}
}

func (templater *templater) templateFuncs(exposedFlags ...string) template.FuncMap {
	return template.FuncMap{
		"trim":                strings.TrimSpace,
		"trimRight":           func(s string) string { return strings.TrimRightFunc(s, unicode.IsSpace) },
		"trimLeft":            func(s string) string { return strings.TrimLeftFunc(s, unicode.IsSpace) },
		"gt":                  cobra.Gt,
		"eq":                  cobra.Eq,
		"rpad":                rpad,
		"appendIfNotPresent":  appendIfNotPresent,
		"flagsNotIntersected": flagsNotIntersected,
		"visibleFlags":        visibleFlags,
		"flagsUsages":         flagsUsages,
		"cmdGroups":           templater.cmdGroups,
		"cmdGroupsString":     templater.cmdGroupsString,
		"rootCmd":             templater.rootCmdName,
		"isRootCmd":           templater.isRootCmd,
		"optionsCmdFor":       templater.optionsCmdFor,
		"usageLine":           templater.usageLine,
		"exposed": func(c *cobra.Command) *flag.FlagSet {
			exposed := flag.NewFlagSet("exposed", flag.ContinueOnError)
			if len(exposedFlags) > 0 {
				for _, name := range exposedFlags {
					if flag := c.Flags().Lookup(name); flag != nil {
						exposed.AddFlag(flag)
					}
				}
			}
			return exposed
		},
	}
}

func (templater *templater) cmdGroups(c *cobra.Command, all []*cobra.Command) []CommandGroup {
	if len(templater.CommandGroups) > 0 && c == templater.RootCmd {
		all = filter(all, templater.Filtered...)
		return AddAdditionalCommands(templater.CommandGroups, "Other Commands:", all)
	}
	all = filter(all, "options")
	return []CommandGroup{
		{
			Message:  "Available Commands:",
			Commands: all,
		},
	}
}

func (t *templater) cmdGroupsString(c *cobra.Command) string {
	groups := []string{}
	for _, cmdGroup := range t.cmdGroups(c, c.Commands()) {
		cmds := []string{cmdGroup.Message}
		for _, cmd := range cmdGroup.Commands {
			if cmd.IsAvailableCommand() {
				cmds = append(cmds, "  "+rpad(cmd.Name(), cmd.NamePadding())+" "+cmd.Short)
			}
		}
		groups = append(groups, strings.Join(cmds, "\n"))
	}
	return strings.Join(groups, "\n\n")
}

func (t *templater) rootCmdName(c *cobra.Command) string {
	return t.rootCmd(c).CommandPath()
}

func (t *templater) isRootCmd(c *cobra.Command) bool {
	return t.rootCmd(c) == c
}

func (t *templater) parents(c *cobra.Command) []*cobra.Command {
	parents := []*cobra.Command{c}
	for current := c; !t.isRootCmd(current) && current.HasParent(); {
		current = current.Parent()
		parents = append(parents, current)
	}
	return parents
}

func (t *templater) rootCmd(c *cobra.Command) *cobra.Command {
	if c != nil && !c.HasParent() {
		return c
	}
	if t.RootCmd == nil {
		panic("nil root cmd")
	}
	return t.RootCmd
}

func (t *templater) optionsCmdFor(c *cobra.Command) string {
	if !c.Runnable() {
		return ""
	}
	rootCmdStructure := t.parents(c)
	for i := len(rootCmdStructure) - 1; i >= 0; i-- {
		cmd := rootCmdStructure[i]
		if _, _, err := cmd.Find([]string{"options"}); err == nil {
			return cmd.CommandPath() + " options"
		}
	}
	return ""
}

func (t *templater) usageLine(c *cobra.Command) string {
	usage := c.UseLine()
	suffix := "[options]"
	if c.HasFlags() && !strings.Contains(usage, suffix) {
		usage += " " + suffix
	}
	return usage
}

func flagsUsages(f *flag.FlagSet) string {
	x := new(bytes.Buffer)

	f.VisitAll(func(flag *flag.Flag) {
		if flag.Hidden {
			return
		}
		format := "--%s=%s: %s\n"

		if flag.Value.Type() == "string" {
			format = "--%s='%s': %s\n"
		}

		if len(flag.Shorthand) > 0 {
			format = "  -%s, " + format
		} else {
			format = "   %s   " + format
		}

		fmt.Fprintf(x, format, flag.Shorthand, flag.Name, flag.DefValue, flag.Usage)
	})

	return x.String()
}

func rpad(s string, padding int) string {
	template := fmt.Sprintf("%%-%ds", padding)
	return fmt.Sprintf(template, s)
}

func appendIfNotPresent(s, stringToAppend string) string {
	if strings.Contains(s, stringToAppend) {
		return s
	}
	return s + " " + stringToAppend
}

func flagsNotIntersected(l *flag.FlagSet, r *flag.FlagSet) *flag.FlagSet {
	f := flag.NewFlagSet("notIntersected", flag.ContinueOnError)
	l.VisitAll(func(flag *flag.Flag) {
		if r.Lookup(flag.Name) == nil {
			f.AddFlag(flag)
		}
	})
	return f
}

func visibleFlags(l *flag.FlagSet) *flag.FlagSet {
	hidden := "help"
	f := flag.NewFlagSet("visible", flag.ContinueOnError)
	l.VisitAll(func(flag *flag.Flag) {
		if flag.Name != hidden {
			f.AddFlag(flag)
		}
	})
	return f
}

func filter(cmds []*cobra.Command, names ...string) []*cobra.Command {
	out := []*cobra.Command{}
	for _, c := range cmds {
		if c.Hidden {
			continue
		}
		skip := false
		for _, name := range names {
			if name == c.Name() {
				skip = true
				break
			}
		}
		if skip {
			continue
		}
		out = append(out, c)
	}
	return out
}
<commit_msg>kubectl: remove usage info from bad flag error msg<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package templates

import (
	"bytes"
	"fmt"
	"strings"
	"text/template"
	"unicode"

	"k8s.io/kubectl/pkg/util/term"

	"github.com/spf13/cobra"
	flag "github.com/spf13/pflag"
)

type FlagExposer interface {
	ExposeFlags(cmd *cobra.Command, flags ...string) FlagExposer
}

func ActsAsRootCommand(cmd *cobra.Command, filters []string, groups ...CommandGroup) FlagExposer {
	if cmd == nil {
		panic("nil root command")
	}
	templater := &templater{
		RootCmd:       cmd,
		UsageTemplate: MainUsageTemplate(),
		HelpTemplate:  MainHelpTemplate(),
		CommandGroups: groups,
		Filtered:      filters,
	}
	cmd.SetFlagErrorFunc(templater.FlagErrorFunc())
	cmd.SetUsageFunc(templater.UsageFunc())
	cmd.SetHelpFunc(templater.HelpFunc())
	return templater
}

func UseOptionsTemplates(cmd *cobra.Command) {
	templater := &templater{
		UsageTemplate: OptionsUsageTemplate(),
		HelpTemplate:  OptionsHelpTemplate(),
	}
	cmd.SetUsageFunc(templater.UsageFunc())
	cmd.SetHelpFunc(templater.HelpFunc())
}

type templater struct {
	UsageTemplate string
	HelpTemplate  string
	RootCmd       *cobra.Command
	CommandGroups
	Filtered []string
}

func (templater *templater) FlagErrorFunc(exposedFlags ...string) func(*cobra.Command, error) error {
	return func(c *cobra.Command, err error) error {
		c.SilenceUsage = true
		switch c.CalledAs() {
		case "options":
			return fmt.Errorf("%s\nRun '%s' without flags.", err, c.CommandPath())
		default:
			return fmt.Errorf("%s\nSee '%s --help' for usage.", err, c.CommandPath())
		}
	}
}

func (templater *templater) ExposeFlags(cmd *cobra.Command, flags ...string) FlagExposer {
	cmd.SetUsageFunc(templater.UsageFunc(flags...))
	return templater
}

func (templater *templater) HelpFunc() func(*cobra.Command, []string) {
	return func(c *cobra.Command, s []string) {
		t := template.New("help")
		t.Funcs(templater.templateFuncs())
		template.Must(t.Parse(templater.HelpTemplate))
		out := term.NewResponsiveWriter(c.OutOrStdout())
		err := t.Execute(out, c)
		if err != nil {
			c.Println(err)
		}
	}
}

func (templater *templater) UsageFunc(exposedFlags ...string) func(*cobra.Command) error {
	return func(c *cobra.Command) error {
		t := template.New("usage")
		t.Funcs(templater.templateFuncs(exposedFlags...))
		template.Must(t.Parse(templater.UsageTemplate))
		out := term.NewResponsiveWriter(c.OutOrStderr())
		return t.Execute(out, c)
	}
}

func (templater *templater) templateFuncs(exposedFlags ...string) template.FuncMap {
	return template.FuncMap{
		"trim":                strings.TrimSpace,
		"trimRight":           func(s string) string { return strings.TrimRightFunc(s, unicode.IsSpace) },
		"trimLeft":            func(s string) string { return strings.TrimLeftFunc(s, unicode.IsSpace) },
		"gt":                  cobra.Gt,
		"eq":                  cobra.Eq,
		"rpad":                rpad,
		"appendIfNotPresent":  appendIfNotPresent,
		"flagsNotIntersected": flagsNotIntersected,
		"visibleFlags":        visibleFlags,
		"flagsUsages":         flagsUsages,
		"cmdGroups":           templater.cmdGroups,
		"cmdGroupsString":     templater.cmdGroupsString,
		"rootCmd":             templater.rootCmdName,
		"isRootCmd":           templater.isRootCmd,
		"optionsCmdFor":       templater.optionsCmdFor,
		"usageLine":           templater.usageLine,
		"exposed": func(c *cobra.Command) *flag.FlagSet {
			exposed := flag.NewFlagSet("exposed", flag.ContinueOnError)
			if len(exposedFlags) > 0 {
				for _, name := range exposedFlags {
					if flag := c.Flags().Lookup(name); flag != nil {
						exposed.AddFlag(flag)
					}
				}
			}
			return exposed
		},
	}
}

func (templater *templater) cmdGroups(c *cobra.Command, all []*cobra.Command) []CommandGroup {
	if len(templater.CommandGroups) > 0 && c == templater.RootCmd {
		all = filter(all, templater.Filtered...)
		return AddAdditionalCommands(templater.CommandGroups, "Other Commands:", all)
	}
	all = filter(all, "options")
	return []CommandGroup{
		{
			Message:  "Available Commands:",
			Commands: all,
		},
	}
}

func (t *templater) cmdGroupsString(c *cobra.Command) string {
	groups := []string{}
	for _, cmdGroup := range t.cmdGroups(c, c.Commands()) {
		cmds := []string{cmdGroup.Message}
		for _, cmd := range cmdGroup.Commands {
			if cmd.IsAvailableCommand() {
				cmds = append(cmds, "  "+rpad(cmd.Name(), cmd.NamePadding())+" "+cmd.Short)
			}
		}
		groups = append(groups, strings.Join(cmds, "\n"))
	}
	return strings.Join(groups, "\n\n")
}

func (t *templater) rootCmdName(c *cobra.Command) string {
	return t.rootCmd(c).CommandPath()
}

func (t *templater) isRootCmd(c *cobra.Command) bool {
	return t.rootCmd(c) == c
}

func (t *templater) parents(c *cobra.Command) []*cobra.Command {
	parents := []*cobra.Command{c}
	for current := c; !t.isRootCmd(current) && current.HasParent(); {
		current = current.Parent()
		parents = append(parents, current)
	}
	return parents
}

func (t *templater) rootCmd(c *cobra.Command) *cobra.Command {
	if c != nil && !c.HasParent() {
		return c
	}
	if t.RootCmd == nil {
		panic("nil root cmd")
	}
	return t.RootCmd
}

func (t *templater) optionsCmdFor(c *cobra.Command) string {
	if !c.Runnable() {
		return ""
	}
	rootCmdStructure := t.parents(c)
	for i := len(rootCmdStructure) - 1; i >= 0; i-- {
		cmd := rootCmdStructure[i]
		if _, _, err := cmd.Find([]string{"options"}); err == nil {
			return cmd.CommandPath() + " options"
		}
	}
	return ""
}

func (t *templater) usageLine(c *cobra.Command) string {
	usage := c.UseLine()
	suffix := "[options]"
	if c.HasFlags() && !strings.Contains(usage, suffix) {
		usage += " " + suffix
	}
	return usage
}

func flagsUsages(f *flag.FlagSet) string {
	x := new(bytes.Buffer)

	f.VisitAll(func(flag *flag.Flag) {
		if flag.Hidden {
			return
		}
		format := "--%s=%s: %s\n"

		if flag.Value.Type() == "string" {
			format = "--%s='%s': %s\n"
		}

		if len(flag.Shorthand) > 0 {
			format = "  -%s, " + format
		} else {
			format = "   %s   " + format
		}

		fmt.Fprintf(x, format, flag.Shorthand, flag.Name, flag.DefValue, flag.Usage)
	})

	return x.String()
}

func rpad(s string, padding int) string {
	template := fmt.Sprintf("%%-%ds", padding)
	return fmt.Sprintf(template, s)
}

func appendIfNotPresent(s, stringToAppend string) string {
	if strings.Contains(s, stringToAppend) {
		return s
	}
	return s + " " + stringToAppend
}

func flagsNotIntersected(l *flag.FlagSet, r *flag.FlagSet) *flag.FlagSet {
	f := flag.NewFlagSet("notIntersected", flag.ContinueOnError)
	l.VisitAll(func(flag *flag.Flag) {
		if r.Lookup(flag.Name) == nil {
			f.AddFlag(flag)
		}
	})
	return f
}

func visibleFlags(l *flag.FlagSet) *flag.FlagSet {
	hidden := "help"
	f := flag.NewFlagSet("visible", flag.ContinueOnError)
	l.VisitAll(func(flag *flag.Flag) {
		if flag.Name != hidden {
			f.AddFlag(flag)
		}
	})
	return f
}

func filter(cmds []*cobra.Command, names ...string) []*cobra.Command {
	out := []*cobra.Command{}
	for _, c := range cmds {
		if c.Hidden {
			continue
		}
		skip := false
		for _, name := range names {
			if name == c.Name() {
				skip = true
				break
			}
		}
		if skip {
			continue
		}
		out = append(out, c)
	}
	return out
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Kubernetes Authors.

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

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

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

package adapter

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	containeranalysis "cloud.google.com/go/containeranalysis/apiv1"
	"cloud.google.com/go/storage"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/html"
	"google.golang.org/api/iterator"
	grafeaspb "google.golang.org/genproto/googleapis/grafeas/v1"
)

func uploadFile(directory, filename, bucket string) error {
	const timeout = 60
	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		return errors.Errorf("storage.NewClient: %v", err)
	}
	defer client.Close()

	// Open local file
	f, err := os.Open(directory + filename)
	if err != nil {
		return errors.Errorf("os.Open: %v", err)
	}
	defer f.Close()

	ctx, cancel := context.WithTimeout(ctx, time.Second*timeout)
	defer cancel()

	// Upload the object with storage.Writer
	wc := client.Bucket(bucket).Object(filename).NewWriter(ctx)
	if _, err = io.Copy(wc, f); err != nil {
		return errors.Errorf("io.Copy: %v", err)
	}

	if err := wc.Close(); err != nil {
		return errors.Errorf("Writer.Close: %v", err)
	}

	return nil
}

// GetAllVulnerabilities gets all of the vulnerability occurrences associated
// with images in a specific project using the Container Analysis Service.
func GetAllVulnerabilities(
	projectID string,
	pageSize int32,
) ([]*grafeaspb.Occurrence, error) {
	ctx := context.Background()
	client, err := containeranalysis.NewClient(ctx)
	if err != nil {
		return nil, errors.Errorf("NewClient: %v", err)
	}
	defer client.Close()

	req := &grafeaspb.ListOccurrencesRequest{
		Parent:   fmt.Sprintf("projects/%s", projectID),
		Filter:   fmt.Sprintf("kind = %q", "VULNERABILITY"),
		PageSize: pageSize,
	}

	logrus.Info("listing the vulnerabilities, will take a while...")
	var occurrenceList []*grafeaspb.Occurrence
	it := client.GetGrafeasClient().ListOccurrences(ctx, req)
	for {
		var occ *grafeaspb.Occurrence
		var err error
		occ, err = it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, errors.Errorf("occurrence iteration error: %v", err)
		}
		occurrenceList = append(occurrenceList, occ)
	}
	logrus.Infof("done listing the vulnerabilities")

	return occurrenceList, err
}

func parseImageResourceURL(resourceURL string) (registryImageName, digest string) {
	FQIN := path.Base(resourceURL)
	splitFQIN := strings.Split(FQIN, "@")
	registryImageName, digest = splitFQIN[0], splitFQIN[1]

	return registryImageName, digest
}

func parseVulnName(noteName string) string {
	return path.Base(noteName)
}

// GenerateVulnerabilityBreakdown parses the a slice of vulnerability
// occurrences into a breakdown that only contains the necessary information
// for each vulnerability.
func GenerateVulnerabilityBreakdown(
	productionVulnerabilities []*grafeaspb.Occurrence,
) map[string]ImageVulnBreakdown {
	vulnBreakdowns := make(map[string]ImageVulnBreakdown)

	for _, occ := range productionVulnerabilities {
		// resourceURI is a url pointing to a specific image
		// in the form gcr.io/project/foo@sha256:111
		if _, found := vulnBreakdowns[occ.ResourceUri]; !found {
			imageName, imageDigest := parseImageResourceURL(occ.ResourceUri)
			vulnBreakdowns[occ.ResourceUri] = ImageVulnBreakdown{
				occ.ResourceUri,
				imageName,
				imageDigest,
				0,
				[]string{},
				[]string{},
			}
		}

		imageVulnBreakdown := vulnBreakdowns[occ.ResourceUri]
		imageVulnBreakdown.NumVulnerabilities++

		vulnName := parseVulnName(occ.NoteName)
		vuln := occ.GetVulnerability()
		if vuln.GetSeverity() == grafeaspb.Severity_CRITICAL {
			imageVulnBreakdown.CriticalVulnerabilities = append(
				imageVulnBreakdown.CriticalVulnerabilities,
				vulnName,
			)
		}
		if vuln.GetFixAvailable() {
			imageVulnBreakdown.FixableVulnerabilities = append(
				imageVulnBreakdown.FixableVulnerabilities,
				vulnName,
			)
		}
		vulnBreakdowns[occ.ResourceUri] = imageVulnBreakdown
	}

	return vulnBreakdowns
}

// UpdateVulnerabilityDashboard updates the vulnerability dashboard by uploading
// the lastest versions of all the vulnerability dashboard's files.
func UpdateVulnerabilityDashboard(
	dashboardPath string,
	vulnProject string,
	dashboardBucket string,
	pageSize int32,
) error {
	dashboardHTML := dashboardPath + "dashboard.html"
	logrus.Infof("opening %s", dashboardHTML)
	htmlReader, openErr := os.Open(dashboardHTML)
	if openErr != nil {
		return errors.Wrap(openErr, "opening dashboard file")
	}

	logrus.Infof("parsing %s", dashboardHTML)
	_, err := html.Parse(htmlReader)
	if err != nil {
		return errors.Errorf("dashboard.html is not valid HTML: %v", err)
	}

	logrus.Infof("uploading %s to gcs", dashboardHTML)
	err = uploadFile(dashboardPath, "dashboard.html", dashboardBucket)
	if err != nil {
		return errors.Errorf("Unable to upload latest version of "+
			"dashboard HTML: %v", err)
	}

	logrus.Info("uploading updated dashboard.js to gcs")
	err = uploadFile(dashboardPath, "dashboard.js", dashboardBucket)
	if err != nil {
		return errors.Errorf("Unable to upload latest version of "+
			"dashboard JS: %v", err)
	}

	logrus.Infof("checking all vulnerabilities for %s", vulnProject)
	productionVulnerabilities, getVulnErr := GetAllVulnerabilities(vulnProject, pageSize)
	if getVulnErr != nil {
		return errors.Wrap(getVulnErr, "getting all vulnerabilities")
	}

	logrus.Infof("parsing the vulnerabilities for %s", vulnProject)
	vulnBreakdowns := GenerateVulnerabilityBreakdown(productionVulnerabilities)
	jsonFile, err := json.MarshalIndent(vulnBreakdowns, "", " ")
	if err != nil {
		return errors.Errorf("Unable to generate dashboard json: %v", err)
	}

	dashboardJSON := dashboardPath + "dashboard.json"
	logrus.Infof("writing the vulnerabilities for %s in the file %s", vulnProject, dashboardJSON)
	err = ioutil.WriteFile(dashboardJSON, jsonFile, 0644)
	if err != nil {
		return errors.Errorf("Unable to create temporary local"+
			"JSON file for the dashboard: %v", err)
	}

	logrus.Infof("uploading updated %s to gcs", dashboardJSON)
	err = uploadFile(dashboardPath, "dashboard.json", dashboardBucket)
	if err != nil {
		return errors.Errorf("Unable to upload latest version of "+
			"dashboard JSON: %v", err)
	}

	return nil
}
<commit_msg>vulndash: set default page size<commit_after>/*
Copyright 2020 The Kubernetes Authors.

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

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

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

package adapter

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	containeranalysis "cloud.google.com/go/containeranalysis/apiv1"
	"cloud.google.com/go/storage"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/html"
	"google.golang.org/api/iterator"
	grafeaspb "google.golang.org/genproto/googleapis/grafeas/v1"
)

// DefaultPageSize to be used in the vulnerabilities list to set how many items
// will be retrieved in each request
const DefaultPageSize = 200

func uploadFile(directory, filename, bucket string) error {
	const timeout = 60
	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		return errors.Errorf("storage.NewClient: %v", err)
	}
	defer client.Close()

	// Open local file
	f, err := os.Open(directory + filename)
	if err != nil {
		return errors.Errorf("os.Open: %v", err)
	}
	defer f.Close()

	ctx, cancel := context.WithTimeout(ctx, time.Second*timeout)
	defer cancel()

	// Upload the object with storage.Writer
	wc := client.Bucket(bucket).Object(filename).NewWriter(ctx)
	if _, err = io.Copy(wc, f); err != nil {
		return errors.Errorf("io.Copy: %v", err)
	}

	if err := wc.Close(); err != nil {
		return errors.Errorf("Writer.Close: %v", err)
	}

	return nil
}

// GetAllVulnerabilities gets all of the vulnerability occurrences associated
// with images in a specific project using the Container Analysis Service.
func GetAllVulnerabilities(
	projectID string,
	pageSize int32,
) ([]*grafeaspb.Occurrence, error) {
	ctx := context.Background()
	client, err := containeranalysis.NewClient(ctx)
	if err != nil {
		return nil, errors.Errorf("NewClient: %v", err)
	}
	defer client.Close()

	if pageSize <= 0 {
		pageSize = DefaultPageSize
	}
	req := &grafeaspb.ListOccurrencesRequest{
		Parent:   fmt.Sprintf("projects/%s", projectID),
		Filter:   fmt.Sprintf("kind = %q", "VULNERABILITY"),
		PageSize: pageSize,
	}

	logrus.Info("listing the vulnerabilities, will take a while...")
	var occurrenceList []*grafeaspb.Occurrence
	it := client.GetGrafeasClient().ListOccurrences(ctx, req)
	for {
		var occ *grafeaspb.Occurrence
		var err error
		occ, err = it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, errors.Errorf("occurrence iteration error: %v", err)
		}
		occurrenceList = append(occurrenceList, occ)
	}
	logrus.Infof("done listing the vulnerabilities")

	return occurrenceList, err
}

func parseImageResourceURL(resourceURL string) (registryImageName, digest string) {
	FQIN := path.Base(resourceURL)
	splitFQIN := strings.Split(FQIN, "@")
	registryImageName, digest = splitFQIN[0], splitFQIN[1]

	return registryImageName, digest
}

func parseVulnName(noteName string) string {
	return path.Base(noteName)
}

// GenerateVulnerabilityBreakdown parses the a slice of vulnerability
// occurrences into a breakdown that only contains the necessary information
// for each vulnerability.
func GenerateVulnerabilityBreakdown(
	productionVulnerabilities []*grafeaspb.Occurrence,
) map[string]ImageVulnBreakdown {
	vulnBreakdowns := make(map[string]ImageVulnBreakdown)

	for _, occ := range productionVulnerabilities {
		// resourceURI is a url pointing to a specific image
		// in the form gcr.io/project/foo@sha256:111
		if _, found := vulnBreakdowns[occ.ResourceUri]; !found {
			imageName, imageDigest := parseImageResourceURL(occ.ResourceUri)
			vulnBreakdowns[occ.ResourceUri] = ImageVulnBreakdown{
				occ.ResourceUri,
				imageName,
				imageDigest,
				0,
				[]string{},
				[]string{},
			}
		}

		imageVulnBreakdown := vulnBreakdowns[occ.ResourceUri]
		imageVulnBreakdown.NumVulnerabilities++

		vulnName := parseVulnName(occ.NoteName)
		vuln := occ.GetVulnerability()
		if vuln.GetSeverity() == grafeaspb.Severity_CRITICAL {
			imageVulnBreakdown.CriticalVulnerabilities = append(
				imageVulnBreakdown.CriticalVulnerabilities,
				vulnName,
			)
		}
		if vuln.GetFixAvailable() {
			imageVulnBreakdown.FixableVulnerabilities = append(
				imageVulnBreakdown.FixableVulnerabilities,
				vulnName,
			)
		}
		vulnBreakdowns[occ.ResourceUri] = imageVulnBreakdown
	}

	return vulnBreakdowns
}

// UpdateVulnerabilityDashboard updates the vulnerability dashboard by uploading
// the lastest versions of all the vulnerability dashboard's files.
func UpdateVulnerabilityDashboard(
	dashboardPath string,
	vulnProject string,
	dashboardBucket string,
	pageSize int32,
) error {
	dashboardHTML := dashboardPath + "dashboard.html"
	logrus.Infof("opening %s", dashboardHTML)
	htmlReader, openErr := os.Open(dashboardHTML)
	if openErr != nil {
		return errors.Wrap(openErr, "opening dashboard file")
	}

	logrus.Infof("parsing %s", dashboardHTML)
	_, err := html.Parse(htmlReader)
	if err != nil {
		return errors.Errorf("dashboard.html is not valid HTML: %v", err)
	}

	logrus.Infof("uploading %s to gcs", dashboardHTML)
	err = uploadFile(dashboardPath, "dashboard.html", dashboardBucket)
	if err != nil {
		return errors.Errorf("Unable to upload latest version of "+
			"dashboard HTML: %v", err)
	}

	logrus.Info("uploading updated dashboard.js to gcs")
	err = uploadFile(dashboardPath, "dashboard.js", dashboardBucket)
	if err != nil {
		return errors.Errorf("Unable to upload latest version of "+
			"dashboard JS: %v", err)
	}

	logrus.Infof("checking all vulnerabilities for %s", vulnProject)
	productionVulnerabilities, getVulnErr := GetAllVulnerabilities(vulnProject, pageSize)
	if getVulnErr != nil {
		return errors.Wrap(getVulnErr, "getting all vulnerabilities")
	}

	logrus.Infof("parsing the vulnerabilities for %s", vulnProject)
	vulnBreakdowns := GenerateVulnerabilityBreakdown(productionVulnerabilities)
	jsonFile, err := json.MarshalIndent(vulnBreakdowns, "", " ")
	if err != nil {
		return errors.Errorf("Unable to generate dashboard json: %v", err)
	}

	dashboardJSON := dashboardPath + "dashboard.json"
	logrus.Infof("writing the vulnerabilities for %s in the file %s", vulnProject, dashboardJSON)
	err = ioutil.WriteFile(dashboardJSON, jsonFile, 0644)
	if err != nil {
		return errors.Errorf("Unable to create temporary local"+
			"JSON file for the dashboard: %v", err)
	}

	logrus.Infof("uploading updated %s to gcs", dashboardJSON)
	err = uploadFile(dashboardPath, "dashboard.json", dashboardBucket)
	if err != nil {
		return errors.Errorf("Unable to upload latest version of "+
			"dashboard JSON: %v", err)
	}

	return nil
}
<|endoftext|>


<commit_before>// Go implementation of the Paice/Husk Stemming algorithm:
// http://www.comp.lancs.ac.uk/computing/research/stemming/Links/paice.htm
package paicehusk

import (
	"regexp"
	"strconv"
	"strings"
)

// A representation of a stemming rule

// A representation of a stemming rule
type rule struct {

	// The suffix the rule is to act on
	suf string

	// True if the stem is required intact for the rule to operate
	intact bool

	// Number of letters to strip off the stem
	num int

	// A suffix to append to the stem
	app string

	// True if the stem should be stemmed further
	cont bool
}

var DefaultRules = NewRuleTable(strings.Split(defaultRules, "\n"))

// RuleTable stores rules based on the final letter of the suffix they act on
type RuleTable struct {
	Table map[string][]*rule
}

// NewRuleTable returns a new RuleTable instance
func NewRuleTable(f []string) (table *RuleTable) {
	table = &RuleTable{Table: make(map[string][]*rule)}
	for _, s := range f {
		if r, ok := ParseRule(s); ok {
			table.Table[r.suf[:1]] = append(table.Table[r.suf[:1]], r)
		}
	}
	return
}

// Regex for ValidRule
var reg = regexp.MustCompile("[a-zA-Z]*\\*?[0-9][a-zA-z]*[.>]")

// Validates a rule
func ValidRule(s string) (rule string, ok bool) {
	ok = true
	// Find the first instance of a rule in the provided string
	rule = reg.FindString(s)
	if rule == "" {
		ok = false
	}
	return
}

// Regexes for ParseRule
var suf = regexp.MustCompile("[a-zA-Z]+")
var intact = regexp.MustCompile("[*]")
var num = regexp.MustCompile("[0-9]")
var app = regexp.MustCompile("[0-9][a-zA-Z]+")

// ParseRule parses a rule in the form:
// |suffix|intact flag|number to strip|Append|Continue flag
//
// Eg, a rule: ht*2. Means if the stem is still intact, strip the
// suffix th and make no further attempts to stem the word.
//
// Rule nois4j> Means strip the sion suffix, append a j and check
// for any more rules to follow
func ParseRule(s string) (r *rule, ok bool) {
	s, ok = ValidRule(s)
	if !ok {
		return nil, false
	}

	r = new(rule)

	r.suf = suf.FindString(s)
	if intact.FindString(s) == "" {
		r.intact = false
	} else {
		r.intact = true
	}
	if i, err := strconv.ParseInt(num.FindString(s), 0, 0); err != nil {
		panic(err)
	} else {
		r.num = int(i)
	}
	if append := app.FindString(s); len(append) > 0 {
		r.app = app.FindString(s)[1:]
	} else {
		r.app = ""
	}

	if s[len(s)-1:] == ">" {
		r.cont = true
	} else {
		r.cont = false
	}
	return r, true
}

// Stem a string, word, based on the rules in *RuleTable r, by following
// the algorithm described at:
// http://www.comp.lancs.ac.uk/computing/research/stemming/Links/paice.htm
func (r *RuleTable) Stem(word string) string {
	stem := []rune(strings.ToLower(word))
	current := stem

	// Intact Flag
	intact := true

	// If the stem is less than 3 chars, there's nothing to do, so return
	if len(stem) < 3 {
		return string(stem)
	}

	// Main Control Loop
	cont := true
	for cont {
		// Lookup the map to see if a rule is available for the
		// given stems last letter
		rules, ok := r.Table[string(stem[len(stem)-1:])]
		if !ok {
			// Stop the loop if a matching rule is not found
			break
		}
		// Loop through the applicable rules
		for _, rule := range rules {

			// the length of the rule is greater than
			// the stem, so don't bother.
			if len(stem) <= len(rule.suf) {
				continue
			}

			// The rule does not match.
			if !strings.HasSuffix(string(stem), reverse(rule.suf)) {
				continue
			}

			// The stem is protected and should be left alone
			if rule.num == 0 {
				break
			}

			// The intact flag is set and the stem
			// has been operated on already.
			if rule.intact && !intact {
				continue
			}

			s := stem[:len(stem)-rule.num]
			// The result of the rule is invalid, so do nothing.
			if !validStem(string(s) + rule.app) {
				continue
			}

			// All criteria passed, the word should be stemmed
			cont = rule.cont
			current = []rune(string(s) + rule.app)

			// Set the intact flag
			intact = false

			// Break and repeat the process for the new stem
			break
		}

		// No rule matched
		if string(current) == string(stem) {
			break
		}

		// Set the new stem
		stem = current
	}
	return string(stem)
}

// Acceptability condition: if the stem begins with a vowel, then it
// must contain at least 2 letters, one of which must be a consonant
//
// If however, it begins with a consonant then it must contain three
// letters and at least one of these must be a vowel or 'y'
func validStem(word string) bool {
	runes := []rune(word)
	// If there's no vowel left in the stem, stem is invalid
	if !hasVowel(runes) {
		return false
	}

	// If the word has a vowel and is longer than 3 letters, stem is valid
	if len(runes) > 3 {
		return true
	}

	// If the first letter is a vowel
	if vowel(runes, 0) {
		if len(runes) > 1 {
			// The Second letter must be a consonant
			if consonant(runes, 1) {
				return true
			}
		} else {
			return false
		}

	} else {
		// If the first letter is a consonant
		// The stem must contain 3 letters, one of which we allready know
		// to be a vowel
		if len(runes) > 2 {
			return true
		}
	}
	return false
}

// consonant returns whether the letter at offset is a consonant
func consonant(word []rune, offset int) bool {
	switch word[offset] {
	case 'A', 'E', 'I', 'O', 'U', 'a', 'e', 'i', 'o', 'u':
		return false
	case 'Y', 'y':
		if offset == 0 {
			return true
		}
		return offset > 0 && !consonant(word, offset-1)
	}
	return true
}

// vowel returns whether the letter at offset is a vowel
func vowel(word []rune, offset int) bool {
	return !consonant(word, offset)
}

// hasVowel returns whether the word contains a vowel
func hasVowel(word []rune) bool {
	for i := 0; i < len(word); i++ {
		if vowel(word, i) {
			return true
		}
	}
	return false
}

// Reverses a string
func reverse(s string) string {
	runes := []rune(s)
	for i, j := 0, len(runes)-1; i < j; i, j = i+1, j-1 {
		runes[i], runes[j] = runes[j], runes[i]
	}
	return string(runes)
}

// Default Paice/Husk Rules
var defaultRules = `
ai*2.     { -ia > -   if intact }
a*1.      { -a > -    if intact }
bb1.      { -bb > -b   }
city3s.   { -ytic > -ys }
ci2>      { -ic > -    }
cn1t>     { -nc > -nt  }
dd1.      { -dd > -d   }
dei3y>    { -ied > -y  }
deec2ss.  { -ceed > -cess }
dee1.     { -eed > -ee }
de2>      { -ed > -    }
dooh4>    { -hood > -  }
e1>       { -e > -     }
feil1v.   { -lief > -liev }
fi2>      { -if > -    }
gni3>     { -ing > -   }
gai3y.    { -iag > -y  }
ga2>      { -ag > -    }
gg1.      { -gg > -g   }
ht*2.     { -th > -   if intact }
hsiug5ct. { -guish > -ct }
hsi3>     { -ish > -   }
i*1.      { -i > -    if intact }
i1y>      { -i > -y    }
ji1d.     { -ij > -id   --  see nois4j> & vis3j> }
juf1s.    { -fuj > -fus }
ju1d.     { -uj > -ud  }
jo1d.     { -oj > -od  }
jeh1r.    { -hej > -her }
jrev1t.   { -verj > -vert }
jsim2t.   { -misj > -mit }
jn1d.     { -nj > -nd  }
j1s.      { -j > -s    }
lbaifi6.  { -ifiabl > - }
lbai4y.   { -iabl > -y }
lba3>     { -abl > -   }
lbi3.     { -ibl > -   }
lib2l>    { -bil > -bl }
lc1.      { -cl > c    }
lufi4y.   { -iful > -y }
luf3>     { -ful > -   }
lu2.      { -ul > -    }
lai3>     { -ial > -   }
lau3>     { -ual > -   }
la2>      { -al > -    }
ll1.      { -ll > -l   }
mui3.     { -ium > -   }
mu*2.     { -um > -   if intact }
msi3>     { -ism > -   }
mm1.      { -mm > -m   }
nois4j>   { -sion > -j }
noix4ct.  { -xion > -ct }
noi3>     { -ion > -   }
nai3>     { -ian > -   }
na2>      { -an > -    }
nee0.     { protect  -een }
ne2>      { -en > -    }
nn1.      { -nn > -n   }
pihs4>    { -ship > -  }
pp1.      { -pp > -p   }
re2>      { -er > -    }
rae0.     { protect  -ear }
ra2.      { -ar > -    }
ro2>      { -or > -    }
ru2>      { -ur > -    }
rr1.      { -rr > -r   }
rt1>      { -tr > -t   }
rei3y>    { -ier > -y  }
sei3y>    { -ies > -y  }
sis2.     { -sis > -s  }
si2>      { -is > -    }
ssen4>    { -ness > -  }
ss0.      { protect  -ss }
suo3>     { -ous > -   }
su*2.     { -us > -   if intact }
s*1>      { -s > -    if intact }
s0.       { -s > -s    }
tacilp4y. { -plicat > -ply }
ta2>      { -at > -    }
tnem4>    { -ment > -  }
tne3>     { -ent > -   }
tna3>     { -ant > -   }
tpir2b.   { -ript > -rib }
tpro2b.   { -orpt > -orb }
tcud1.    { -duct > -duc }
tpmus2.   { -sumpt > -sum }
tpec2iv.  { -cept > -ceiv }
tulo2v.   { -olut > -olv }
tsis0.    { protect  -sist }
tsi3>     { -ist > -   }
tt1.      { -tt > -t   }
uqi3.     { -iqu > -   }
ugo1.     { -ogu > -og }
vis3j>    { -siv > -j  }
vie0.     { protect  -eiv }
vi2>      { -iv > -    }
ylb1>     { -bly > -bl }
yli3y>    { -ily > -y  }
ylp0.     { protect  -ply }
yl2>      { -ly > -    }
ygo1.     { -ogy > -og }
yhp1.     { -phy > -ph }
ymo1.     { -omy > -om }
ypo1.     { -opy > -op }
yti3>     { -ity > -   }
yte3>     { -ety > -   }
ytl2.     { -lty > -l  }
yrtsi5.   { -istry > - }
yra3>     { -ary > -   }
yro3>     { -ory > -   }
yfi3.     { -ify > -   }
ycn2t>    { -ncy > -nt }
yca3>     { -acy > -   }
zi2>      { -iz > -    }
zy1s.     { -yz > -ys  }
end0.
`
<commit_msg>Comment Change<commit_after>// Go implementation of the Paice/Husk Stemming algorithm:
// http://www.comp.lancs.ac.uk/computing/research/stemming/Links/paice.htm
package paicehusk

import (
	"regexp"
	"strconv"
	"strings"
)

// A representation of a stemming rule

// A representation of a stemming rule
type rule struct {

	// The suffix the rule is to act on
	suf string

	// True if the stem is required intact for the rule to operate
	intact bool

	// Number of letters to strip off the stem
	num int

	// A suffix to append to the stem
	app string

	// True if the stem should be stemmed further
	cont bool
}

// DefaultRules is a default ruleset for the english language.
var DefaultRules = NewRuleTable(strings.Split(defaultRules, "\n"))

// RuleTable stores rules based on the final letter of the suffix they
// act on allowing for easy lookup.
type RuleTable struct {
	Table map[string][]*rule
}

// NewRuleTable returns a new RuleTable instance
func NewRuleTable(f []string) (table *RuleTable) {
	table = &RuleTable{Table: make(map[string][]*rule)}
	for _, s := range f {
		if r, ok := ParseRule(s); ok {
			table.Table[r.suf[:1]] = append(table.Table[r.suf[:1]], r)
		}
	}
	return
}

// Regex for ValidRule
var reg = regexp.MustCompile("[a-zA-Z]*\\*?[0-9][a-zA-z]*[.>]")

// Validates a rule
func ValidRule(s string) (rule string, ok bool) {
	ok = true
	// Find the first instance of a rule in the provided string
	rule = reg.FindString(s)
	if rule == "" {
		ok = false
	}
	return
}

// Regexes for ParseRule
var suf = regexp.MustCompile("[a-zA-Z]+")
var intact = regexp.MustCompile("[*]")
var num = regexp.MustCompile("[0-9]")
var app = regexp.MustCompile("[0-9][a-zA-Z]+")

// ParseRule parses a rule in the form:
// |suffix|intact flag|number to strip|Append|Continue flag
//
// Eg, a rule: ht*2. Means if the stem is still intact, strip the
// suffix th and make no further attempts to stem the word.
//
// Rule nois4j> Means strip the sion suffix, append a j and check
// for any more rules to follow
func ParseRule(s string) (r *rule, ok bool) {
	s, ok = ValidRule(s)
	if !ok {
		return nil, false
	}

	r = new(rule)

	r.suf = suf.FindString(s)
	if intact.FindString(s) == "" {
		r.intact = false
	} else {
		r.intact = true
	}
	if i, err := strconv.ParseInt(num.FindString(s), 0, 0); err != nil {
		panic(err)
	} else {
		r.num = int(i)
	}
	if append := app.FindString(s); len(append) > 0 {
		r.app = app.FindString(s)[1:]
	} else {
		r.app = ""
	}

	if s[len(s)-1:] == ">" {
		r.cont = true
	} else {
		r.cont = false
	}
	return r, true
}

// Stem a string, word, based on the rules in *RuleTable r, by following
// the algorithm described at:
// http://www.comp.lancs.ac.uk/computing/research/stemming/Links/paice.htm
func (r *RuleTable) Stem(word string) string {
	stem := []rune(strings.ToLower(word))
	current := stem

	// Intact Flag
	intact := true

	// If the stem is less than 3 chars, there's nothing to do, so return
	if len(stem) < 3 {
		return string(stem)
	}

	// Main Control Loop
	cont := true
	for cont {
		// Lookup the map to see if a rule is available for the
		// given stems last letter
		rules, ok := r.Table[string(stem[len(stem)-1:])]
		if !ok {
			// Stop the loop if a matching rule is not found
			break
		}
		// Loop through the applicable rules
		for _, rule := range rules {

			// the length of the rule is greater than
			// the stem, so don't bother.
			if len(stem) <= len(rule.suf) {
				continue
			}

			// The rule does not match.
			if !strings.HasSuffix(string(stem), reverse(rule.suf)) {
				continue
			}

			// The stem is protected and should be left alone
			if rule.num == 0 {
				break
			}

			// The intact flag is set and the stem
			// has been operated on already.
			if rule.intact && !intact {
				continue
			}

			s := stem[:len(stem)-rule.num]
			// The result of the rule is invalid, so do nothing.
			if !validStem(string(s) + rule.app) {
				continue
			}

			// All criteria passed, the word should be stemmed
			cont = rule.cont
			current = []rune(string(s) + rule.app)

			// Set the intact flag
			intact = false

			// Break and repeat the process for the new stem
			break
		}

		// No rule matched
		if string(current) == string(stem) {
			break
		}

		// Set the new stem
		stem = current
	}
	return string(stem)
}

// Acceptability condition: if the stem begins with a vowel, then it
// must contain at least 2 letters, one of which must be a consonant
//
// If however, it begins with a consonant then it must contain three
// letters and at least one of these must be a vowel or 'y'
func validStem(word string) bool {
	runes := []rune(word)
	// If there's no vowel left in the stem, stem is invalid
	if !hasVowel(runes) {
		return false
	}

	// If the word has a vowel and is longer than 3 letters, stem is valid
	if len(runes) > 3 {
		return true
	}

	// If the first letter is a vowel
	if vowel(runes, 0) {
		if len(runes) > 1 {
			// The Second letter must be a consonant
			if consonant(runes, 1) {
				return true
			}
		} else {
			return false
		}

	} else {
		// If the first letter is a consonant
		// The stem must contain 3 letters, one of which we allready know
		// to be a vowel
		if len(runes) > 2 {
			return true
		}
	}
	return false
}

// consonant returns whether the letter at offset is a consonant
func consonant(word []rune, offset int) bool {
	switch word[offset] {
	case 'A', 'E', 'I', 'O', 'U', 'a', 'e', 'i', 'o', 'u':
		return false
	case 'Y', 'y':
		if offset == 0 {
			return true
		}
		return offset > 0 && !consonant(word, offset-1)
	}
	return true
}

// vowel returns whether the letter at offset is a vowel
func vowel(word []rune, offset int) bool {
	return !consonant(word, offset)
}

// hasVowel returns whether the word contains a vowel
func hasVowel(word []rune) bool {
	for i := 0; i < len(word); i++ {
		if vowel(word, i) {
			return true
		}
	}
	return false
}

// Reverses a string
func reverse(s string) string {
	runes := []rune(s)
	for i, j := 0, len(runes)-1; i < j; i, j = i+1, j-1 {
		runes[i], runes[j] = runes[j], runes[i]
	}
	return string(runes)
}

// Default Paice/Husk Rules
var defaultRules = `
ai*2.     { -ia > -   if intact }
a*1.      { -a > -    if intact }
bb1.      { -bb > -b   }
city3s.   { -ytic > -ys }
ci2>      { -ic > -    }
cn1t>     { -nc > -nt  }
dd1.      { -dd > -d   }
dei3y>    { -ied > -y  }
deec2ss.  { -ceed > -cess }
dee1.     { -eed > -ee }
de2>      { -ed > -    }
dooh4>    { -hood > -  }
e1>       { -e > -     }
feil1v.   { -lief > -liev }
fi2>      { -if > -    }
gni3>     { -ing > -   }
gai3y.    { -iag > -y  }
ga2>      { -ag > -    }
gg1.      { -gg > -g   }
ht*2.     { -th > -   if intact }
hsiug5ct. { -guish > -ct }
hsi3>     { -ish > -   }
i*1.      { -i > -    if intact }
i1y>      { -i > -y    }
ji1d.     { -ij > -id   --  see nois4j> & vis3j> }
juf1s.    { -fuj > -fus }
ju1d.     { -uj > -ud  }
jo1d.     { -oj > -od  }
jeh1r.    { -hej > -her }
jrev1t.   { -verj > -vert }
jsim2t.   { -misj > -mit }
jn1d.     { -nj > -nd  }
j1s.      { -j > -s    }
lbaifi6.  { -ifiabl > - }
lbai4y.   { -iabl > -y }
lba3>     { -abl > -   }
lbi3.     { -ibl > -   }
lib2l>    { -bil > -bl }
lc1.      { -cl > c    }
lufi4y.   { -iful > -y }
luf3>     { -ful > -   }
lu2.      { -ul > -    }
lai3>     { -ial > -   }
lau3>     { -ual > -   }
la2>      { -al > -    }
ll1.      { -ll > -l   }
mui3.     { -ium > -   }
mu*2.     { -um > -   if intact }
msi3>     { -ism > -   }
mm1.      { -mm > -m   }
nois4j>   { -sion > -j }
noix4ct.  { -xion > -ct }
noi3>     { -ion > -   }
nai3>     { -ian > -   }
na2>      { -an > -    }
nee0.     { protect  -een }
ne2>      { -en > -    }
nn1.      { -nn > -n   }
pihs4>    { -ship > -  }
pp1.      { -pp > -p   }
re2>      { -er > -    }
rae0.     { protect  -ear }
ra2.      { -ar > -    }
ro2>      { -or > -    }
ru2>      { -ur > -    }
rr1.      { -rr > -r   }
rt1>      { -tr > -t   }
rei3y>    { -ier > -y  }
sei3y>    { -ies > -y  }
sis2.     { -sis > -s  }
si2>      { -is > -    }
ssen4>    { -ness > -  }
ss0.      { protect  -ss }
suo3>     { -ous > -   }
su*2.     { -us > -   if intact }
s*1>      { -s > -    if intact }
s0.       { -s > -s    }
tacilp4y. { -plicat > -ply }
ta2>      { -at > -    }
tnem4>    { -ment > -  }
tne3>     { -ent > -   }
tna3>     { -ant > -   }
tpir2b.   { -ript > -rib }
tpro2b.   { -orpt > -orb }
tcud1.    { -duct > -duc }
tpmus2.   { -sumpt > -sum }
tpec2iv.  { -cept > -ceiv }
tulo2v.   { -olut > -olv }
tsis0.    { protect  -sist }
tsi3>     { -ist > -   }
tt1.      { -tt > -t   }
uqi3.     { -iqu > -   }
ugo1.     { -ogu > -og }
vis3j>    { -siv > -j  }
vie0.     { protect  -eiv }
vi2>      { -iv > -    }
ylb1>     { -bly > -bl }
yli3y>    { -ily > -y  }
ylp0.     { protect  -ply }
yl2>      { -ly > -    }
ygo1.     { -ogy > -og }
yhp1.     { -phy > -ph }
ymo1.     { -omy > -om }
ypo1.     { -opy > -op }
yti3>     { -ity > -   }
yte3>     { -ety > -   }
ytl2.     { -lty > -l  }
yrtsi5.   { -istry > - }
yra3>     { -ary > -   }
yro3>     { -ory > -   }
yfi3.     { -ify > -   }
ycn2t>    { -ncy > -nt }
yca3>     { -acy > -   }
zi2>      { -iz > -    }
zy1s.     { -yz > -ys  }
end0.
`
<|endoftext|>


<commit_before>package pagerduty_test

import (
	"fmt"
	"net/http"
	"reflect"
	"testing"

	. "github.com/danryan/go-pagerduty/pagerduty"
)

func TestSchedule_marshal(t *testing.T) {
	testJSONMarshal(t, &Schedule{}, "{}")

	s := &Schedule{
		ID:                   "SCHED1",
		Name:                 "Primary",
		Timezone:             "Pacific Time (US & Canada)",
		Today:                "2014-12-27",
		EscalationPolicies:   []*EscalationPolicy{(&EscalationPolicy{ID: "ESCPOL1", Name: "Escalation Policy 1"})},
		ScheduleLayers:       []*ScheduleLayer{(&ScheduleLayer{Name: "Schedule Layer 1"})},
		OverridesSubschedule: &ScheduleLayer{Name: "Override Schedule Layer"},
		FinalSchedule:        &ScheduleLayer{Name: "Final Schedule Layer"},
	}

	want := `{
		"id": "SCHED1",
		"name": "Primary",
                "time_zone": "Pacific Time (US & Canada)",
                "today": "2014-12-27",
                "escalation_policies": [{"id": "ESCPOL1", "name": "Escalation Policy 1"}],
                "schedule_layers": [{"name": "Schedule Layer 1"}],
                "overrides_subschedule": {"name": "Override Schedule Layer"},
                "final_schedule": {"name": "Final Schedule Layer"}
	}`

	testJSONMarshal(t, s, want)
}

func TestSchedulesService_List(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/schedules", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		fmt.Fprint(w, `{"schedules": [{"id": "SCHED1", "name": "Schedule 1"}], "limit": 100, "offset": 0, "total": 1}`)
	})

	schedules, _, err := client.Schedules.List(&SchedulesOptions{})
	if err != nil {
		t.Errorf("Schedules.List returned error: %v", err)
	}

	want := &Schedules{Schedules: []*Schedule{(&Schedule{ID: "SCHED1", Name: "Schedule 1"})}, Limit: 100, Offset: 0, Total: 1}
	if !reflect.DeepEqual(schedules, want) {
		t.Errorf("Schedules.List returned %+v, want %%+v", schedules, want)
	}
}

func TestSchedulesService_Get(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/schedules/s", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		fmt.Fprint(w, `{"schedule": {"id": "ABCDEF"}}`)
	})

	schedule, _, err := client.Schedules.Get("s")
	if err != nil {
		t.Errorf("Schedules.Get returned error: %v", err)
	}

	want := &Schedule{ID: "ABCDEF"}
	if !reflect.DeepEqual(schedule, want) {
		t.Errorf("Schedules.Get returned %+v, want %%+v", schedule, want)
	}
}

func TestSchedulesService_Entries(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/schedules/s/entries", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		fmt.Fprint(w, `{"total": 1, "entries": [{"user": {"id": "USR1", "name": "Test User"}, "end": "2012-08-19T12:00:00-04:00", "start": "2012-08-19T00:00:00-04:00"}]}`)
	})

	entries, _, err := client.Schedules.Entries("s", &ScheduleEntriesOptions{})
	if err != nil {
		t.Errorf("Schedules.Entries returned error: %v", err)
	}

	want := &ScheduleEntries{ScheduleEntries: []*ScheduleEntry{(&ScheduleEntry{User: &User{ID: "USR1", Name: "Test User"}, Start: "2012-08-19T00:00:00-04:00", End: "2012-08-19T12:00:00-04:00"})}, Total: 1}
	if !reflect.DeepEqual(entries, want) {
		t.Errorf("Schedules.Entries returned %+v, want %%+v", entries, want)
	}
}
<commit_msg>tiny test fix (silly unicode)<commit_after>package pagerduty_test

import (
	"fmt"
	"net/http"
	"reflect"
	"testing"

	. "github.com/danryan/go-pagerduty/pagerduty"
)

func TestSchedule_marshal(t *testing.T) {
	testJSONMarshal(t, &Schedule{}, "{}")

	s := &Schedule{
		ID:                   "SCHED1",
		Name:                 "Primary",
		Timezone:             "Pacific Time (US & Canada)",
		Today:                "2014-12-27",
		EscalationPolicies:   []*EscalationPolicy{(&EscalationPolicy{ID: "ESCPOL1", Name: "Escalation Policy 1"})},
		ScheduleLayers:       []*ScheduleLayer{(&ScheduleLayer{Name: "Schedule Layer 1"})},
		OverridesSubschedule: &ScheduleLayer{Name: "Override Schedule Layer"},
		FinalSchedule:        &ScheduleLayer{Name: "Final Schedule Layer"},
	}

	want := `{
		"id": "SCHED1",
		"name": "Primary",
		"time_zone": "Pacific Time (US \u0026 Canada)",
		"today": "2014-12-27",
		"escalation_policies": [{"id": "ESCPOL1", "name": "Escalation Policy 1"}],
		"schedule_layers": [{"name": "Schedule Layer 1"}],
		"overrides_subschedule": {"name": "Override Schedule Layer"},
		"final_schedule": {"name": "Final Schedule Layer"}
	}`

	testJSONMarshal(t, s, want)
}

func TestSchedulesService_List(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/schedules", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		fmt.Fprint(w, `{"schedules": [{"id": "SCHED1", "name": "Schedule 1"}], "limit": 100, "offset": 0, "total": 1}`)
	})

	schedules, _, err := client.Schedules.List(&SchedulesOptions{})
	if err != nil {
		t.Errorf("Schedules.List returned error: %v", err)
	}

	want := &Schedules{Schedules: []*Schedule{(&Schedule{ID: "SCHED1", Name: "Schedule 1"})}, Limit: 100, Offset: 0, Total: 1}
	if !reflect.DeepEqual(schedules, want) {
		t.Errorf("Schedules.List returned %+v, want %%+v", schedules, want)
	}
}

func TestSchedulesService_Get(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/schedules/s", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		fmt.Fprint(w, `{"schedule": {"id": "ABCDEF"}}`)
	})

	schedule, _, err := client.Schedules.Get("s")
	if err != nil {
		t.Errorf("Schedules.Get returned error: %v", err)
	}

	want := &Schedule{ID: "ABCDEF"}
	if !reflect.DeepEqual(schedule, want) {
		t.Errorf("Schedules.Get returned %+v, want %%+v", schedule, want)
	}
}

func TestSchedulesService_Entries(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/schedules/s/entries", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		fmt.Fprint(w, `{"total": 1, "entries": [{"user": {"id": "USR1", "name": "Test User"}, "end": "2012-08-19T12:00:00-04:00", "start": "2012-08-19T00:00:00-04:00"}]}`)
	})

	entries, _, err := client.Schedules.Entries("s", &ScheduleEntriesOptions{})
	if err != nil {
		t.Errorf("Schedules.Entries returned error: %v", err)
	}

	want := &ScheduleEntries{ScheduleEntries: []*ScheduleEntry{(&ScheduleEntry{User: &User{ID: "USR1", Name: "Test User"}, Start: "2012-08-19T00:00:00-04:00", End: "2012-08-19T12:00:00-04:00"})}, Total: 1}
	if !reflect.DeepEqual(entries, want) {
		t.Errorf("Schedules.Entries returned %+v, want %%+v", entries, want)
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"flag"
	"fmt"
	"io"
	"os"
	"regexp"
	"runtime"
	"testing"
	"time"
)

func newLocalListener(t *testing.T) Listener {
	ln, err := Listen("tcp", "127.0.0.1:0")
	if err != nil {
		ln, err = Listen("tcp6", "[::1]:0")
	}
	if err != nil {
		t.Fatal(err)
	}
	return ln
}

func TestDialTimeout(t *testing.T) {
	ln := newLocalListener(t)
	defer ln.Close()

	errc := make(chan error)

	numConns := listenerBacklog + 10

	// TODO(bradfitz): It's hard to test this in a portable
	// way. This is unfortunate, but works for now.
	switch runtime.GOOS {
	case "linux":
		// The kernel will start accepting TCP connections before userspace
		// gets a chance to not accept them, so fire off a bunch to fill up
		// the kernel's backlog.  Then we test we get a failure after that.
		for i := 0; i < numConns; i++ {
			go func() {
				_, err := DialTimeout("tcp", ln.Addr().String(), 200*time.Millisecond)
				errc <- err
			}()
		}
	case "darwin", "windows":
		// At least OS X 10.7 seems to accept any number of
		// connections, ignoring listen's backlog, so resort
		// to connecting to a hopefully-dead 127/8 address.
		// Same for windows.
		//
		// Use an IANA reserved port (49151) instead of 80, because
		// on our 386 builder, this Dial succeeds, connecting
		// to an IIS web server somewhere.  The data center
		// or VM or firewall must be stealing the TCP connection.
		//
		// IANA Service Name and Transport Protocol Port Number Registry
		// <http://www.iana.org/assignments/service-names-port-numbers/service-names-port-numbers.xml>
		go func() {
			c, err := DialTimeout("tcp", "127.0.71.111:49151", 200*time.Millisecond)
			if err == nil {
				err = fmt.Errorf("unexpected: connected to %s!", c.RemoteAddr())
				c.Close()
			}
			errc <- err
		}()
	default:
		// TODO(bradfitz):
		// OpenBSD may have a reject route to 127/8 except 127.0.0.1/32
		// by default. FreeBSD likely works, but is untested.
		// TODO(rsc):
		// The timeout never happens on Windows.  Why?  Issue 3016.
		t.Logf("skipping test on %q; untested.", runtime.GOOS)
		return
	}

	connected := 0
	for {
		select {
		case <-time.After(15 * time.Second):
			t.Fatal("too slow")
		case err := <-errc:
			if err == nil {
				connected++
				if connected == numConns {
					t.Fatal("all connections connected; expected some to time out")
				}
			} else {
				terr, ok := err.(timeout)
				if !ok {
					t.Fatalf("got error %q; want error with timeout interface", err)
				}
				if !terr.Timeout() {
					t.Fatalf("got error %q; not a timeout", err)
				}
				// Pass. We saw a timeout error.
				return
			}
		}
	}
}

func TestSelfConnect(t *testing.T) {
	if runtime.GOOS == "windows" {
		// TODO(brainman): do not know why it hangs.
		t.Logf("skipping known-broken test on windows")
		return
	}
	// Test that Dial does not honor self-connects.
	// See the comment in DialTCP.

	// Find a port that would be used as a local address.
	l, err := Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	c, err := Dial("tcp", l.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	addr := c.LocalAddr().String()
	c.Close()
	l.Close()

	// Try to connect to that address repeatedly.
	n := 100000
	if testing.Short() {
		n = 1000
	}
	switch runtime.GOOS {
	case "darwin", "freebsd", "netbsd", "openbsd", "plan9", "windows":
		// Non-Linux systems take a long time to figure
		// out that there is nothing listening on localhost.
		n = 100
	}
	for i := 0; i < n; i++ {
		c, err := Dial("tcp", addr)
		if err == nil {
			c.Close()
			t.Errorf("#%d: Dial %q succeeded", i, addr)
		}
	}
}

var runErrorTest = flag.Bool("run_error_test", false, "let TestDialError check for dns errors")

type DialErrorTest struct {
	Net     string
	Raddr   string
	Pattern string
}

var dialErrorTests = []DialErrorTest{
	{
		"datakit", "mh/astro/r70",
		"dial datakit mh/astro/r70: unknown network datakit",
	},
	{
		"tcp", "127.0.0.1:☺",
		"dial tcp 127.0.0.1:☺: unknown port tcp/☺",
	},
	{
		"tcp", "no-such-name.google.com.:80",
		"dial tcp no-such-name.google.com.:80: lookup no-such-name.google.com.( on .*)?: no (.*)",
	},
	{
		"tcp", "no-such-name.no-such-top-level-domain.:80",
		"dial tcp no-such-name.no-such-top-level-domain.:80: lookup no-such-name.no-such-top-level-domain.( on .*)?: no (.*)",
	},
	{
		"tcp", "no-such-name:80",
		`dial tcp no-such-name:80: lookup no-such-name\.(.*\.)?( on .*)?: no (.*)`,
	},
	{
		"tcp", "mh/astro/r70:http",
		"dial tcp mh/astro/r70:http: lookup mh/astro/r70: invalid domain name",
	},
	{
		"unix", "/etc/file-not-found",
		"dial unix /etc/file-not-found: no such file or directory",
	},
	{
		"unix", "/etc/",
		"dial unix /etc/: (permission denied|socket operation on non-socket|connection refused)",
	},
	{
		"unixpacket", "/etc/file-not-found",
		"dial unixpacket /etc/file-not-found: no such file or directory",
	},
	{
		"unixpacket", "/etc/",
		"dial unixpacket /etc/: (permission denied|socket operation on non-socket|connection refused)",
	},
}

var duplicateErrorPattern = `dial (.*) dial (.*)`

func TestDialError(t *testing.T) {
	if !*runErrorTest {
		t.Logf("test disabled; use -run_error_test to enable")
		return
	}
	for i, tt := range dialErrorTests {
		c, err := Dial(tt.Net, tt.Raddr)
		if c != nil {
			c.Close()
		}
		if err == nil {
			t.Errorf("#%d: nil error, want match for %#q", i, tt.Pattern)
			continue
		}
		s := err.Error()
		match, _ := regexp.MatchString(tt.Pattern, s)
		if !match {
			t.Errorf("#%d: %q, want match for %#q", i, s, tt.Pattern)
		}
		match, _ = regexp.MatchString(duplicateErrorPattern, s)
		if match {
			t.Errorf("#%d: %q, duplicate error return from Dial", i, s)
		}
	}
}

func TestDialTimeoutFDLeak(t *testing.T) {
	if runtime.GOOS != "linux" {
		// TODO(bradfitz): test on other platforms
		t.Logf("skipping test on %s", runtime.GOOS)
		return
	}

	ln := newLocalListener(t)
	defer ln.Close()

	type connErr struct {
		conn Conn
		err  error
	}
	dials := listenerBacklog + 100
	maxGoodConnect := listenerBacklog + 5 // empirically 131 good ones (of 128). who knows?
	resc := make(chan connErr)
	for i := 0; i < dials; i++ {
		go func() {
			conn, err := DialTimeout("tcp", ln.Addr().String(), 500*time.Millisecond)
			resc <- connErr{conn, err}
		}()
	}

	var firstErr string
	var ngood int
	var toClose []io.Closer
	for i := 0; i < dials; i++ {
		ce := <-resc
		if ce.err == nil {
			ngood++
			if ngood > maxGoodConnect {
				t.Errorf("%d good connects; expected at most %d", ngood, maxGoodConnect)
			}
			toClose = append(toClose, ce.conn)
			continue
		}
		err := ce.err
		if firstErr == "" {
			firstErr = err.Error()
		} else if err.Error() != firstErr {
			t.Fatalf("inconsistent error messages: first was %q, then later %q", firstErr, err)
		}
	}
	for _, c := range toClose {
		c.Close()
	}
	for i := 0; i < 100; i++ {
		if got := numFD(); got < dials {
			// Test passes.
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	if got := numFD(); got >= dials {
		t.Errorf("num fds after %d timeouts = %d; want <%d", dials, got, dials)
	}
}

func numFD() int {
	if runtime.GOOS == "linux" {
		f, err := os.Open("/proc/self/fd")
		if err != nil {
			panic(err)
		}
		defer f.Close()
		names, err := f.Readdirnames(0)
		if err != nil {
			panic(err)
		}
		return len(names)
	}
	// All tests using this should be skipped anyway, but:
	panic("numFDs not implemented on " + runtime.GOOS)
}
<commit_msg>net: TestDialTimeoutFDLeak failure<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"flag"
	"fmt"
	"io"
	"os"
	"regexp"
	"runtime"
	"testing"
	"time"
)

func newLocalListener(t *testing.T) Listener {
	ln, err := Listen("tcp", "127.0.0.1:0")
	if err != nil {
		ln, err = Listen("tcp6", "[::1]:0")
	}
	if err != nil {
		t.Fatal(err)
	}
	return ln
}

func TestDialTimeout(t *testing.T) {
	ln := newLocalListener(t)
	defer ln.Close()

	errc := make(chan error)

	numConns := listenerBacklog + 10

	// TODO(bradfitz): It's hard to test this in a portable
	// way. This is unfortunate, but works for now.
	switch runtime.GOOS {
	case "linux":
		// The kernel will start accepting TCP connections before userspace
		// gets a chance to not accept them, so fire off a bunch to fill up
		// the kernel's backlog.  Then we test we get a failure after that.
		for i := 0; i < numConns; i++ {
			go func() {
				_, err := DialTimeout("tcp", ln.Addr().String(), 200*time.Millisecond)
				errc <- err
			}()
		}
	case "darwin", "windows":
		// At least OS X 10.7 seems to accept any number of
		// connections, ignoring listen's backlog, so resort
		// to connecting to a hopefully-dead 127/8 address.
		// Same for windows.
		//
		// Use an IANA reserved port (49151) instead of 80, because
		// on our 386 builder, this Dial succeeds, connecting
		// to an IIS web server somewhere.  The data center
		// or VM or firewall must be stealing the TCP connection.
		//
		// IANA Service Name and Transport Protocol Port Number Registry
		// <http://www.iana.org/assignments/service-names-port-numbers/service-names-port-numbers.xml>
		go func() {
			c, err := DialTimeout("tcp", "127.0.71.111:49151", 200*time.Millisecond)
			if err == nil {
				err = fmt.Errorf("unexpected: connected to %s!", c.RemoteAddr())
				c.Close()
			}
			errc <- err
		}()
	default:
		// TODO(bradfitz):
		// OpenBSD may have a reject route to 127/8 except 127.0.0.1/32
		// by default. FreeBSD likely works, but is untested.
		// TODO(rsc):
		// The timeout never happens on Windows.  Why?  Issue 3016.
		t.Logf("skipping test on %q; untested.", runtime.GOOS)
		return
	}

	connected := 0
	for {
		select {
		case <-time.After(15 * time.Second):
			t.Fatal("too slow")
		case err := <-errc:
			if err == nil {
				connected++
				if connected == numConns {
					t.Fatal("all connections connected; expected some to time out")
				}
			} else {
				terr, ok := err.(timeout)
				if !ok {
					t.Fatalf("got error %q; want error with timeout interface", err)
				}
				if !terr.Timeout() {
					t.Fatalf("got error %q; not a timeout", err)
				}
				// Pass. We saw a timeout error.
				return
			}
		}
	}
}

func TestSelfConnect(t *testing.T) {
	if runtime.GOOS == "windows" {
		// TODO(brainman): do not know why it hangs.
		t.Logf("skipping known-broken test on windows")
		return
	}
	// Test that Dial does not honor self-connects.
	// See the comment in DialTCP.

	// Find a port that would be used as a local address.
	l, err := Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	c, err := Dial("tcp", l.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	addr := c.LocalAddr().String()
	c.Close()
	l.Close()

	// Try to connect to that address repeatedly.
	n := 100000
	if testing.Short() {
		n = 1000
	}
	switch runtime.GOOS {
	case "darwin", "freebsd", "netbsd", "openbsd", "plan9", "windows":
		// Non-Linux systems take a long time to figure
		// out that there is nothing listening on localhost.
		n = 100
	}
	for i := 0; i < n; i++ {
		c, err := Dial("tcp", addr)
		if err == nil {
			c.Close()
			t.Errorf("#%d: Dial %q succeeded", i, addr)
		}
	}
}

var runErrorTest = flag.Bool("run_error_test", false, "let TestDialError check for dns errors")

type DialErrorTest struct {
	Net     string
	Raddr   string
	Pattern string
}

var dialErrorTests = []DialErrorTest{
	{
		"datakit", "mh/astro/r70",
		"dial datakit mh/astro/r70: unknown network datakit",
	},
	{
		"tcp", "127.0.0.1:☺",
		"dial tcp 127.0.0.1:☺: unknown port tcp/☺",
	},
	{
		"tcp", "no-such-name.google.com.:80",
		"dial tcp no-such-name.google.com.:80: lookup no-such-name.google.com.( on .*)?: no (.*)",
	},
	{
		"tcp", "no-such-name.no-such-top-level-domain.:80",
		"dial tcp no-such-name.no-such-top-level-domain.:80: lookup no-such-name.no-such-top-level-domain.( on .*)?: no (.*)",
	},
	{
		"tcp", "no-such-name:80",
		`dial tcp no-such-name:80: lookup no-such-name\.(.*\.)?( on .*)?: no (.*)`,
	},
	{
		"tcp", "mh/astro/r70:http",
		"dial tcp mh/astro/r70:http: lookup mh/astro/r70: invalid domain name",
	},
	{
		"unix", "/etc/file-not-found",
		"dial unix /etc/file-not-found: no such file or directory",
	},
	{
		"unix", "/etc/",
		"dial unix /etc/: (permission denied|socket operation on non-socket|connection refused)",
	},
	{
		"unixpacket", "/etc/file-not-found",
		"dial unixpacket /etc/file-not-found: no such file or directory",
	},
	{
		"unixpacket", "/etc/",
		"dial unixpacket /etc/: (permission denied|socket operation on non-socket|connection refused)",
	},
}

var duplicateErrorPattern = `dial (.*) dial (.*)`

func TestDialError(t *testing.T) {
	if !*runErrorTest {
		t.Logf("test disabled; use -run_error_test to enable")
		return
	}
	for i, tt := range dialErrorTests {
		c, err := Dial(tt.Net, tt.Raddr)
		if c != nil {
			c.Close()
		}
		if err == nil {
			t.Errorf("#%d: nil error, want match for %#q", i, tt.Pattern)
			continue
		}
		s := err.Error()
		match, _ := regexp.MatchString(tt.Pattern, s)
		if !match {
			t.Errorf("#%d: %q, want match for %#q", i, s, tt.Pattern)
		}
		match, _ = regexp.MatchString(duplicateErrorPattern, s)
		if match {
			t.Errorf("#%d: %q, duplicate error return from Dial", i, s)
		}
	}
}

func TestDialTimeoutFDLeak(t *testing.T) {
	if runtime.GOOS != "linux" {
		// TODO(bradfitz): test on other platforms
		t.Logf("skipping test on %s", runtime.GOOS)
		return
	}

	ln := newLocalListener(t)
	defer ln.Close()

	type connErr struct {
		conn Conn
		err  error
	}
	dials := listenerBacklog + 100
	// used to be listenerBacklog + 5, but was found to be unreliable, issue 4384.
	maxGoodConnect := 150
	resc := make(chan connErr)
	for i := 0; i < dials; i++ {
		go func() {
			conn, err := DialTimeout("tcp", ln.Addr().String(), 500*time.Millisecond)
			resc <- connErr{conn, err}
		}()
	}

	var firstErr string
	var ngood int
	var toClose []io.Closer
	for i := 0; i < dials; i++ {
		ce := <-resc
		if ce.err == nil {
			ngood++
			if ngood > maxGoodConnect {
				t.Errorf("%d good connects; expected at most %d", ngood, maxGoodConnect)
			}
			toClose = append(toClose, ce.conn)
			continue
		}
		err := ce.err
		if firstErr == "" {
			firstErr = err.Error()
		} else if err.Error() != firstErr {
			t.Fatalf("inconsistent error messages: first was %q, then later %q", firstErr, err)
		}
	}
	for _, c := range toClose {
		c.Close()
	}
	for i := 0; i < 100; i++ {
		if got := numFD(); got < dials {
			// Test passes.
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	if got := numFD(); got >= dials {
		t.Errorf("num fds after %d timeouts = %d; want <%d", dials, got, dials)
	}
}

func numFD() int {
	if runtime.GOOS == "linux" {
		f, err := os.Open("/proc/self/fd")
		if err != nil {
			panic(err)
		}
		defer f.Close()
		names, err := f.Readdirnames(0)
		if err != nil {
			panic(err)
		}
		return len(names)
	}
	// All tests using this should be skipped anyway, but:
	panic("numFDs not implemented on " + runtime.GOOS)
}
<|endoftext|>