<commit_before>package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"path"
	"path/filepath"
	"time"

	"golang.org/x/sync/errgroup"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/gogo/protobuf/proto"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/client/version"
	"github.com/pachyderm/pachyderm/src/server/pkg/cmdutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/worker"
	ppsserver "github.com/pachyderm/pachyderm/src/server/pps/server"
	"google.golang.org/grpc"
)

// appEnv stores the environment variables that this worker needs
type appEnv struct {
	// Address of etcd, so that worker can write its own IP there for discoverh
	EtcdAddress string `env:"ETCD_PORT_2379_TCP_ADDR,required"`

	// Address for connecting to pachd (so this can download input data)
	PachdAddress string `env:"PACHD_PORT_650_TCP_ADDR,required"`

	// Prefix in etcd for all pachd-related records
	PPSPrefix string `env:"PPS_ETCD_PREFIX,required"`

	// worker gets its own IP here, via the k8s downward API. It then writes that
	// IP back to etcd so that pachd can discover it
	PPSWorkerIP string `env:"PPS_WORKER_IP,required"`

	// Either pipeline name or job name must be set
	PPSPipelineName string `env:"PPS_PIPELINE_NAME"`
	PPSJobID        string `env:"PPS_JOB_ID"`
	PodName         string `env:"PPS_POD_NAME,required"`
}

func main() {
	cmdutil.Main(do, &appEnv{})
}

func validateEnv(appEnv *appEnv) error {
	if appEnv.PPSPipelineName == "" && appEnv.PPSJobID == "" {
		return fmt.Errorf("worker must recieve either pipeline name or job ID, but got neither")
	} else if appEnv.PPSPipelineName != "" && appEnv.PPSJobID != "" {
		return fmt.Errorf("worker must recieve either pipeline name or job ID, but got both")
	}
	return nil
}

// getPipelineInfo gets the PipelineInfo proto describing the pipeline that this
// worker is part of
func getPipelineInfo(etcdClient *etcd.Client, appEnv *appEnv) (*pps.PipelineInfo, error) {
	ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
	resp, err := etcdClient.Get(ctx, path.Join(appEnv.PPSPrefix, "pipelines", appEnv.PPSPipelineName))
	if err != nil {
		return nil, err
	}
	if len(resp.Kvs) != 1 {
		return nil, fmt.Errorf("expected to find 1 pipeline, got %d: %v", len(resp.Kvs), resp)
	}
	pipelineInfo := new(pps.PipelineInfo)
	if err := proto.UnmarshalText(string(resp.Kvs[0].Value), pipelineInfo); err != nil {
		return nil, err
	}
	return pipelineInfo, nil
}

func getJobInfo(etcdClient *etcd.Client, appEnv *appEnv) (*pps.JobInfo, error) {
	ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
	resp, err := etcdClient.Get(ctx, path.Join(appEnv.PPSPrefix, "jobs", appEnv.PPSJobID))
	if err != nil {
		return nil, err
	}
	if len(resp.Kvs) != 1 {
		return nil, fmt.Errorf("expected to find 1 job, got %d: %v", len(resp.Kvs), resp)
	}
	jobInfo := new(pps.JobInfo)
	if err := proto.UnmarshalText(string(resp.Kvs[0].Value), jobInfo); err != nil {
		return nil, err
	}
	return jobInfo, nil
}

func do(appEnvObj interface{}) error {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	appEnv := appEnvObj.(*appEnv)
	if err := validateEnv(appEnv); err != nil {
		return err
	}

	// get pachd client, so we can upload output data from the user binary
	pachClient, err := client.NewFromAddress(fmt.Sprintf("%v:650", appEnv.PachdAddress))
	if err != nil {
		return err
	}
	go pachClient.KeepConnected(make(chan bool)) // we never cancel the connection

	// Get etcd client, so we can register our IP (so pachd can discover us)
	etcdClient, err := etcd.New(etcd.Config{
		Endpoints:   []string{fmt.Sprintf("%s:2379", appEnv.EtcdAddress)},
		DialTimeout: 15 * time.Second,
	})
	if err != nil {
		return err
	}

	// Construct worker API server. Get relevant pipeline or job info, and then
	// use that to create a worker.APIServer.
	var workerRcName string
	var apiServer *worker.APIServer
	if appEnv.PPSPipelineName != "" {
		pipelineInfo, err := getPipelineInfo(etcdClient, appEnv)
		if err != nil {
			return err
		}
		workerRcName = ppsserver.PipelineRcName(pipelineInfo.Pipeline.Name, pipelineInfo.Version)
		apiServer = worker.NewPipelineAPIServer(pachClient, pipelineInfo, appEnv.PodName)
	} else if appEnv.PPSJobID != "" {
		jobInfo, err := getJobInfo(etcdClient, appEnv)
		if err != nil {
			return err
		}
		workerRcName = ppsserver.JobRcName(jobInfo.Job.ID)
		apiServer = worker.NewJobAPIServer(pachClient, jobInfo, appEnv.PodName)
	}

	// Setup the hostPath mount to use a unique directory for this worker
	workerDir := filepath.Join(client.PPSHostPath, appEnv.PodName)
	if err := os.MkdirAll(workerDir, 0777); err != nil {
		return err
	}
	if err := os.Symlink(workerDir, client.PPSInputPrefix); err != nil {
		return err
	}

	// Start worker api server
	eg := errgroup.Group{}
	ready := make(chan error)
	eg.Go(func() error {
		return grpcutil.Serve(
			func(s *grpc.Server) {
				worker.RegisterWorkerServer(s, apiServer)
				close(ready)
			},
			grpcutil.ServeOptions{
				Version:    version.Version,
				MaxMsgSize: grpcutil.MaxMsgSize,
			},
			grpcutil.ServeEnv{
				GRPCPort: client.PPSWorkerPort,
			},
		)
	})

	// Wait until server is ready, then put our IP address into etcd, so pachd can
	// discover us
	<-ready
	key := path.Join(appEnv.PPSPrefix, "workers", workerRcName, appEnv.PPSWorkerIP)

	// Prepare to write "key" into etcd by creating lease -- if worker dies, our
	// IP will be removed from etcd
	ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
	resp, err := etcdClient.Grant(ctx, 60 /* seconds */)
	if err != nil {
		return err
	}
	etcdClient.KeepAlive(context.Background(), resp.ID) // keepalive forever

	// Actually write "key" into etcd
	ctx, _ = context.WithTimeout(context.Background(), 30*time.Second) // new ctx
	if _, err := etcdClient.Put(ctx, key, "", etcd.WithLease(resp.ID)); err != nil {
		return err
	}

	// If server ever exits, return error
	if err := eg.Wait(); err != nil {
		return err
	}
	return nil
}
<commit_msg>Lower worker's etcd lease value so they expire faster once they are scaled down<commit_after>package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"path"
	"path/filepath"
	"time"

	"golang.org/x/sync/errgroup"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/gogo/protobuf/proto"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/client/version"
	"github.com/pachyderm/pachyderm/src/server/pkg/cmdutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/worker"
	ppsserver "github.com/pachyderm/pachyderm/src/server/pps/server"
	"google.golang.org/grpc"
)

// appEnv stores the environment variables that this worker needs
type appEnv struct {
	// Address of etcd, so that worker can write its own IP there for discoverh
	EtcdAddress string `env:"ETCD_PORT_2379_TCP_ADDR,required"`

	// Address for connecting to pachd (so this can download input data)
	PachdAddress string `env:"PACHD_PORT_650_TCP_ADDR,required"`

	// Prefix in etcd for all pachd-related records
	PPSPrefix string `env:"PPS_ETCD_PREFIX,required"`

	// worker gets its own IP here, via the k8s downward API. It then writes that
	// IP back to etcd so that pachd can discover it
	PPSWorkerIP string `env:"PPS_WORKER_IP,required"`

	// Either pipeline name or job name must be set
	PPSPipelineName string `env:"PPS_PIPELINE_NAME"`
	PPSJobID        string `env:"PPS_JOB_ID"`
	PodName         string `env:"PPS_POD_NAME,required"`
}

func main() {
	cmdutil.Main(do, &appEnv{})
}

func validateEnv(appEnv *appEnv) error {
	if appEnv.PPSPipelineName == "" && appEnv.PPSJobID == "" {
		return fmt.Errorf("worker must recieve either pipeline name or job ID, but got neither")
	} else if appEnv.PPSPipelineName != "" && appEnv.PPSJobID != "" {
		return fmt.Errorf("worker must recieve either pipeline name or job ID, but got both")
	}
	return nil
}

// getPipelineInfo gets the PipelineInfo proto describing the pipeline that this
// worker is part of
func getPipelineInfo(etcdClient *etcd.Client, appEnv *appEnv) (*pps.PipelineInfo, error) {
	ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
	resp, err := etcdClient.Get(ctx, path.Join(appEnv.PPSPrefix, "pipelines", appEnv.PPSPipelineName))
	if err != nil {
		return nil, err
	}
	if len(resp.Kvs) != 1 {
		return nil, fmt.Errorf("expected to find 1 pipeline, got %d: %v", len(resp.Kvs), resp)
	}
	pipelineInfo := new(pps.PipelineInfo)
	if err := proto.UnmarshalText(string(resp.Kvs[0].Value), pipelineInfo); err != nil {
		return nil, err
	}
	return pipelineInfo, nil
}

func getJobInfo(etcdClient *etcd.Client, appEnv *appEnv) (*pps.JobInfo, error) {
	ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
	resp, err := etcdClient.Get(ctx, path.Join(appEnv.PPSPrefix, "jobs", appEnv.PPSJobID))
	if err != nil {
		return nil, err
	}
	if len(resp.Kvs) != 1 {
		return nil, fmt.Errorf("expected to find 1 job, got %d: %v", len(resp.Kvs), resp)
	}
	jobInfo := new(pps.JobInfo)
	if err := proto.UnmarshalText(string(resp.Kvs[0].Value), jobInfo); err != nil {
		return nil, err
	}
	return jobInfo, nil
}

func do(appEnvObj interface{}) error {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	appEnv := appEnvObj.(*appEnv)
	if err := validateEnv(appEnv); err != nil {
		return err
	}

	// get pachd client, so we can upload output data from the user binary
	pachClient, err := client.NewFromAddress(fmt.Sprintf("%v:650", appEnv.PachdAddress))
	if err != nil {
		return err
	}
	go pachClient.KeepConnected(make(chan bool)) // we never cancel the connection

	// Get etcd client, so we can register our IP (so pachd can discover us)
	etcdClient, err := etcd.New(etcd.Config{
		Endpoints:   []string{fmt.Sprintf("%s:2379", appEnv.EtcdAddress)},
		DialTimeout: 15 * time.Second,
	})
	if err != nil {
		return err
	}

	// Construct worker API server. Get relevant pipeline or job info, and then
	// use that to create a worker.APIServer.
	var workerRcName string
	var apiServer *worker.APIServer
	if appEnv.PPSPipelineName != "" {
		pipelineInfo, err := getPipelineInfo(etcdClient, appEnv)
		if err != nil {
			return err
		}
		workerRcName = ppsserver.PipelineRcName(pipelineInfo.Pipeline.Name, pipelineInfo.Version)
		apiServer = worker.NewPipelineAPIServer(pachClient, pipelineInfo, appEnv.PodName)
	} else if appEnv.PPSJobID != "" {
		jobInfo, err := getJobInfo(etcdClient, appEnv)
		if err != nil {
			return err
		}
		workerRcName = ppsserver.JobRcName(jobInfo.Job.ID)
		apiServer = worker.NewJobAPIServer(pachClient, jobInfo, appEnv.PodName)
	}

	// Setup the hostPath mount to use a unique directory for this worker
	workerDir := filepath.Join(client.PPSHostPath, appEnv.PodName)
	if err := os.MkdirAll(workerDir, 0777); err != nil {
		return err
	}
	if err := os.Symlink(workerDir, client.PPSInputPrefix); err != nil {
		return err
	}

	// Start worker api server
	eg := errgroup.Group{}
	ready := make(chan error)
	eg.Go(func() error {
		return grpcutil.Serve(
			func(s *grpc.Server) {
				worker.RegisterWorkerServer(s, apiServer)
				close(ready)
			},
			grpcutil.ServeOptions{
				Version:    version.Version,
				MaxMsgSize: grpcutil.MaxMsgSize,
			},
			grpcutil.ServeEnv{
				GRPCPort: client.PPSWorkerPort,
			},
		)
	})

	// Wait until server is ready, then put our IP address into etcd, so pachd can
	// discover us
	<-ready
	key := path.Join(appEnv.PPSPrefix, "workers", workerRcName, appEnv.PPSWorkerIP)

	// Prepare to write "key" into etcd by creating lease -- if worker dies, our
	// IP will be removed from etcd
	ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
	resp, err := etcdClient.Grant(ctx, 10 /* seconds */)
	if err != nil {
		return err
	}
	etcdClient.KeepAlive(context.Background(), resp.ID) // keepalive forever

	// Actually write "key" into etcd
	ctx, _ = context.WithTimeout(context.Background(), 10*time.Second) // new ctx
	if _, err := etcdClient.Put(ctx, key, "", etcd.WithLease(resp.ID)); err != nil {
		return err
	}

	// If server ever exits, return error
	if err := eg.Wait(); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package e2e

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"

	habv1 "github.com/kinvolk/habitat-operator/pkg/apis/habitat/v1"
	utils "github.com/kinvolk/habitat-operator/test/e2e/framework"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	waitForPorts  = 1 * time.Minute
	configMapName = "peer-watch-file"

	nodejsImage = "kinvolk/nodejs-hab:test"
)

// TestFunction tests that the operator correctly created two Habitat Services and bound them together.
func TestFunction(t *testing.T) {
	// Get Habitat object from Habitat go example.
	habitatGo, err := utils.ConvertHabitat("resources/bind-config/habitat-go.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(habitatGo); err != nil {
		t.Fatal(err)
	}

	// Get Habitat object from Habitat db example.
	habitatDB, err := utils.ConvertHabitat("resources/bind-config/habitat-postgresql.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(habitatDB); err != nil {
		t.Fatal(err)
	}

	// Get Service object from example file.
	service, err := utils.ConvertService("resources/bind-config/service.yml")
	if err != nil {
		t.Fatal(err)
	}

	// Create Service.
	_, err = framework.KubeClient.CoreV1().Services(utils.TestNs).Create(service)
	if err != nil {
		t.Fatal(err)
	}

	// Get Secret object from example file.
	secret, err := utils.ConvertSecret("resources/bind-config/secret.yml")
	if err != nil {
		t.Fatal(err)
	}

	// Create Secret.
	_, err = framework.KubeClient.CoreV1().Secrets(utils.TestNs).Create(secret)
	if err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be ready.
	if err := framework.WaitForResources(habv1.HabitatNameLabel, habitatGo.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}
	if err := framework.WaitForResources(habv1.HabitatNameLabel, habitatDB.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}

	// Wait until endpoints are ready.
	if err := framework.WaitForEndpoints(service.ObjectMeta.Name); err != nil {
		t.Fatal(err)
	}

	time.Sleep(waitForPorts)

	// Get response from Habitat Service.
	url := fmt.Sprintf("http://%s:30001/", framework.ExternalIP)
	resp, err := http.Get(url)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Fatal("Habitat Service did not start correctly.")
	}

	bodyBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	// This msg is set in the config of the kinvolk/bindgo-hab Go Habitat Service.
	expectedMsg := "hello from port: 4444"
	actualMsg := string(bodyBytes)
	// actualMsg can contain whitespace and newlines or different formatting,
	// the only thing we need to check is it contains the expectedMsg.
	if !strings.Contains(actualMsg, expectedMsg) {
		t.Fatalf("Habitat Service msg does not match one in default.toml. Expected: *%s* got: *%s*", expectedMsg, actualMsg)
	}

	// Delete Service so it doesn't interfere with other tests.
	if err := framework.DeleteService(service.ObjectMeta.Name); err != nil {
		t.Fatal(err)
	}
}

// TestHabitatDelete tests Habitat deletion.
func TestHabitatDelete(t *testing.T) {
	// Get Habitat object from Habitat go example.
	habitat, err := utils.ConvertHabitat("resources/standalone/habitat.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(habitat); err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be ready.
	if err := framework.WaitForResources(habv1.HabitatNameLabel, habitat.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}

	// Delete Habitat.
	if err := framework.DeleteHabitat(habitat.ObjectMeta.Name); err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be deleted.
	if err := framework.WaitForResources(habv1.HabitatNameLabel, habitat.ObjectMeta.Name, 0); err != nil {
		t.Fatal(err)
	}

	// Check if all the resources the operator creates are deleted.
	// We do not care about secrets being deleted, as the user needs to delete those manually.
	d, err := framework.KubeClient.AppsV1beta1().Deployments(utils.TestNs).Get(habitat.ObjectMeta.Name, metav1.GetOptions{})
	if err == nil && d != nil {
		t.Fatal("Deployment was not deleted.")
	}

	// The CM with the peer IP should still be alive, despite the Habitat being deleted as it was created outside of the scope of a Habitat.
	_, err = framework.KubeClient.CoreV1().ConfigMaps(utils.TestNs).Get(configMapName, metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}
}
<commit_msg>Use quotes in error message, instead of *<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package e2e

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"

	habv1 "github.com/kinvolk/habitat-operator/pkg/apis/habitat/v1"
	utils "github.com/kinvolk/habitat-operator/test/e2e/framework"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	waitForPorts  = 1 * time.Minute
	configMapName = "peer-watch-file"

	nodejsImage = "kinvolk/nodejs-hab:test"
)

// TestFunction tests that the operator correctly created two Habitat Services and bound them together.
func TestFunction(t *testing.T) {
	// Get Habitat object from Habitat go example.
	habitatGo, err := utils.ConvertHabitat("resources/bind-config/habitat-go.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(habitatGo); err != nil {
		t.Fatal(err)
	}

	// Get Habitat object from Habitat db example.
	habitatDB, err := utils.ConvertHabitat("resources/bind-config/habitat-postgresql.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(habitatDB); err != nil {
		t.Fatal(err)
	}

	// Get Service object from example file.
	service, err := utils.ConvertService("resources/bind-config/service.yml")
	if err != nil {
		t.Fatal(err)
	}

	// Create Service.
	_, err = framework.KubeClient.CoreV1().Services(utils.TestNs).Create(service)
	if err != nil {
		t.Fatal(err)
	}

	// Get Secret object from example file.
	secret, err := utils.ConvertSecret("resources/bind-config/secret.yml")
	if err != nil {
		t.Fatal(err)
	}

	// Create Secret.
	_, err = framework.KubeClient.CoreV1().Secrets(utils.TestNs).Create(secret)
	if err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be ready.
	if err := framework.WaitForResources(habv1.HabitatNameLabel, habitatGo.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}
	if err := framework.WaitForResources(habv1.HabitatNameLabel, habitatDB.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}

	// Wait until endpoints are ready.
	if err := framework.WaitForEndpoints(service.ObjectMeta.Name); err != nil {
		t.Fatal(err)
	}

	time.Sleep(waitForPorts)

	// Get response from Habitat Service.
	url := fmt.Sprintf("http://%s:30001/", framework.ExternalIP)
	resp, err := http.Get(url)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Fatal("Habitat Service did not start correctly.")
	}

	bodyBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	// This msg is set in the config of the kinvolk/bindgo-hab Go Habitat Service.
	expectedMsg := "hello from port: 4444"
	actualMsg := string(bodyBytes)
	// actualMsg can contain whitespace and newlines or different formatting,
	// the only thing we need to check is it contains the expectedMsg.
	if !strings.Contains(actualMsg, expectedMsg) {
		t.Fatalf("Habitat Service msg does not match one in default.toml. Expected: \"%s\", got: \"%s\"", expectedMsg, actualMsg)
	}

	// Delete Service so it doesn't interfere with other tests.
	if err := framework.DeleteService(service.ObjectMeta.Name); err != nil {
		t.Fatal(err)
	}
}

// TestHabitatDelete tests Habitat deletion.
func TestHabitatDelete(t *testing.T) {
	// Get Habitat object from Habitat go example.
	habitat, err := utils.ConvertHabitat("resources/standalone/habitat.yml")
	if err != nil {
		t.Fatal(err)
	}

	if err := framework.CreateHabitat(habitat); err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be ready.
	if err := framework.WaitForResources(habv1.HabitatNameLabel, habitat.ObjectMeta.Name, 1); err != nil {
		t.Fatal(err)
	}

	// Delete Habitat.
	if err := framework.DeleteHabitat(habitat.ObjectMeta.Name); err != nil {
		t.Fatal(err)
	}

	// Wait for resources to be deleted.
	if err := framework.WaitForResources(habv1.HabitatNameLabel, habitat.ObjectMeta.Name, 0); err != nil {
		t.Fatal(err)
	}

	// Check if all the resources the operator creates are deleted.
	// We do not care about secrets being deleted, as the user needs to delete those manually.
	d, err := framework.KubeClient.AppsV1beta1().Deployments(utils.TestNs).Get(habitat.ObjectMeta.Name, metav1.GetOptions{})
	if err == nil && d != nil {
		t.Fatal("Deployment was not deleted.")
	}

	// The CM with the peer IP should still be alive, despite the Habitat being deleted as it was created outside of the scope of a Habitat.
	_, err = framework.KubeClient.CoreV1().ConfigMaps(utils.TestNs).Get(configMapName, metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package govcr

import (
	"net/http"
)

// vcrTransport is the heart of VCR. It provides
// an http.RoundTripper that wraps over the default
// one provided by Go's http package or a custom one
// if specified when calling NewVCR.
type vcrTransport struct {
	PCB      *pcb
	Cassette *cassette
}

// RoundTrip is an implementation of http.RoundTripper.
func (t *vcrTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	// Note: by convention resp should be nil if an error occurs with HTTP
	var resp *http.Response

	// copy the request before the body is closed by the HTTP server.
	copiedReq, err := copyRequest(req)
	if err != nil {
		t.PCB.Logger.Println(err)
		return nil, err
	}

	request, err := newRequest(req, t.PCB.Logger)
	if err != nil {
		return nil, err
	}
	request = t.PCB.RequestFilter(request)

	// attempt to use a track from the cassette that matches
	// the request if one exists.
	if trackNumber := t.PCB.seekTrack(t.Cassette, request); trackNumber != trackNotFound {
		// Only the played back response is filtered.
		// The live request and response should NOT EVER be changed!
		request, _ = newRequest(req, t.PCB.Logger)
		resp = t.PCB.filterResponse(t.Cassette.replayResponse(trackNumber, copiedReq), request)
	} else {
		// no recorded track was found so execute the request live
		t.PCB.Logger.Printf("INFO - Cassette '%s' - Executing request to live server for %s %s\n", t.Cassette.Name, req.Method, req.URL.String())

		resp, err = t.PCB.Transport.RoundTrip(req)

		if !t.PCB.DisableRecording {
			// the VCR is not in read-only mode so
			// record the HTTP traffic into a new track on the cassette
			t.PCB.Logger.Printf("INFO - Cassette '%s' - Recording new track for %s %s as %s %s\n", t.Cassette.Name, req.Method, req.URL.String(), copiedReq.Method, copiedReq.URL)
			if err := recordNewTrackToCassette(t.Cassette, copiedReq, resp, err); err != nil {
				t.PCB.Logger.Println(err)
			}
		}
	}

	return resp, err
}
<commit_msg>Refactored if .. else<commit_after>package govcr

import (
	"net/http"
)

// vcrTransport is the heart of VCR. It provides
// an http.RoundTripper that wraps over the default
// one provided by Go's http package or a custom one
// if specified when calling NewVCR.
type vcrTransport struct {
	PCB      *pcb
	Cassette *cassette
}

// RoundTrip is an implementation of http.RoundTripper.
func (t *vcrTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	// Note: by convention resp should be nil if an error occurs with HTTP
	var resp *http.Response

	// copy the request before the body is closed by the HTTP server.
	copiedReq, err := copyRequest(req)
	if err != nil {
		t.PCB.Logger.Println(err)
		return nil, err
	}

	request, err := newRequest(req, t.PCB.Logger)
	if err != nil {
		return nil, err
	}
	request = t.PCB.RequestFilter(request)

	// attempt to use a track from the cassette that matches
	// the request if one exists.
	if trackNumber := t.PCB.seekTrack(t.Cassette, request); trackNumber != trackNotFound {
		// Only the played back response is filtered.
		// The live request and response should NOT EVER be changed!
		request, _ = newRequest(req, t.PCB.Logger)
		resp = t.PCB.filterResponse(t.Cassette.replayResponse(trackNumber, copiedReq), request)

		return resp, err
	}

	// no recorded track was found so execute the request live
	t.PCB.Logger.Printf("INFO - Cassette '%s' - Executing request to live server for %s %s\n", t.Cassette.Name, req.Method, req.URL.String())

	resp, err = t.PCB.Transport.RoundTrip(req)

	if !t.PCB.DisableRecording {
		// the VCR is not in read-only mode so
		// record the HTTP traffic into a new track on the cassette
		t.PCB.Logger.Printf("INFO - Cassette '%s' - Recording new track for %s %s as %s %s\n", t.Cassette.Name, req.Method, req.URL.String(), copiedReq.Method, copiedReq.URL)
		if err := recordNewTrackToCassette(t.Cassette, copiedReq, resp, err); err != nil {
			t.PCB.Logger.Println(err)
		}
	}

	return resp, err
}
<|endoftext|>


<commit_before>// Copyright 2017 Google Inc.
//
// 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 main

import (
	"fmt"
	"log"

	"github.com/google/shenzhen-go/api"
	"github.com/google/shenzhen-go/jsutil"
	"github.com/gopherjs/gopherjs/js"
)

const (
	activeColour = "#09f"
	normalColour = "#000"
	errorColour  = "#f06"

	errRectStyle = "fill: #fee; fill-opacity: 0.5; stroke: #533; stroke-width:1"
	errTextStyle = "font-family:Go; font-size:16; user-select:none; pointer-events:none"

	pinRadius = 5
	lineWidth = 2
	snapQuad  = 144
)

var (
	graphPath = jsutil.MustGetGlobal("graphPath").String()

	graphPropertiesPanel = jsutil.MustGetElement("graph-properties")
	nodePropertiesPanel  = jsutil.MustGetElement("node-properties")
	rhsPanel             = graphPropertiesPanel

	client api.ShenzhenGoClient
)

func showRHSPanel(p *jsutil.Element) {
	if p == rhsPanel {
		return
	}
	rhsPanel.Get("style").Set("display", "none")
	rhsPanel = p
	rhsPanel.Get("style").Set("display", nil)
}

func main() {
	client = api.NewShenzhenGoClient(jsutil.MustGetGlobal("apiURL").String())

	d := &diagram{
		Element: jsutil.MustGetElement("diagram"),
	}
	d.errLabel = newTextBox(d, "", errTextStyle, errRectStyle, 0, 0, 0, 32)
	d.errLabel.hide()

	g, err := loadGraph(d, jsutil.MustGetGlobal("GraphJSON").String())
	if err != nil {
		log.Fatalf("Couldn't load graph: %v", err)
	}
	d.graph = g

	d.
		AddEventListener("mousedown", d.mouseDown).
		AddEventListener("mousemove", d.mouseMove).
		AddEventListener("mouseup", d.mouseUp)

	jsutil.MustGetElement("graph-save").
		AddEventListener("click", g.save)
	jsutil.MustGetElement("graph-properties-save").
		AddEventListener("click", g.saveProperties)

	jsutil.MustGetElement("node-save-link").
		AddEventListener("click", d.saveSelected)
	jsutil.MustGetElement("node-metadata-link").
		AddEventListener("click", func(*js.Object) {
			d.selectedItem.(*Node).showSubPanel(nodeMetadataSubpanel)
		})

	for n, e := range nodePartEditors {
		for m, p := range e.Panels {
			p := p
			jsutil.MustGetElement(fmt.Sprintf("node-%s-%s-link", n, m)).
				AddEventListener("click",
					func(*js.Object) {
						d.selectedItem.(*Node).showSubPanel(p)
					})
		}
	}
}
<commit_msg>Stop breaking on TravisCI?<commit_after>// Copyright 2017 Google Inc.
//
// 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.

// +build js

package main

import (
	"fmt"
	"log"

	"github.com/google/shenzhen-go/api"
	"github.com/google/shenzhen-go/jsutil"
	"github.com/gopherjs/gopherjs/js"
)

const (
	activeColour = "#09f"
	normalColour = "#000"
	errorColour  = "#f06"

	errRectStyle = "fill: #fee; fill-opacity: 0.5; stroke: #533; stroke-width:1"
	errTextStyle = "font-family:Go; font-size:16; user-select:none; pointer-events:none"

	pinRadius = 5
	lineWidth = 2
	snapQuad  = 144
)

var (
	graphPath = jsutil.MustGetGlobal("graphPath").String()

	graphPropertiesPanel = jsutil.MustGetElement("graph-properties")
	nodePropertiesPanel  = jsutil.MustGetElement("node-properties")
	rhsPanel             = graphPropertiesPanel

	client api.ShenzhenGoClient
)

func showRHSPanel(p *jsutil.Element) {
	if p == rhsPanel {
		return
	}
	rhsPanel.Get("style").Set("display", "none")
	rhsPanel = p
	rhsPanel.Get("style").Set("display", nil)
}

func main() {
	client = api.NewShenzhenGoClient(jsutil.MustGetGlobal("apiURL").String())

	d := &diagram{
		Element: jsutil.MustGetElement("diagram"),
	}
	d.errLabel = newTextBox(d, "", errTextStyle, errRectStyle, 0, 0, 0, 32)
	d.errLabel.hide()

	g, err := loadGraph(d, jsutil.MustGetGlobal("GraphJSON").String())
	if err != nil {
		log.Fatalf("Couldn't load graph: %v", err)
	}
	d.graph = g

	d.
		AddEventListener("mousedown", d.mouseDown).
		AddEventListener("mousemove", d.mouseMove).
		AddEventListener("mouseup", d.mouseUp)

	jsutil.MustGetElement("graph-save").
		AddEventListener("click", g.save)
	jsutil.MustGetElement("graph-properties-save").
		AddEventListener("click", g.saveProperties)

	jsutil.MustGetElement("node-save-link").
		AddEventListener("click", d.saveSelected)
	jsutil.MustGetElement("node-metadata-link").
		AddEventListener("click", func(*js.Object) {
			d.selectedItem.(*Node).showSubPanel(nodeMetadataSubpanel)
		})

	for n, e := range nodePartEditors {
		for m, p := range e.Panels {
			p := p
			jsutil.MustGetElement(fmt.Sprintf("node-%s-%s-link", n, m)).
				AddEventListener("click",
					func(*js.Object) {
						d.selectedItem.(*Node).showSubPanel(p)
					})
		}
	}
}
<|endoftext|>


<commit_before>// +build e2e

/*
Copyright 2020 The Knative 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 ha

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"

	"knative.dev/pkg/ptr"
	"knative.dev/pkg/system"

	appsv1 "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	pkgnet "knative.dev/pkg/network"
	pkgTest "knative.dev/pkg/test"
	"knative.dev/serving/pkg/apis/autoscaling"
	"knative.dev/serving/pkg/networking"
	revisionresourcenames "knative.dev/serving/pkg/reconciler/revision/resources/names"
	rtesting "knative.dev/serving/pkg/testing/v1"
	"knative.dev/serving/test"
	"knative.dev/serving/test/e2e"
)

const (
	activatorDeploymentName = "activator"
	minProbes               = 400 // We want to send at least 400 requests.
)

func TestActivatorHAGraceful(t *testing.T) {
	testActivatorHA(t, nil, 1)
}

func TestActivatorHANonGraceful(t *testing.T) {
	// For non-graceful tests, we want the pod to receive a SIGKILL straight away.
	testActivatorHA(t, ptr.Int64(0), 0.90)
}

// The Activator does not have leader election enabled.
// The test ensures that stopping one of the activator pods doesn't affect user applications.
// One service is probed during activator restarts and another service is used for testing
// that we can scale from zero after activator restart.
func testActivatorHA(t *testing.T, gracePeriod *int64, slo float64) {
	clients := e2e.Setup(t)
	ctx := context.Background()

	podDeleteOptions := metav1.DeleteOptions{GracePeriodSeconds: gracePeriod}

	desiredScale, err := waitForActivatorScale(ctx, clients.KubeClient)
	if err != nil {
		t.Fatalf("Deployment %s not scaled > 1: %s", activatorDeploymentName, err)
	}

	if err != nil {
		t.Fatal("Failed to get activator pods:", err)
	}

	// Create first service that we will continually probe during activator restart.
	names, resources := createPizzaPlanetService(t,
		rtesting.WithConfigAnnotations(map[string]string{
			autoscaling.MinScaleAnnotationKey:  "1",  // Make sure we don't scale to zero during the test.
			autoscaling.TargetBurstCapacityKey: "-1", // Make sure all requests go through the activator.
		}),
	)
	test.EnsureTearDown(t, clients, &names)

	// Create second service that will be scaled to zero and after stopping the activator we'll
	// ensure it can be scaled back from zero.
	namesScaleToZero, resourcesScaleToZero := createPizzaPlanetService(t,
		rtesting.WithConfigAnnotations(map[string]string{
			autoscaling.WindowAnnotationKey:    autoscaling.WindowMin.String(), // Make sure we scale to zero quickly.
			autoscaling.TargetBurstCapacityKey: "-1",                           // Make sure all requests go through the activator.
		}),
	)
	test.EnsureTearDown(t, clients, &namesScaleToZero)

	t.Log("Starting prober")
	prober := test.NewProberManager(t.Logf, clients, minProbes, test.AddRootCAtoTransport(context.Background(), t.Logf, clients, test.ServingFlags.HTTPS))

	prober.Spawn(resources.Service.Status.URL.URL())
	defer assertSLO(t, prober, slo)

	activatorPods, err := gatherBackingActivators(ctx, clients.KubeClient, test.ServingNamespace, resources.Revision.Name, resourcesScaleToZero.Revision.Name)
	if err != nil {
		t.Fatal("failed to gather backing activators:", err)
	}

	t.Logf("Rolling %d activators", len(activatorPods))
	i := 0
	for name, ip := range activatorPods {
		t.Logf("Waiting for %s to scale to zero", namesScaleToZero.Revision)
		if err := e2e.WaitForScaleToZero(t, revisionresourcenames.Deployment(resourcesScaleToZero.Revision), clients); err != nil {
			t.Fatal("Failed to scale to zero:", err)
		}

		t.Logf("Deleting activator%d (%s)", i, name)
		if err := clients.KubeClient.CoreV1().Pods(system.Namespace()).Delete(ctx, name, podDeleteOptions); err != nil {
			t.Fatalf("Failed to delete pod %s: %v", name, err)
		}

		// Wait for the killed activator to disappear from the knative service's endpoints.
		if err := waitForEndpointsState(clients.KubeClient, resourcesScaleToZero.Revision.Name, test.ServingNamespace, readyEndpointsDoNotContain(ip)); err != nil {
			t.Fatal("Failed to wait for the service to update its endpoints:", err)
		}
		if gracePeriod != nil && *gracePeriod == 0 {
			t.Log("Allow the network to notice the missing endpoint")
			time.Sleep(pkgnet.DefaultDrainTimeout)
		}

		t.Log("Test if service still works")
		assertServiceEventuallyWorks(t, clients, namesScaleToZero, resourcesScaleToZero.Service.Status.URL.URL(), test.PizzaPlanetText1, test.ServingFlags.ResolvableDomain)

		t.Logf("Wait for activator%d (%s) to vanish", i, name)
		if err := pkgTest.WaitForPodDeleted(ctx, clients.KubeClient, name, system.Namespace()); err != nil {
			t.Fatalf("Did not observe %s to actually be deleted: %v", name, err)
		}
		// Check for the endpoint to appear in the activator's endpoints, since this revision may pick a subset of those endpoints.
		t.Logf("Wait for a new activator to spin up")
		if err := pkgTest.WaitForServiceEndpoints(ctx, clients.KubeClient, networking.ActivatorServiceName, system.Namespace(), desiredScale); err != nil {
			t.Fatalf("Deployment %s failed to scale up: %v", activatorDeploymentName, err)
		}
		if gracePeriod != nil && *gracePeriod == 0 {
			t.Log("Allow the network to notice the new endpoint")
			time.Sleep(pkgnet.DefaultDrainTimeout)
		}
		i++
	}
}

func assertSLO(t *testing.T, p test.Prober, slo float64) {
	t.Helper()
	if err := p.Stop(); err != nil {
		t.Error("Failed to stop prober:", err)
	}
	if err := test.CheckSLO(slo, t.Name(), p); err != nil {
		t.Error("CheckSLO failed:", err)
	}
}

type podIPs map[string]string

func gatherBackingActivators(ctx context.Context, client kubernetes.Interface, namespace string, revs ...string) (podIPs, error) {
	pods := make(podIPs)

	for _, rev := range revs {
		endpoints := client.CoreV1().Endpoints(namespace)
		e, err := endpoints.Get(ctx, rev, metav1.GetOptions{})

		if err != nil {
			return nil, fmt.Errorf("failed to gather %s endpoints: %w", rev, err)
		}

		for _, subset := range e.Subsets {
			for _, address := range subset.Addresses {
				if address.TargetRef == nil {
					return nil, fmt.Errorf("%s service is not pointing to a pod", rev)
				}

				name := address.TargetRef.Name
				if !strings.HasPrefix(name, activatorDeploymentName) {
					return nil, fmt.Errorf("%s service is not pointing to an activator pod but: %s", rev, address.TargetRef.Name)
				}

				pods[name] = address.IP
			}
		}
	}

	return pods, nil
}

func waitForActivatorScale(ctx context.Context, client kubernetes.Interface) (int, error) {
	desiredScale := 0
	check := func(d *appsv1.Deployment) (bool, error) {
		if *d.Spec.Replicas < 2 {
			return false, errors.New("spec.replicaCount should be > 1")
		}
		desiredScale = int(*d.Spec.Replicas)
		for _, c := range d.Status.Conditions {
			if c.Type == appsv1.DeploymentAvailable {
				return c.Status == v1.ConditionTrue, nil
			}
		}
		return false, nil
	}

	err := pkgTest.WaitForDeploymentState(
		ctx,
		client,
		activatorDeploymentName,
		check,
		"ActivatorIsScaled",
		system.Namespace(),
		time.Minute,
	)

	return desiredScale, err
}
<commit_msg>Remove unreachable code in testActivatorHA (#11549)<commit_after>// +build e2e

/*
Copyright 2020 The Knative 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 ha

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"

	"knative.dev/pkg/ptr"
	"knative.dev/pkg/system"

	appsv1 "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	pkgnet "knative.dev/pkg/network"
	pkgTest "knative.dev/pkg/test"
	"knative.dev/serving/pkg/apis/autoscaling"
	"knative.dev/serving/pkg/networking"
	revisionresourcenames "knative.dev/serving/pkg/reconciler/revision/resources/names"
	rtesting "knative.dev/serving/pkg/testing/v1"
	"knative.dev/serving/test"
	"knative.dev/serving/test/e2e"
)

const (
	activatorDeploymentName = "activator"
	minProbes               = 400 // We want to send at least 400 requests.
)

func TestActivatorHAGraceful(t *testing.T) {
	testActivatorHA(t, nil, 1)
}

func TestActivatorHANonGraceful(t *testing.T) {
	// For non-graceful tests, we want the pod to receive a SIGKILL straight away.
	testActivatorHA(t, ptr.Int64(0), 0.90)
}

// The Activator does not have leader election enabled.
// The test ensures that stopping one of the activator pods doesn't affect user applications.
// One service is probed during activator restarts and another service is used for testing
// that we can scale from zero after activator restart.
func testActivatorHA(t *testing.T, gracePeriod *int64, slo float64) {
	clients := e2e.Setup(t)
	ctx := context.Background()

	podDeleteOptions := metav1.DeleteOptions{GracePeriodSeconds: gracePeriod}

	desiredScale, err := waitForActivatorScale(ctx, clients.KubeClient)
	if err != nil {
		t.Fatalf("Deployment %s not scaled > 1: %s", activatorDeploymentName, err)
	}

	// Create first service that we will continually probe during activator restart.
	names, resources := createPizzaPlanetService(t,
		rtesting.WithConfigAnnotations(map[string]string{
			autoscaling.MinScaleAnnotationKey:  "1",  // Make sure we don't scale to zero during the test.
			autoscaling.TargetBurstCapacityKey: "-1", // Make sure all requests go through the activator.
		}),
	)
	test.EnsureTearDown(t, clients, &names)

	// Create second service that will be scaled to zero and after stopping the activator we'll
	// ensure it can be scaled back from zero.
	namesScaleToZero, resourcesScaleToZero := createPizzaPlanetService(t,
		rtesting.WithConfigAnnotations(map[string]string{
			autoscaling.WindowAnnotationKey:    autoscaling.WindowMin.String(), // Make sure we scale to zero quickly.
			autoscaling.TargetBurstCapacityKey: "-1",                           // Make sure all requests go through the activator.
		}),
	)
	test.EnsureTearDown(t, clients, &namesScaleToZero)

	t.Log("Starting prober")
	prober := test.NewProberManager(t.Logf, clients, minProbes, test.AddRootCAtoTransport(context.Background(), t.Logf, clients, test.ServingFlags.HTTPS))

	prober.Spawn(resources.Service.Status.URL.URL())
	defer assertSLO(t, prober, slo)

	activatorPods, err := gatherBackingActivators(ctx, clients.KubeClient, test.ServingNamespace, resources.Revision.Name, resourcesScaleToZero.Revision.Name)
	if err != nil {
		t.Fatal("failed to gather backing activators:", err)
	}

	t.Logf("Rolling %d activators", len(activatorPods))
	i := 0
	for name, ip := range activatorPods {
		t.Logf("Waiting for %s to scale to zero", namesScaleToZero.Revision)
		if err := e2e.WaitForScaleToZero(t, revisionresourcenames.Deployment(resourcesScaleToZero.Revision), clients); err != nil {
			t.Fatal("Failed to scale to zero:", err)
		}

		t.Logf("Deleting activator%d (%s)", i, name)
		if err := clients.KubeClient.CoreV1().Pods(system.Namespace()).Delete(ctx, name, podDeleteOptions); err != nil {
			t.Fatalf("Failed to delete pod %s: %v", name, err)
		}

		// Wait for the killed activator to disappear from the knative service's endpoints.
		if err := waitForEndpointsState(clients.KubeClient, resourcesScaleToZero.Revision.Name, test.ServingNamespace, readyEndpointsDoNotContain(ip)); err != nil {
			t.Fatal("Failed to wait for the service to update its endpoints:", err)
		}
		if gracePeriod != nil && *gracePeriod == 0 {
			t.Log("Allow the network to notice the missing endpoint")
			time.Sleep(pkgnet.DefaultDrainTimeout)
		}

		t.Log("Test if service still works")
		assertServiceEventuallyWorks(t, clients, namesScaleToZero, resourcesScaleToZero.Service.Status.URL.URL(), test.PizzaPlanetText1, test.ServingFlags.ResolvableDomain)

		t.Logf("Wait for activator%d (%s) to vanish", i, name)
		if err := pkgTest.WaitForPodDeleted(ctx, clients.KubeClient, name, system.Namespace()); err != nil {
			t.Fatalf("Did not observe %s to actually be deleted: %v", name, err)
		}
		// Check for the endpoint to appear in the activator's endpoints, since this revision may pick a subset of those endpoints.
		t.Logf("Wait for a new activator to spin up")
		if err := pkgTest.WaitForServiceEndpoints(ctx, clients.KubeClient, networking.ActivatorServiceName, system.Namespace(), desiredScale); err != nil {
			t.Fatalf("Deployment %s failed to scale up: %v", activatorDeploymentName, err)
		}
		if gracePeriod != nil && *gracePeriod == 0 {
			t.Log("Allow the network to notice the new endpoint")
			time.Sleep(pkgnet.DefaultDrainTimeout)
		}
		i++
	}
}

func assertSLO(t *testing.T, p test.Prober, slo float64) {
	t.Helper()
	if err := p.Stop(); err != nil {
		t.Error("Failed to stop prober:", err)
	}
	if err := test.CheckSLO(slo, t.Name(), p); err != nil {
		t.Error("CheckSLO failed:", err)
	}
}

type podIPs map[string]string

func gatherBackingActivators(ctx context.Context, client kubernetes.Interface, namespace string, revs ...string) (podIPs, error) {
	pods := make(podIPs)

	for _, rev := range revs {
		endpoints := client.CoreV1().Endpoints(namespace)
		e, err := endpoints.Get(ctx, rev, metav1.GetOptions{})

		if err != nil {
			return nil, fmt.Errorf("failed to gather %s endpoints: %w", rev, err)
		}

		for _, subset := range e.Subsets {
			for _, address := range subset.Addresses {
				if address.TargetRef == nil {
					return nil, fmt.Errorf("%s service is not pointing to a pod", rev)
				}

				name := address.TargetRef.Name
				if !strings.HasPrefix(name, activatorDeploymentName) {
					return nil, fmt.Errorf("%s service is not pointing to an activator pod but: %s", rev, address.TargetRef.Name)
				}

				pods[name] = address.IP
			}
		}
	}

	return pods, nil
}

func waitForActivatorScale(ctx context.Context, client kubernetes.Interface) (int, error) {
	desiredScale := 0
	check := func(d *appsv1.Deployment) (bool, error) {
		if *d.Spec.Replicas < 2 {
			return false, errors.New("spec.replicaCount should be > 1")
		}
		desiredScale = int(*d.Spec.Replicas)
		for _, c := range d.Status.Conditions {
			if c.Type == appsv1.DeploymentAvailable {
				return c.Status == v1.ConditionTrue, nil
			}
		}
		return false, nil
	}

	err := pkgTest.WaitForDeploymentState(
		ctx,
		client,
		activatorDeploymentName,
		check,
		"ActivatorIsScaled",
		system.Namespace(),
		time.Minute,
	)

	return desiredScale, err
}
<|endoftext|>


<commit_before>// Copyright 2017 Google Inc.
//
// 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 meminfo

import (
	"bytes"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/spf13/afero"
	"github.com/stretchrcom/testify/assert"

	"github.com/soumya92/barista/base"
	testBar "github.com/soumya92/barista/testing/bar"
	"github.com/soumya92/barista/timing"
)

type meminfo map[string]int

func shouldReturn(info meminfo) {
	var out bytes.Buffer
	for key, value := range info {
		out.WriteString(fmt.Sprintf("%s:\t%d kB\n", key, value))
	}
	afero.WriteFile(fs, "/proc/meminfo", out.Bytes(), 0644)
}

func TestMeminfo(t *testing.T) {
	assert := assert.New(t)
	fs = afero.NewMemMapFs()
	testBar.New(t)
	currentInfo = base.ErrorValue{}
	once = sync.Once{}

	shouldReturn(meminfo{
		"MemAvailable": 2048,
		"MemTotal":     4096,
		"MemFree":      1024,
	})

	avail := New().Template(`{{.Available.Kibibytes}}`)
	free := New().Template(`{{.FreeFrac "Mem"}}`)

	testBar.Run(avail, free)
	testBar.LatestOutput().AssertText(
		[]string{"2048", "0.25"}, "on start")

	shouldReturn(meminfo{
		"MemAvailable": 1024,
		"MemTotal":     4096,
		"MemFree":      256,
		"Cached":       512,
	})
	testBar.Tick()

	testBar.LatestOutput().AssertText(
		[]string{"1024", "0.0625"}, "on tick")

	shouldReturn(meminfo{
		"Cached":   1024,
		"Buffers":  512,
		"MemTotal": 4096,
		"MemFree":  512,
	})
	testBar.Tick()

	testBar.LatestOutput().AssertText(
		[]string{"2048", "0.125"}, "on tick")

	beforeTick := timing.Now()
	RefreshInterval(time.Minute)
	testBar.Tick()
	assert.Equal(time.Minute, timing.Now().Sub(beforeTick), "RefreshInterval change")

	testBar.LatestOutput().Expect("on refresh interval change")
}

func TestErrors(t *testing.T) {
	fs = afero.NewMemMapFs()
	testBar.New(t)
	currentInfo = base.ErrorValue{}
	once = sync.Once{}

	availFrac := New().Template(`{{.AvailFrac}}`)
	free := New().Template(`{{.MemFree | ibytesize}}`)
	total := New().Template(`{{.MemTotal | bytesize}}`)

	testBar.Run(availFrac, free, total)
	testBar.LatestOutput().AssertError("on start if missing meminfo")

	afero.WriteFile(fs, "/proc/meminfo", []byte(`
	-- Lines in weird formats --
	Empty:

	Valid line:
	MemAvailable: 1024 kB
	Missing colon:
	MemFree 1024 kB
	Fields are non-numeric:
	MemTotal: ABCD kB
	`), 0644)
	testBar.Tick()
	out := testBar.LatestOutput()
	out.At(1).AssertError("non-numeric value")
	out.At(2).AssertError("non-numeric value")
	// MemAvailable is parsed, but total is 0.
	out.At(0).AssertText("+Inf")

	afero.WriteFile(fs, "/proc/meminfo", []byte(`
	MemAvailable: 1024 kB
	MemFree: 1024 kB
	MemTotal: 2048 kB
	`), 0644)
	testBar.Tick()
	testBar.LatestOutput().AssertText(
		[]string{"0.5", "1.0 MiB", "2.1 MB"},
		"when meminfo is back to normal")
}
<commit_msg>Add output format test to meminfo<commit_after>// Copyright 2017 Google Inc.
//
// 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 meminfo

import (
	"bytes"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/spf13/afero"
	"github.com/stretchrcom/testify/assert"

	"github.com/soumya92/barista/base"
	testBar "github.com/soumya92/barista/testing/bar"
	"github.com/soumya92/barista/timing"
)

type meminfo map[string]int

func shouldReturn(info meminfo) {
	var out bytes.Buffer
	for key, value := range info {
		out.WriteString(fmt.Sprintf("%s:\t%d kB\n", key, value))
	}
	afero.WriteFile(fs, "/proc/meminfo", out.Bytes(), 0644)
}

func TestMeminfo(t *testing.T) {
	assert := assert.New(t)
	fs = afero.NewMemMapFs()
	testBar.New(t)
	currentInfo = base.ErrorValue{}
	once = sync.Once{}

	shouldReturn(meminfo{
		"MemAvailable": 2048,
		"MemTotal":     4096,
		"MemFree":      1024,
	})

	def := New()
	avail := New().Template(`{{.Available.Kibibytes}}`)
	free := New().Template(`{{.FreeFrac "Mem"}}`)

	testBar.Run(def, avail, free)
	testBar.LatestOutput().AssertText(
		[]string{"Mem: 2.0 MiB", "2048", "0.25"}, "on start")

	shouldReturn(meminfo{
		"MemAvailable": 1024,
		"MemTotal":     4096,
		"MemFree":      256,
		"Cached":       512,
	})
	testBar.Tick()

	testBar.LatestOutput().AssertText(
		[]string{"Mem: 1.0 MiB", "1024", "0.0625"}, "on tick")

	shouldReturn(meminfo{
		"Cached":   1024,
		"Buffers":  512,
		"MemTotal": 4096,
		"MemFree":  512,
	})
	testBar.Tick()

	testBar.LatestOutput().AssertText(
		[]string{"Mem: 2.0 MiB", "2048", "0.125"}, "on tick")

	def.Template(`{{.Buffers.Mebibytes}}`)
	testBar.LatestOutput().AssertText(
		[]string{"0.5", "2048", "0.125"}, "on template change")

	beforeTick := timing.Now()
	RefreshInterval(time.Minute)
	testBar.Tick()
	assert.Equal(time.Minute, timing.Now().Sub(beforeTick), "RefreshInterval change")

	testBar.LatestOutput().Expect("on refresh interval change")
}

func TestErrors(t *testing.T) {
	fs = afero.NewMemMapFs()
	testBar.New(t)
	currentInfo = base.ErrorValue{}
	once = sync.Once{}

	availFrac := New().Template(`{{.AvailFrac}}`)
	free := New().Template(`{{.MemFree | ibytesize}}`)
	total := New().Template(`{{.MemTotal | bytesize}}`)

	testBar.Run(availFrac, free, total)
	testBar.LatestOutput().AssertError("on start if missing meminfo")

	afero.WriteFile(fs, "/proc/meminfo", []byte(`
	-- Lines in weird formats --
	Empty:

	Valid line:
	MemAvailable: 1024 kB
	Missing colon:
	MemFree 1024 kB
	Fields are non-numeric:
	MemTotal: ABCD kB
	`), 0644)
	testBar.Tick()
	out := testBar.LatestOutput()
	out.At(1).AssertError("non-numeric value")
	out.At(2).AssertError("non-numeric value")
	// MemAvailable is parsed, but total is 0.
	out.At(0).AssertText("+Inf")

	afero.WriteFile(fs, "/proc/meminfo", []byte(`
	MemAvailable: 1024 kB
	MemFree: 1024 kB
	MemTotal: 2048 kB
	`), 0644)
	testBar.Tick()
	testBar.LatestOutput().AssertText(
		[]string{"0.5", "1.0 MiB", "2.1 MB"},
		"when meminfo is back to normal")
}
<|endoftext|>


<commit_before>package plugins

import (
	"strings"

	"net/url"

	"time"

	"bytes"

	"github.com/Krognol/go-wolfram"
	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/helpers"
	"github.com/Seklfreak/Robyul2/metrics"
	"github.com/bwmarrin/discordgo"
)

type WolframAlpha struct{}

func (m *WolframAlpha) Commands() []string {
	return []string{
		"wolfram",
		"w",
		"ask",
	}
}

func (m *WolframAlpha) Init(session *discordgo.Session) {

}

func (m *WolframAlpha) TypingLoop(channelID string, quitChannel chan int) {
	for {
		select {
		case <-quitChannel:
			return
		default:
			cache.GetSession().ChannelTyping(channelID)
			time.Sleep(5 * time.Second)
		}
	}
}

func (m *WolframAlpha) Action(command string, content string, msg *discordgo.Message, session *discordgo.Session) {
	var err error
	var res string
	var imageSearch bool
	if strings.HasPrefix(content, "image ") || strings.HasPrefix(content, "img ") {
		content = strings.TrimLeft(content, "image ")
		content = strings.TrimLeft(content, "img ")
		content = strings.TrimSpace(content)

		imageSearch = true
	}

	if content == "" || content == "img" {
		return
	}

	quitChannel := make(chan int)
	defer func() { quitChannel <- 0 }()

	go m.TypingLoop(msg.ChannelID, quitChannel)

	wolframClient := &wolfram.Client{AppID: helpers.GetConfig().Path("wolframalpha.appid").Data().(string)}

	if !imageSearch {
		metrics.WolframAlphaRequests.Add(1)
		res, err = wolframClient.GetShortAnswerQuery(content, wolfram.Metric, 10)
		helpers.Relax(err)

		if res == "No short answer available" {
			imageSearch = true
		}
	}

	if imageSearch {
		urlValues := url.Values{}
		//urlValues.Add("foreground", "white")
		//urlValues.Add("background", "35393E")
		urlValues.Add("layout", "labelbar")
		urlValues.Add("timeout", "30")

		metrics.WolframAlphaRequests.Add(1)
		image, _, err := wolframClient.GetSimpleQuery(content, urlValues)
		helpers.Relax(err)

		buf := new(bytes.Buffer)
		buf.ReadFrom(image)
		if strings.Contains(buf.String(), "Wolfram|Alpha did not understand your input") {
			helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.wolframalpha.error"))
			return
		}

		_, err = helpers.SendComplex(
			msg.ChannelID, &discordgo.MessageSend{
				Files: []*discordgo.File{
					{
						Name:   "wolframalpha-robyul.png",
						Reader: bytes.NewReader(buf.Bytes()),
					},
				},
			})
		helpers.RelaxEmbed(err, msg.ChannelID, msg.ID)
		return
	}

	if res == "" || res == "Wolfram|Alpha did not understand your input" {
		helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.wolframalpha.error"))
		return
	}

	_, err = helpers.SendMessage(msg.ChannelID, res)
	helpers.RelaxMessage(err, msg.ChannelID, msg.ID)
}
<commit_msg>[wolframalpha] refactoring<commit_after>package plugins

import (
	"strings"

	"net/url"

	"bytes"

	"github.com/Krognol/go-wolfram"
	"github.com/Seklfreak/Robyul2/helpers"
	"github.com/Seklfreak/Robyul2/metrics"
	"github.com/bwmarrin/discordgo"
)

type WolframAlpha struct{}

func (m *WolframAlpha) Commands() []string {
	return []string{
		"wolfram",
		"w",
		"ask",
	}
}

func (m *WolframAlpha) Init(session *discordgo.Session) {

}

func (m *WolframAlpha) Action(command string, content string, msg *discordgo.Message, session *discordgo.Session) {
	var err error
	var res string
	var imageSearch bool
	if strings.HasPrefix(content, "image ") || strings.HasPrefix(content, "img ") {
		content = strings.TrimLeft(content, "image ")
		content = strings.TrimLeft(content, "img ")
		content = strings.TrimSpace(content)

		imageSearch = true
	}

	if content == "" || content == "img" {
		return
	}

	quitChannel := helpers.StartTypingLoop(msg.ChannelID)
	defer func() { quitChannel <- 0 }()

	wolframClient := &wolfram.Client{AppID: helpers.GetConfig().Path("wolframalpha.appid").Data().(string)}

	if !imageSearch {
		metrics.WolframAlphaRequests.Add(1)
		res, err = wolframClient.GetShortAnswerQuery(content, wolfram.Metric, 10)
		helpers.Relax(err)

		if res == "No short answer available" {
			imageSearch = true
		}
	}

	if imageSearch {
		urlValues := url.Values{}
		//urlValues.Add("foreground", "white")
		//urlValues.Add("background", "35393E")
		urlValues.Add("layout", "labelbar")
		urlValues.Add("timeout", "30")

		metrics.WolframAlphaRequests.Add(1)
		image, _, err := wolframClient.GetSimpleQuery(content, urlValues)
		helpers.Relax(err)

		buf := new(bytes.Buffer)
		buf.ReadFrom(image)
		if strings.Contains(buf.String(), "Wolfram|Alpha did not understand your input") {
			helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.wolframalpha.error"))
			return
		}

		quitChannel <- 0
		_, err = helpers.SendComplex(
			msg.ChannelID, &discordgo.MessageSend{
				Files: []*discordgo.File{
					{
						Name:   "wolframalpha-robyul.png",
						Reader: bytes.NewReader(buf.Bytes()),
					},
				},
			})
		helpers.RelaxEmbed(err, msg.ChannelID, msg.ID)
		return
	}

	if res == "" || res == "Wolfram|Alpha did not understand your input" {
		quitChannel <- 0
		helpers.SendMessage(msg.ChannelID, helpers.GetText("plugins.wolframalpha.error"))
		return
	}

	quitChannel <- 0
	_, err = helpers.SendMessage(msg.ChannelID, res)
	helpers.RelaxMessage(err, msg.ChannelID, msg.ID)
}
<|endoftext|>


<commit_before>package commands

import (
	"crypto/sha256"
	"database/sql"
	"fmt"
	"log"
	"strings"

	"github.com/bwmarrin/discordgo"
	"github.com/camd67/moebot/moebot_bot/bot/permissions"
	"github.com/camd67/moebot/moebot_bot/util"
	"github.com/camd67/moebot/moebot_bot/util/db"
	"github.com/camd67/moebot/moebot_bot/util/moeDiscord"
)

const roleCodeSearchText = "[code]"
const roleCodeLength = 6

type RoleCommand struct {
	ComPrefix   string
	PermChecker permissions.PermissionChecker
}

func (rc *RoleCommand) Execute(pack *CommPackage) {
	server, err := db.ServerQueryOrInsert(pack.guild.ID)
	if err != nil {
		pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, there was an error loading server information!")
		return
	}
	var vetRole *discordgo.Role
	if server.VeteranRole.Valid {
		vetRole = moeDiscord.FindRoleById(pack.guild.Roles, server.VeteranRole.String)
	}
	if len(pack.params) == 0 {
		printAllRoles(server, vetRole, pack)
	} else {
		var role *discordgo.Role
		var roleGroup db.RoleGroup
		var dbRole db.Role
		var confirmCodes []string
		// TODO: convert this to use the actual veteran name
		if strings.EqualFold(pack.params[0], "veteran") {
			// before anything, if the server doesn't have a rank or role bail out
			if !server.VeteranRank.Valid || !server.VeteranRole.Valid {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, this server isn't setup to handle veteran role yet! Contact the server admins.")
				return
			}

			// make some placeholder veteran role tables
			roleGroup = db.RoleGroup{
				Name: "Veteran",
				Type: db.GroupTypeExclusive,
			}
			dbRole = db.Role{
				RoleUid: server.VeteranRole.String,
				Trigger: sql.NullString{
					String: "veteran",
					Valid:  true,
				},
			}
			usr, err := db.UserServerRankQuery(pack.message.Author.ID, pack.guild.ID)
			var pointCountMessage string
			if usr != nil {
				pointCountMessage = fmt.Sprintf("%.2f%% of the way to veteran", float64(usr.Rank)/float64(server.VeteranRank.Int64)*100)
			} else {
				pointCountMessage = "Unranked"
			}
			if err != nil {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you don't have enough veteran points yet! You're currently: "+pointCountMessage)
				return
			}
			if int64(usr.Rank) >= server.VeteranRank.Int64 {
				role = vetRole
			} else {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you don't have enough veteran points yet! You're currently: "+pointCountMessage)
				return
			}
		} else {
			// load up the trigger to see if it exists, stripping out anything prefixed with - (our security text)
			var roleNameBuf strings.Builder
			for _, param := range pack.params {
				if !strings.HasPrefix(param, "-") {
					roleNameBuf.WriteString(param)
					roleNameBuf.WriteString(" ")
				} else {
					confirmCodes = append(confirmCodes, param)
				}
			}
			roleNameString := strings.TrimSpace(roleNameBuf.String())
			if len(roleNameString) == 0 {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you must provide a valid role name")
				return
			}
			dbRole, err = db.RoleQueryTrigger(roleNameString, server.Id)
			// an invalid trigger should pretty much never happen, but checking for it anyways
			// however an error may indicate that there were simply no roles in the result set
			if err != nil || !dbRole.Trigger.Valid {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, there was an issue fetching the role, or the role you provided doesn't exist. "+
					"Please provide a valid role. `"+rc.ComPrefix+" role` to list all roles for this server.")
				return
			}
			role = moeDiscord.FindRoleById(pack.guild.Roles, dbRole.RoleUid)
			if role == nil {
				log.Println("Nil dbRole when searching for dbRole id:" + dbRole.RoleUid)
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, there was an issue finding that role in this server. It may have been deleted.")
				return
			}
			roleGroup, err = db.RoleGroupQueryId(dbRole.GroupId)
			if err != nil {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, There was an issue finding that role group. This is an issue with moebot "+
					"and not discord.")
				return
			}
		}
		// process the role to see if it has a confirmation message, then decide if we need to bail out or continue to the role update phase
		if !rc.processRoleConfirmation(dbRole, role, pack, confirmCodes) {
			return
		}

		rc.updateUserRoles(pack, role, roleGroup)
	}
}

/*
Actually go through and update the roles for this user based on the given role and role group
*/
func (rc *RoleCommand) updateUserRoles(pack *CommPackage, role *discordgo.Role, group db.RoleGroup) {
	if util.StrContains(pack.member.Roles, role.ID, util.CaseSensitive) {
		if group.Type == db.GroupTypeExclusiveNoRemove {
			pack.session.ChannelMessageSend(pack.channel.ID, "You've already got that role! You can change roles but can't remove them in the `"+
				group.Name+"` group.")
		} else {
			pack.session.GuildMemberRoleRemove(pack.guild.ID, pack.message.Author.ID, role.ID)
			pack.session.ChannelMessageSend(pack.channel.ID, "Removed role "+role.Name+" for "+pack.message.Author.Mention())
		}
	} else {
		if group.Type == db.GroupTypeAny {
			pack.session.GuildMemberRoleAdd(pack.guild.ID, pack.message.Author.ID, role.ID)
			pack.session.ChannelMessageSend(pack.channel.ID, "Added role "+role.Name+" for "+pack.message.Author.Mention())
		} else {
			// This case needs to check to see if the user has any other roles from this group, since they may not be allowed to add more
			// (GroupTypeExclusive, GroupTypeExclusiveNoRemove)
			fullGroupRoles, err := db.RoleQueryGroup(group.Id)
			if err != nil {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, There was an issue finding that role group. This is an issue with moebot "+
					"and not discord.")
				return
			}
			// we'll always be adding a role here
			pack.session.GuildMemberRoleAdd(pack.guild.ID, pack.message.Author.ID, role.ID)
			var message strings.Builder
			message.WriteString("Added role `")
			message.WriteString(role.Name)
			message.WriteString("` for ")
			message.WriteString(pack.message.Author.Mention())
			// we should only find one other role, but just in case
			foundOtherRole := false
			for _, dbGroupRole := range fullGroupRoles {
				// only send a message that we removed the role if they actually have it and it's not the one we just added
				if dbGroupRole.RoleUid != role.ID && util.StrContains(pack.member.Roles, dbGroupRole.RoleUid, util.CaseSensitive) {
					roleToRemove := moeDiscord.FindRoleById(pack.guild.Roles, dbGroupRole.RoleUid)
					// The user already has this role, remove it and tell them
					if !foundOtherRole {
						message.WriteString("\nAlso removed:")
						foundOtherRole = true
					}
					message.WriteString(" `")
					message.WriteString(roleToRemove.Name)
					message.WriteString("`")
					pack.session.GuildMemberRoleRemove(pack.guild.ID, pack.message.Author.ID, dbGroupRole.RoleUid)
				}
			}
			pack.session.ChannelMessageSend(pack.channel.ID, message.String())
		}
	}
}

func (rc *RoleCommand) processRoleConfirmation(dbRole db.Role, roleToAdd *discordgo.Role, pack *CommPackage, confirmCodes []string) (shouldProceed bool) {
	// we only want to check for a confirmation when we have an actual confirmation message and they don't already have the role
	if dbRole.ConfirmationMessage.Valid && dbRole.ConfirmationMessage.String != "" && !util.StrContains(pack.member.Roles, roleToAdd.ID, util.CaseSensitive) {
		// no confirm codes provided, given them their confirmation code
		if len(confirmCodes) <= 0 {
			err := rc.sendConfirmationMessage(pack.session, pack.channel, dbRole, pack.message.Author)
			if err == nil {
				pack.session.ChannelMessageSend(pack.channel.ID, pack.message.Author.Mention()+" check your PM's for further instructions!")
			} else {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, I couldn't send you a PM! Please check your settings to allow direct messages from "+
					"users on this server.")
			}
			return false
		}

		// they gave us some sort of confirmation code, delete their message
		pack.session.ChannelMessageDelete(pack.channel.ID, pack.message.ID)

		// Optionally check for a security answer, since we can have just a confirmation code and no security
		if dbRole.ConfirmationSecurityAnswer.Valid && dbRole.ConfirmationSecurityAnswer.String != "" {
			if len(confirmCodes) != 2 {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you need to insert a confirmation code and security answer to access "+
					"this role. Use `"+rc.ComPrefix+" "+dbRole.Trigger.String+"` to receive a DM containing detailed instructions.")
				return false
			}
			if !util.StrContains(confirmCodes, dbRole.ConfirmationSecurityAnswer.String, util.CaseSensitive) ||
				!util.StrContains(confirmCodes, "-"+rc.getRoleCode(roleToAdd.ID, pack.message.Author.ID), util.CaseSensitive) {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you need to insert the correct confirmation code to access this role.")
				return false
			}
		} else {
			if len(confirmCodes) != 1 {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you need to insert a confirmation code to access this role. Use `"+
					rc.ComPrefix+" "+dbRole.Trigger.String+"` to receive a DM containing detailed instructions.")
				return false
			}
			if !util.StrContains(confirmCodes, "-"+rc.getRoleCode(roleToAdd.ID, pack.message.Author.ID), util.CaseSensitive) {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you need to insert the correct confirmation code to access this role.")
				return false
			}
		}
	}
	return true
}
func (rc *RoleCommand) sendConfirmationMessage(session *discordgo.Session, channel *discordgo.Channel, role db.Role, user *discordgo.User) error {
	userChannel, err := session.UserChannelCreate(user.ID)
	if err != nil {
		// could log error creating user channel, but seems like it'll clutter the logs for a valid scenario..
		return err
	}
	roleCode := rc.getRoleCode(role.RoleUid, user.ID)
	var messageText string
	if strings.Contains(strings.ToLower(role.ConfirmationMessage.String), roleCodeSearchText) {
		messageText = strings.Replace(role.ConfirmationMessage.String, roleCodeSearchText, roleCode, -1)
	} else {
		messageText = role.ConfirmationMessage.String + "\nYour confirmation code: `-" + roleCode + "`"
	}
	_, err = session.ChannelMessageSend(userChannel.ID, messageText)
	return err
}

/*
Returns a 6 character role code string that is unique per user and per role.
This should NOT be used for security features as it is not a secure algorithm.
*/
func (rc *RoleCommand) getRoleCode(roleUID, userUID string) string {
	hash := sha256.New()
	hash.Write([]byte(roleUID + userUID))
	return string(fmt.Sprintf("%x", hash.Sum(nil))[0:roleCodeLength])
}

func (rc *RoleCommand) GetPermLevel() db.Permission {
	return db.PermAll
}

func (rc *RoleCommand) GetCommandKeys() []string {
	return []string{"ROLE"}
}

func (rc *RoleCommand) GetCommandHelp(commPrefix string) string {
	return fmt.Sprintf("`%[1]s role <role name>` - Changes your role to one of the approved roles. `%[1]s role` to list all the roles", commPrefix)
}
func printAllRoles(server db.Server, vetRole *discordgo.Role, pack *CommPackage) {
	triggersByGroup := make(map[string][]string)
	// go find all the roles for this server
	roles, err := db.RoleQueryServer(server)
	if err != nil {
		pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, there was an issue fetching the server. This is an issue with moebot!")
		return
	}
	// Then find all the groups for the server
	roleGroups, err := db.RoleGroupQueryServer(server)
	if err != nil {
		pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, there was an issue fetching the roles for this server. This is an issue with moebot!")
		return
	}
	if vetRole != nil {
		// triggers = append(triggers, vetRole.Name)
		// TODO: this should be the name of the role, but role is restricted to one word right now...
		triggersByGroup["veteran"] = append(triggersByGroup["veteran"], "veteran")
	}
	for _, role := range roles {
		if !role.Trigger.Valid {
			// skip any invalid triggers. We don't want people thinking that they can choose roles they actually can't
			continue
		}
		// Could maybe make a map here, but the group size is going to be pretty small
		foundGroup := false
		for _, group := range roleGroups {
			if role.GroupId == group.Id {
				triggersByGroup[group.Name] = append(triggersByGroup[group.Name], role.Trigger.String)
				foundGroup = true
			}
		}
		if !foundGroup {
			log.Println("!!! WARNING !!! Failed to find group for a role! This is most likely a programming error")
			triggersByGroup["uncategorized"] = append(triggersByGroup["uncategorized"], role.Trigger.String)
		}
	}
	var message strings.Builder
	if len(triggersByGroup) == 0 {
		message.WriteString("Looks like there aren't any roles I can assign to you in this server!")
	} else {
		message.WriteString("This server's roles (highlighted `like this`): ")
		for groupName, triggerList := range triggersByGroup {
			message.WriteString("\nGroup (")
			message.WriteString(groupName)
			// TODO: add group type string here
			message.WriteString("): `")
			message.WriteString(strings.Join(triggerList, "`, `"))
			message.WriteString("`. ")
		}
	}
	pack.session.ChannelMessageSend(pack.channel.ID, message.String())
}
<commit_msg>Adds better logging for role addition/removal<commit_after>package commands

import (
	"crypto/sha256"
	"database/sql"
	"fmt"
	"log"
	"strings"

	"github.com/bwmarrin/discordgo"
	"github.com/camd67/moebot/moebot_bot/bot/permissions"
	"github.com/camd67/moebot/moebot_bot/util"
	"github.com/camd67/moebot/moebot_bot/util/db"
	"github.com/camd67/moebot/moebot_bot/util/moeDiscord"
)

const roleCodeSearchText = "[code]"
const roleCodeLength = 6

type RoleCommand struct {
	ComPrefix   string
	PermChecker permissions.PermissionChecker
}

func (rc *RoleCommand) Execute(pack *CommPackage) {
	server, err := db.ServerQueryOrInsert(pack.guild.ID)
	if err != nil {
		pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, there was an error loading server information!")
		return
	}
	var vetRole *discordgo.Role
	if server.VeteranRole.Valid {
		vetRole = moeDiscord.FindRoleById(pack.guild.Roles, server.VeteranRole.String)
	}
	if len(pack.params) == 0 {
		printAllRoles(server, vetRole, pack)
	} else {
		var role *discordgo.Role
		var roleGroup db.RoleGroup
		var dbRole db.Role
		var confirmCodes []string
		// TODO: convert this to use the actual veteran name
		if strings.EqualFold(pack.params[0], "veteran") {
			// before anything, if the server doesn't have a rank or role bail out
			if !server.VeteranRank.Valid || !server.VeteranRole.Valid {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, this server isn't setup to handle veteran role yet! Contact the server admins.")
				return
			}

			// make some placeholder veteran role tables
			roleGroup = db.RoleGroup{
				Name: "Veteran",
				Type: db.GroupTypeExclusive,
			}
			dbRole = db.Role{
				RoleUid: server.VeteranRole.String,
				Trigger: sql.NullString{
					String: "veteran",
					Valid:  true,
				},
			}
			usr, err := db.UserServerRankQuery(pack.message.Author.ID, pack.guild.ID)
			var pointCountMessage string
			if usr != nil {
				pointCountMessage = fmt.Sprintf("%.2f%% of the way to veteran", float64(usr.Rank)/float64(server.VeteranRank.Int64)*100)
			} else {
				pointCountMessage = "Unranked"
			}
			if err != nil {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you don't have enough veteran points yet! You're currently: "+pointCountMessage)
				return
			}
			if int64(usr.Rank) >= server.VeteranRank.Int64 {
				role = vetRole
			} else {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you don't have enough veteran points yet! You're currently: "+pointCountMessage)
				return
			}
		} else {
			// load up the trigger to see if it exists, stripping out anything prefixed with - (our security text)
			var roleNameBuf strings.Builder
			for _, param := range pack.params {
				if !strings.HasPrefix(param, "-") {
					roleNameBuf.WriteString(param)
					roleNameBuf.WriteString(" ")
				} else {
					confirmCodes = append(confirmCodes, param)
				}
			}
			roleNameString := strings.TrimSpace(roleNameBuf.String())
			if len(roleNameString) == 0 {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you must provide a valid role name")
				return
			}
			dbRole, err = db.RoleQueryTrigger(roleNameString, server.Id)
			// an invalid trigger should pretty much never happen, but checking for it anyways
			// however an error may indicate that there were simply no roles in the result set
			if err != nil || !dbRole.Trigger.Valid {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, there was an issue fetching the role, or the role you provided doesn't exist. "+
					"Please provide a valid role. `"+rc.ComPrefix+" role` to list all roles for this server.")
				return
			}
			role = moeDiscord.FindRoleById(pack.guild.Roles, dbRole.RoleUid)
			if role == nil {
				log.Println("Nil dbRole when searching for dbRole id:" + dbRole.RoleUid)
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, there was an issue finding that role in this server. It may have been deleted.")
				return
			}
			roleGroup, err = db.RoleGroupQueryId(dbRole.GroupId)
			if err != nil {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, There was an issue finding that role group. This is an issue with moebot "+
					"and not discord.")
				return
			}
		}
		// process the role to see if it has a confirmation message, then decide if we need to bail out or continue to the role update phase
		if !rc.processRoleConfirmation(dbRole, role, pack, confirmCodes) {
			return
		}

		rc.updateUserRoles(pack, role, roleGroup)
	}
}

/*
Actually go through and update the roles for this user based on the given role and role group
*/
func (rc *RoleCommand) updateUserRoles(pack *CommPackage, role *discordgo.Role, group db.RoleGroup) {
	if util.StrContains(pack.member.Roles, role.ID, util.CaseSensitive) {
		if group.Type == db.GroupTypeExclusiveNoRemove {
			pack.session.ChannelMessageSend(pack.channel.ID, "You've already got that role! You can change roles but can't remove them in the `"+
				group.Name+"` group.")
		} else {
			pack.session.GuildMemberRoleRemove(pack.guild.ID, pack.message.Author.ID, role.ID)
			pack.session.ChannelMessageSend(pack.channel.ID, "Removed role "+role.Name+" for "+pack.message.Author.Mention())
		}
	} else {
		if group.Type == db.GroupTypeAny {
			pack.session.GuildMemberRoleAdd(pack.guild.ID, pack.message.Author.ID, role.ID)
			pack.session.ChannelMessageSend(pack.channel.ID, "Added role "+role.Name+" for "+pack.message.Author.Mention())
		} else {
			// This case needs to check to see if the user has any other roles from this group, since they may not be allowed to add more
			// (GroupTypeExclusive, GroupTypeExclusiveNoRemove)
			fullGroupRoles, err := db.RoleQueryGroup(group.Id)
			if err != nil {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, There was an issue finding that role group. This is an issue with moebot "+
					"and not discord.")
				return
			}
			// we'll always be adding a role here
			err = pack.session.GuildMemberRoleAdd(pack.guild.ID, pack.message.Author.ID, role.ID)
			if err != nil {
				// We can add better handling here for what kind of error we got and notify the server appropriately
				log.Println("error adding role to user "+pack.message.Author.ID+" with role "+role.ID+" with error: ", err)
				pack.session.ChannelMessageSend(pack.message.ChannelID, "Sorry there was an error adding your new role! Make sure moebot has "+
					"permission to edit roles")
				return
			}
			var message strings.Builder
			message.WriteString("Added role `")
			message.WriteString(role.Name)
			message.WriteString("` for ")
			message.WriteString(pack.message.Author.Mention())
			// we should only find one other role, but just in case
			foundOtherRole := false
			for _, dbGroupRole := range fullGroupRoles {
				// only send a message that we removed the role if they actually have it and it's not the one we just added
				if dbGroupRole.RoleUid != role.ID && util.StrContains(pack.member.Roles, dbGroupRole.RoleUid, util.CaseSensitive) {
					roleToRemove := moeDiscord.FindRoleById(pack.guild.Roles, dbGroupRole.RoleUid)

					// Check for an error first
					err = pack.session.GuildMemberRoleRemove(pack.guild.ID, pack.message.Author.ID, dbGroupRole.RoleUid)
					if err != nil {
						message.WriteString("\nFailed to remove: `")
						message.WriteString(roleToRemove.Name)
						message.WriteString("`")
						log.Println("error removing role from "+pack.message.Author.ID+" with role "+role.ID+" with error: ", err)
						continue
					}

					// The user already has this role, remove it and tell them
					if !foundOtherRole {
						message.WriteString("\nAlso removed:")
						foundOtherRole = true
					}
					message.WriteString(" `")
					message.WriteString(roleToRemove.Name)
					message.WriteString("`")
				}
			}
			pack.session.ChannelMessageSend(pack.channel.ID, message.String())
		}
	}
}

func (rc *RoleCommand) processRoleConfirmation(dbRole db.Role, roleToAdd *discordgo.Role, pack *CommPackage, confirmCodes []string) (shouldProceed bool) {
	// we only want to check for a confirmation when we have an actual confirmation message and they don't already have the role
	if dbRole.ConfirmationMessage.Valid && dbRole.ConfirmationMessage.String != "" && !util.StrContains(pack.member.Roles, roleToAdd.ID, util.CaseSensitive) {
		// no confirm codes provided, given them their confirmation code
		if len(confirmCodes) <= 0 {
			err := rc.sendConfirmationMessage(pack.session, pack.channel, dbRole, pack.message.Author)
			if err == nil {
				pack.session.ChannelMessageSend(pack.channel.ID, pack.message.Author.Mention()+" check your PM's for further instructions!")
			} else {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, I couldn't send you a PM! Please check your settings to allow direct messages from "+
					"users on this server.")
			}
			return false
		}

		// they gave us some sort of confirmation code, delete their message
		pack.session.ChannelMessageDelete(pack.channel.ID, pack.message.ID)

		// Optionally check for a security answer, since we can have just a confirmation code and no security
		if dbRole.ConfirmationSecurityAnswer.Valid && dbRole.ConfirmationSecurityAnswer.String != "" {
			if len(confirmCodes) != 2 {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you need to insert a confirmation code and security answer to access "+
					"this role. Use `"+rc.ComPrefix+" "+dbRole.Trigger.String+"` to receive a DM containing detailed instructions.")
				return false
			}
			if !util.StrContains(confirmCodes, dbRole.ConfirmationSecurityAnswer.String, util.CaseSensitive) ||
				!util.StrContains(confirmCodes, "-"+rc.getRoleCode(roleToAdd.ID, pack.message.Author.ID), util.CaseSensitive) {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you need to insert the correct confirmation code to access this role.")
				return false
			}
		} else {
			if len(confirmCodes) != 1 {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you need to insert a confirmation code to access this role. Use `"+
					rc.ComPrefix+" "+dbRole.Trigger.String+"` to receive a DM containing detailed instructions.")
				return false
			}
			if !util.StrContains(confirmCodes, "-"+rc.getRoleCode(roleToAdd.ID, pack.message.Author.ID), util.CaseSensitive) {
				pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, you need to insert the correct confirmation code to access this role.")
				return false
			}
		}
	}
	return true
}
func (rc *RoleCommand) sendConfirmationMessage(session *discordgo.Session, channel *discordgo.Channel, role db.Role, user *discordgo.User) error {
	userChannel, err := session.UserChannelCreate(user.ID)
	if err != nil {
		// could log error creating user channel, but seems like it'll clutter the logs for a valid scenario..
		return err
	}
	roleCode := rc.getRoleCode(role.RoleUid, user.ID)
	var messageText string
	if strings.Contains(strings.ToLower(role.ConfirmationMessage.String), roleCodeSearchText) {
		messageText = strings.Replace(role.ConfirmationMessage.String, roleCodeSearchText, roleCode, -1)
	} else {
		messageText = role.ConfirmationMessage.String + "\nYour confirmation code: `-" + roleCode + "`"
	}
	_, err = session.ChannelMessageSend(userChannel.ID, messageText)
	return err
}

/*
Returns a 6 character role code string that is unique per user and per role.
This should NOT be used for security features as it is not a secure algorithm.
*/
func (rc *RoleCommand) getRoleCode(roleUID, userUID string) string {
	hash := sha256.New()
	hash.Write([]byte(roleUID + userUID))
	return string(fmt.Sprintf("%x", hash.Sum(nil))[0:roleCodeLength])
}

func (rc *RoleCommand) GetPermLevel() db.Permission {
	return db.PermAll
}

func (rc *RoleCommand) GetCommandKeys() []string {
	return []string{"ROLE"}
}

func (rc *RoleCommand) GetCommandHelp(commPrefix string) string {
	return fmt.Sprintf("`%[1]s role <role name>` - Changes your role to one of the approved roles. `%[1]s role` to list all the roles", commPrefix)
}
func printAllRoles(server db.Server, vetRole *discordgo.Role, pack *CommPackage) {
	triggersByGroup := make(map[string][]string)
	// go find all the roles for this server
	roles, err := db.RoleQueryServer(server)
	if err != nil {
		pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, there was an issue fetching the server. This is an issue with moebot!")
		return
	}
	// Then find all the groups for the server
	roleGroups, err := db.RoleGroupQueryServer(server)
	if err != nil {
		pack.session.ChannelMessageSend(pack.channel.ID, "Sorry, there was an issue fetching the roles for this server. This is an issue with moebot!")
		return
	}
	if vetRole != nil {
		// triggers = append(triggers, vetRole.Name)
		// TODO: this should be the name of the role, but role is restricted to one word right now...
		triggersByGroup["veteran"] = append(triggersByGroup["veteran"], "veteran")
	}
	for _, role := range roles {
		if !role.Trigger.Valid {
			// skip any invalid triggers. We don't want people thinking that they can choose roles they actually can't
			continue
		}
		// Could maybe make a map here, but the group size is going to be pretty small
		foundGroup := false
		for _, group := range roleGroups {
			if role.GroupId == group.Id {
				triggersByGroup[group.Name] = append(triggersByGroup[group.Name], role.Trigger.String)
				foundGroup = true
			}
		}
		if !foundGroup {
			log.Println("!!! WARNING !!! Failed to find group for a role! This is most likely a programming error")
			triggersByGroup["uncategorized"] = append(triggersByGroup["uncategorized"], role.Trigger.String)
		}
	}
	var message strings.Builder
	if len(triggersByGroup) == 0 {
		message.WriteString("Looks like there aren't any roles I can assign to you in this server!")
	} else {
		message.WriteString("This server's roles (highlighted `like this`): ")
		for groupName, triggerList := range triggersByGroup {
			message.WriteString("\nGroup (")
			message.WriteString(groupName)
			// TODO: add group type string here
			message.WriteString("): `")
			message.WriteString(strings.Join(triggerList, "`, `"))
			message.WriteString("`. ")
		}
	}
	pack.session.ChannelMessageSend(pack.channel.ID, message.String())
}
<|endoftext|>


<commit_before>package jsonrpc1

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"net/http"
	"net/rpc"
	"strings"
)

const contentType = "application/json"

type contextKey int

var httpRequestContextKey contextKey = 0

// HTTPRequestFromContext returns HTTP request related to this RPC (if
// you use HTTPHander to serve JSON RPC 1.0 over HTTP) or nil otherwise.
func HTTPRequestFromContext(ctx context.Context) *http.Request {
	req, _ := ctx.Value(httpRequestContextKey).(*http.Request)
	return req
}

type httpServerConn struct {
	req     io.Reader
	res     io.Writer
	replied bool
}

func (conn *httpServerConn) Read(buf []byte) (int, error) {
	return conn.req.Read(buf)
}

func (conn *httpServerConn) Write(buf []byte) (int, error) {
	conn.replied = true
	return conn.res.Write(buf)
}

func (conn *httpServerConn) Close() error {
	return nil
}

type httpHandler struct {
	rpc *rpc.Server
}

// HTTPHandler returns handler for HTTP requests which will execute
// incoming JSON-RPC 1.0 over HTTP using srv.
//
// If srv is nil then rpc.DefaultServer will be used.
//
// Specification: http://www.simple-is-better.org/json-rpc/transport_http.html
func HTTPHandler(srv *rpc.Server) http.Handler {
	if srv == nil {
		srv = rpc.DefaultServer
	}
	return &httpHandler{srv}
}

func (h *httpHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", contentType)

	if req.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	mediaType, _, _ := mime.ParseMediaType(req.Header.Get("Content-Type"))
	if mediaType != contentType || req.Header.Get("Accept") != contentType {
		w.WriteHeader(http.StatusUnsupportedMediaType)
		return
	}

	ctx := context.WithValue(context.Background(), httpRequestContextKey, req)
	conn := &httpServerConn{req: req.Body, res: w}
	h.rpc.ServeRequest(NewServerCodecContext(ctx, conn, h.rpc))
	if !conn.replied {
		w.WriteHeader(http.StatusNoContent)
	}
}

// Doer is an interface for doing HTTP requests.
type Doer interface {
	Do(req *http.Request) (resp *http.Response, err error)
}

// The DoerFunc type is an adapter to allow the use of ordinary
// functions as HTTP clients. If f is a function with the appropriate
// signature, DoerFunc(f) is a client that calls f.
type DoerFunc func(req *http.Request) (resp *http.Response, err error)

// DoerFunc calls f(req).
func (f DoerFunc) Do(req *http.Request) (resp *http.Response, err error) {
	return f(req)
}

type httpClientConn struct {
	url   string
	doer  Doer
	ready chan io.ReadCloser
	body  io.ReadCloser
}

func (conn *httpClientConn) Read(buf []byte) (int, error) {
	if conn.body == nil {
		conn.body = <-conn.ready
	}
	n, err := conn.body.Read(buf)
	if err == io.EOF || strings.Contains(string(buf), "result") {
		conn.body.Close()
		conn.body = nil
		err = nil
		if n == 0 {
			return conn.Read(buf)
		}
	}
	return n, err
}

func (conn *httpClientConn) Write(buf []byte) (int, error) {
	b := make([]byte, len(buf))
	copy(b, buf)
	go func() {
		req, err := http.NewRequest("POST", conn.url, bytes.NewReader(b))
		if err == nil {
			req.Header.Add("Content-Type", contentType)
			req.Header.Add("Accept", contentType)
			var resp *http.Response
			resp, err = conn.doer.Do(req)
			const maxBodySlurpSize = 32 * 1024
			if err != nil {
			} else if resp.Header.Get("Content-Type") != contentType {
				err = fmt.Errorf("bad HTTP Content-Type: %s", resp.Header.Get("Content-Type"))
			} else if resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusInternalServerError {
				conn.ready <- resp.Body
				return
			} else if resp.StatusCode == http.StatusNoContent || resp.StatusCode == http.StatusAccepted {
				// Read the body if small so underlying TCP connection will be re-used.
				// No need to check for errors: if it fails, Transport won't reuse it anyway.
				if resp.ContentLength == -1 || resp.ContentLength <= maxBodySlurpSize {
					io.CopyN(ioutil.Discard, resp.Body, maxBodySlurpSize)
				}
				resp.Body.Close()
				return
			} else {
				err = fmt.Errorf("bad HTTP Status: %s", resp.Status)
			}
			if resp != nil {
				// Read the body if small so underlying TCP connection will be re-used.
				// No need to check for errors: if it fails, Transport won't reuse it anyway.
				if resp.ContentLength == -1 || resp.ContentLength <= maxBodySlurpSize {
					io.CopyN(ioutil.Discard, resp.Body, maxBodySlurpSize)
				}
				resp.Body.Close()
			}
		}
		var res clientResponse
		if json.Unmarshal(b, &res) == nil && res.ID == nil {
			return // ignore error from Notification
		}
		res.Error = NewError(errInternal.Code, err.Error())
		buf := &bytes.Buffer{}
		json.NewEncoder(buf).Encode(res)
		conn.ready <- ioutil.NopCloser(buf)
	}()
	return len(buf), nil
}

func (conn *httpClientConn) Close() error {
	return nil
}

// NewHTTPClient returns a new Client to handle requests to the
// set of services at the given url.
func NewHTTPClient(url string) *Client {
	return NewCustomHTTPClient(url, nil)
}

// NewCustomHTTPClient returns a new Client to handle requests to the
// set of services at the given url using provided doer (&http.Client{} by
// default).
//
// Use doer to customize HTTP authorization/headers/etc. sent with each
// request (it method Do() will receive already configured POST request
// with url, all required headers and body set according to specification).
func NewCustomHTTPClient(url string, doer Doer) *Client {
	if doer == nil {
		doer = &http.Client{}
	}
	return NewClient(&httpClientConn{
		url:   url,
		doer:  doer,
		ready: make(chan io.ReadCloser, 16),
	})
}
<commit_msg>Fix response rpc bug<commit_after>package jsonrpc1

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"net/http"
	"net/rpc"
	"strings"
)

const contentType = "application/json"

type contextKey int

var httpRequestContextKey contextKey = 0

// HTTPRequestFromContext returns HTTP request related to this RPC (if
// you use HTTPHander to serve JSON RPC 1.0 over HTTP) or nil otherwise.
func HTTPRequestFromContext(ctx context.Context) *http.Request {
	req, _ := ctx.Value(httpRequestContextKey).(*http.Request)
	return req
}

type httpServerConn struct {
	req     io.Reader
	res     io.Writer
	replied bool
}

func (conn *httpServerConn) Read(buf []byte) (int, error) {
	return conn.req.Read(buf)
}

func (conn *httpServerConn) Write(buf []byte) (int, error) {
	conn.replied = true
	return conn.res.Write(buf)
}

func (conn *httpServerConn) Close() error {
	return nil
}

type httpHandler struct {
	rpc *rpc.Server
}

// HTTPHandler returns handler for HTTP requests which will execute
// incoming JSON-RPC 1.0 over HTTP using srv.
//
// If srv is nil then rpc.DefaultServer will be used.
//
// Specification: http://www.simple-is-better.org/json-rpc/transport_http.html
func HTTPHandler(srv *rpc.Server) http.Handler {
	if srv == nil {
		srv = rpc.DefaultServer
	}
	return &httpHandler{srv}
}

func (h *httpHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", contentType)

	if req.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	mediaType, _, _ := mime.ParseMediaType(req.Header.Get("Content-Type"))
	if mediaType != contentType || req.Header.Get("Accept") != contentType {
		w.WriteHeader(http.StatusUnsupportedMediaType)
		return
	}

	ctx := context.WithValue(context.Background(), httpRequestContextKey, req)
	conn := &httpServerConn{req: req.Body, res: w}
	h.rpc.ServeRequest(NewServerCodecContext(ctx, conn, h.rpc))
	if !conn.replied {
		w.WriteHeader(http.StatusNoContent)
	}
}

// Doer is an interface for doing HTTP requests.
type Doer interface {
	Do(req *http.Request) (resp *http.Response, err error)
}

// The DoerFunc type is an adapter to allow the use of ordinary
// functions as HTTP clients. If f is a function with the appropriate
// signature, DoerFunc(f) is a client that calls f.
type DoerFunc func(req *http.Request) (resp *http.Response, err error)

// DoerFunc calls f(req).
func (f DoerFunc) Do(req *http.Request) (resp *http.Response, err error) {
	return f(req)
}

type httpClientConn struct {
	url   string
	doer  Doer
	ready chan io.ReadCloser
	body  io.ReadCloser
}

func (conn *httpClientConn) Read(buf []byte) (int, error) {
	if conn.body == nil {
		conn.body = <-conn.ready
	}
	n, err := conn.body.Read(buf)
	if err == io.EOF {
		conn.body.Close()
		conn.body = nil
		err = nil
		if n == 0 {
			return conn.Read(buf)
		}
	}
	return n, err
}

func (conn *httpClientConn) Write(buf []byte) (int, error) {
	b := make([]byte, len(buf))
	copy(b, buf)
	go func() {
		req, err := http.NewRequest("POST", conn.url, bytes.NewReader(b))
		if err == nil {
			req.Header.Add("Content-Type", contentType)
			req.Header.Add("Accept", contentType)
			var resp *http.Response
			resp, err = conn.doer.Do(req)
			const maxBodySlurpSize = 32 * 1024
			if err != nil {
			} else if resp.Header.Get("Content-Type") != contentType {
				err = fmt.Errorf("bad HTTP Content-Type: %s", resp.Header.Get("Content-Type"))
			} else if resp.StatusCode == http.StatusOK {
				conn.ready <- resp.Body
				return
			} else if resp.StatusCode == http.StatusNoContent || resp.StatusCode == http.StatusAccepted {
				// Read the body if small so underlying TCP connection will be re-used.
				// No need to check for errors: if it fails, Transport won't reuse it anyway.
				if resp.ContentLength == -1 || resp.ContentLength <= maxBodySlurpSize {
					io.CopyN(ioutil.Discard, resp.Body, maxBodySlurpSize)
				}
				resp.Body.Close()
				return
			} else {
				err = fmt.Errorf("bad HTTP Status: %s", resp.Status)
			}
			if resp != nil {
				// Read the body if small so underlying TCP connection will be re-used.
				// No need to check for errors: if it fails, Transport won't reuse it anyway.
				if resp.ContentLength == -1 || resp.ContentLength <= maxBodySlurpSize {
					io.CopyN(ioutil.Discard, resp.Body, maxBodySlurpSize)
				}
				resp.Body.Close()
			}
		}
		var res clientResponse
		if json.Unmarshal(b, &res) == nil && res.ID == nil {
			return // ignore error from Notification
		}
		res.Error = NewError(errInternal.Code, err.Error())
		buf := &bytes.Buffer{}
		json.NewEncoder(buf).Encode(res)
		conn.ready <- ioutil.NopCloser(buf)
	}()
	return len(buf), nil
}

func (conn *httpClientConn) Close() error {
	return nil
}

// NewHTTPClient returns a new Client to handle requests to the
// set of services at the given url.
func NewHTTPClient(url string) *Client {
	return NewCustomHTTPClient(url, nil)
}

// NewCustomHTTPClient returns a new Client to handle requests to the
// set of services at the given url using provided doer (&http.Client{} by
// default).
//
// Use doer to customize HTTP authorization/headers/etc. sent with each
// request (it method Do() will receive already configured POST request
// with url, all required headers and body set according to specification).
func NewCustomHTTPClient(url string, doer Doer) *Client {
	if doer == nil {
		doer = &http.Client{}
	}
	return NewClient(&httpClientConn{
		url:   url,
		doer:  doer,
		ready: make(chan io.ReadCloser, 16),
	})
}
<|endoftext|>