<commit_before>// +build k8srequired

package integration

import (
	"fmt"
	"html/template"
	"log"
	"os"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"

	"github.com/giantswarm/aws-operator/service/keyv2"
	"github.com/giantswarm/microerror"
)

const (
	awsOperatorValuesFile  = "/tmp/aws-operator-values.yaml"
	awsOperatorChartValues = `Installation:
  V1:
    Name: ci-awsop
    Provider:
      AWS:
        Region: ${AWS_REGION}
    Secret:
      AWSOperator:
        IDRSAPub: ${IDRSA_PUB}
        SecretYaml: |
          service:
            aws:
              accesskey:
                id: ${AWS_ACCESS_KEY_ID}
                secret: ${AWS_SECRET_ACCESS_KEY}
                token: ${AWS_SESSION_TOKEN}
              hostaccesskey:
                id: ""
                secret: ""
      Registry:
        PullSecret:
          DockerConfigJSON: "{\"auths\":{\"quay.io\":{\"auth\":\"${REGISTRY_PULL_SECRET}\"}}}"
`
	awsResourceValuesFile  = "/tmp/aws-operator-values.yaml"
	awsResourceChartValues = `commonDomain: ${COMMON_DOMAIN}
clusterName: ${CLUSTER_NAME}
clusterVersion: v_0_1_0
sshPublicKey: ${IDRSA_PUB}
versionBundleVersion: ${VERSION_BUNDLE_VERSION}
aws:
  networkCIDR: "{{.NetworkCIDR}}"
  privateSubnetCIDR: "{{.PrivateSubnetCIDR}}"
  publicSubnetCIDR: "{{.PublicSubnetCIDR}}"
  region: ${AWS_REGION}
  apiHostedZone: ${AWS_API_HOSTED_ZONE}
  ingressHostedZone: ${AWS_INGRESS_HOSTED_ZONE}
  routeTable0: ${AWS_ROUTE_TABLE_0}
  routeTable1: ${AWS_ROUTE_TABLE_1}
  vpcPeerId: ${AWS_VPC_PEER_ID}
`
)

type aWSClient struct {
	EC2 *ec2.EC2
}

func newAWSClient() aWSClient {
	awsCfg := &aws.Config{
		Credentials: credentials.NewStaticCredentials(
			os.Getenv("AWS_ACCESS_KEY_ID"),
			os.Getenv("AWS_SECRET_ACCESS_KEY"),
			os.Getenv("AWS_SESSION_TOKEN")),
		Region: aws.String(os.Getenv("AWS_REGION")),
	}
	s := session.New(awsCfg)
	clients := aWSClient{
		EC2: ec2.New(s),
	}

	return clients
}

var (
	f *framework
	c aWSClient
)

// TestMain allows us to have common setup and teardown steps that are run
// once for all the tests https://golang.org/pkg/testing/#hdr-Main.
func TestMain(m *testing.M) {
	var v int
	var err error
	f, err = newFramework()
	if err != nil {
		log.Printf("unexpected error: %v\n", err)
		os.Exit(1)
	}

	c = newAWSClient()

	if err := f.SetUp(); err != nil {
		log.Printf("unexpected error: %v\n", err)
		v = 1
	}

	if err := operatorSetup(); err != nil {
		log.Printf("unexpected error: %v\n", err)
		v = 1
	}

	if v == 0 {
		v = m.Run()
	}

	f.DeleteGuestCluster()
	operatorTearDown()
	f.TearDown()

	os.Exit(v)
}

func TestGuestReadyAfterMasterReboot(t *testing.T) {
	log.Println("getting master ID")
	describeInput := &ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name: aws.String("tag:Name"),
				Values: []*string{
					aws.String(fmt.Sprintf("%s-master", os.Getenv("CLUSTER_NAME"))),
				},
			},
			&ec2.Filter{
				Name: aws.String("instance-state-name"),
				Values: []*string{
					aws.String("running"),
				},
			},
		},
	}
	res, err := c.EC2.DescribeInstances(describeInput)
	if err != nil {
		t.Errorf("unexpected error getting master id %v", err)
	}
	if len(res.Reservations) != 1 {
		t.Errorf("unexpected number of reservations %d", len(res.Reservations))
	}
	if len(res.Reservations[0].Instances) != 1 {
		t.Errorf("unexpected number of instances %d", len(res.Reservations[0].Instances))
	}
	masterID := res.Reservations[0].Instances[0].InstanceId

	log.Println("rebooting master")
	rebootInput := &ec2.RebootInstancesInput{
		InstanceIds: []*string{
			masterID,
		},
	}
	_, err = c.EC2.RebootInstances(rebootInput)
	if err != nil {
		t.Errorf("unexpected error rebooting  master %v", err)
	}

	if err := f.WaitForAPIDown(); err != nil {
		t.Errorf("unexpected error waiting for master shutting down %v", err)
	}

	if err := f.WaitForGuestReady(); err != nil {
		t.Errorf("unexpected error waiting for guest cluster ready, %v", err)
	}
}

func TestWorkersScaling(t *testing.T) {
	currentWorkers, err := numberOfWorkers(os.Getenv("CLUSTER_NAME"))
	if err != nil {
		t.Fatalf("unexpected error getting number of workers %v", err)
	}
	currentMasters, err := numberOfMasters(os.Getenv("CLUSTER_NAME"))
	if err != nil {
		t.Fatalf("unexpected error getting number of masters %v", err)
	}

	// increase number of workers
	expectedWorkers := currentWorkers + 1
	log.Printf("Increasing the number of workers to %d", expectedWorkers)
	err = addWorker(os.Getenv("CLUSTER_NAME"))
	if err != nil {
		t.Fatalf("unexpected error setting number of workers to %d, %v", expectedWorkers, err)
	}

	if err := f.WaitForNodesUp(currentMasters + expectedWorkers); err != nil {
		t.Fatalf("unexpected error waiting for %d nodes up, %v", expectedWorkers, err)
	}
	log.Printf("%d worker nodes ready", expectedWorkers)

	// decrease number of workers
	expectedWorkers--
	log.Printf("Decreasing the number of workers to %d", expectedWorkers)
	err = removeWorker(os.Getenv("CLUSTER_NAME"))
	if err != nil {
		t.Fatalf("unexpected error setting number of workers to %d, %v", expectedWorkers, err)
	}

	if err := f.WaitForNodesUp(currentMasters + expectedWorkers); err != nil {
		t.Fatalf("unexpected error waiting for %d nodes up, %v", expectedWorkers, err)
	}
	log.Printf("%d worker nodes ready", expectedWorkers)
}

func operatorSetup() error {
	if err := f.InstallCertOperator(); err != nil {
		return microerror.Mask(err)
	}

	if err := f.InstallCertResource(); err != nil {
		return microerror.Mask(err)
	}

	if err := f.InstallAwsOperator(); err != nil {
		return microerror.Mask(err)
	}

	if err := writeAWSResourceValues(); err != nil {
		return microerror.Maskf(err, "unexpected error writing aws-resource-lab values file")
	}

	if err := runCmd("helm registry install quay.io/giantswarm/aws-resource-lab-chart:stable -- -n aws-resource-lab --values " + awsOperatorValuesFile); err != nil {
		return microerror.Maskf(err, "unexpected error installing aws-resource-lab chart: %v")
	}

	logEntry := "cluster '${CLUSTER_NAME}' processed"
	if os.Getenv("VERSION_BUNDLE_VERSION") == "0.2.0" {
		logEntry = "creating AWS cloudformation stack: created"
	}

	operatorPodName, err := f.PodName("giantswarm", "app=aws-operator")
	if err != nil {
		return microerror.Maskf(err, "unexpected error getting operator pod name: %v")
	}

	if err := f.WaitForPodLog("giantswarm", logEntry, operatorPodName); err != nil {
		return microerror.Maskf(err, "unexpected error waiting for guest cluster installed: %v")
	}

	if err := f.WaitForGuestReady(); err != nil {
		return microerror.Maskf(err, "unexpected error waiting for guest cluster ready")
	}

	return nil
}

func operatorTearDown() {
	runCmd("helm delete cert-resource-lab --purge")
	runCmd("helm delete cert-operator --purge")
	runCmd("helm delete aws-resource-lab --purge")
	runCmd("helm delete aws-operator --purge")
}

func writeAWSResourceValues() error {
	awsResourceChartValuesEnv := os.ExpandEnv(awsResourceChartValues)

	tmpl, err := template.New("awsResource").Parse(awsResourceChartValuesEnv)
	if err != nil {
		return microerror.Mask(err)
	}

	f, err := os.Create(awsResourceValuesFile)
	if err != nil {
		return microerror.Mask(err)
	}
	defer f.Close()

	vpc, err := newAWSVPCBlock(c)
	if err != nil {
		return microerror.Mask(err)
	}

	err = tmpl.Execute(f, vpc)
	if err != nil {
		return microerror.Mask(err)
	}

	return nil
}

func numberOfWorkers(clusterName string) (int, error) {
	cluster, err := f.AWSCluster(clusterName)
	if err != nil {
		return 0, microerror.Mask(err)
	}

	return keyv2.WorkerCount(cluster), nil
}

func numberOfMasters(clusterName string) (int, error) {
	cluster, err := f.AWSCluster(clusterName)
	if err != nil {
		return 0, microerror.Mask(err)
	}

	return keyv2.MasterCount(cluster), nil
}

func addWorker(clusterName string) error {
	cluster, err := f.AWSCluster(clusterName)
	if err != nil {
		return microerror.Mask(err)
	}

	newWorker := cluster.Spec.AWS.Workers[0]

	patch := make([]PatchSpec, 1)
	patch[0].Op = "add"
	patch[0].Path = "/spec/aws/workers/-"
	patch[0].Value = newWorker

	return f.ApplyAWSConfigPatch(patch, clusterName)
}

func removeWorker(clusterName string) error {
	patch := make([]PatchSpec, 1)
	patch[0].Op = "remove"
	patch[0].Path = "/spec/aws/workers/1"

	return f.ApplyAWSConfigPatch(patch, clusterName)
}
<commit_msg>fix e2e build (#673)<commit_after>// +build k8srequired

package integration

import (
	"fmt"
	"html/template"
	"log"
	"os"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"

	"github.com/giantswarm/aws-operator/service/keyv2"
	"github.com/giantswarm/microerror"
)

const (
	awsOperatorValuesFile  = "/tmp/aws-operator-values.yaml"
	awsOperatorChartValues = `Installation:
  V1:
    Name: ci-awsop
    Provider:
      AWS:
        Region: ${AWS_REGION}
    Secret:
      AWSOperator:
        IDRSAPub: ${IDRSA_PUB}
        SecretYaml: |
          service:
            aws:
              accesskey:
                id: ${AWS_ACCESS_KEY_ID}
                secret: ${AWS_SECRET_ACCESS_KEY}
                token: ${AWS_SESSION_TOKEN}
              hostaccesskey:
                id: ""
                secret: ""
      Registry:
        PullSecret:
          DockerConfigJSON: "{\"auths\":{\"quay.io\":{\"auth\":\"${REGISTRY_PULL_SECRET}\"}}}"
`
	awsResourceValuesFile  = "/tmp/aws-operator-values.yaml"
	awsResourceChartValues = `commonDomain: ${COMMON_DOMAIN}
clusterName: ${CLUSTER_NAME}
clusterVersion: v_0_1_0
sshPublicKey: ${IDRSA_PUB}
versionBundleVersion: ${VERSION_BUNDLE_VERSION}
aws:
  networkCIDR: "{{.NetworkCIDR}}"
  privateSubnetCIDR: "{{.PrivateSubnetCIDR}}"
  publicSubnetCIDR: "{{.PublicSubnetCIDR}}"
  region: ${AWS_REGION}
  apiHostedZone: ${AWS_API_HOSTED_ZONE}
  ingressHostedZone: ${AWS_INGRESS_HOSTED_ZONE}
  routeTable0: ${AWS_ROUTE_TABLE_0}
  routeTable1: ${AWS_ROUTE_TABLE_1}
  vpcPeerId: ${AWS_VPC_PEER_ID}
`
)

type aWSClient struct {
	EC2 *ec2.EC2
}

func newAWSClient() aWSClient {
	awsCfg := &aws.Config{
		Credentials: credentials.NewStaticCredentials(
			os.Getenv("AWS_ACCESS_KEY_ID"),
			os.Getenv("AWS_SECRET_ACCESS_KEY"),
			os.Getenv("AWS_SESSION_TOKEN")),
		Region: aws.String(os.Getenv("AWS_REGION")),
	}
	s := session.New(awsCfg)
	clients := aWSClient{
		EC2: ec2.New(s),
	}

	return clients
}

var (
	f *framework
	c aWSClient
)

// TestMain allows us to have common setup and teardown steps that are run
// once for all the tests https://golang.org/pkg/testing/#hdr-Main.
func TestMain(m *testing.M) {
	var v int
	var err error
	f, err = newFramework()
	if err != nil {
		log.Printf("unexpected error: %v\n", err)
		os.Exit(1)
	}

	c = newAWSClient()

	if err := f.SetUp(); err != nil {
		log.Printf("unexpected error: %v\n", err)
		v = 1
	}

	if err := operatorSetup(); err != nil {
		log.Printf("unexpected error: %v\n", err)
		v = 1
	}

	if v == 0 {
		v = m.Run()
	}

	f.DeleteGuestCluster()
	operatorTearDown()
	f.TearDown()

	os.Exit(v)
}

func TestGuestReadyAfterMasterReboot(t *testing.T) {
	log.Println("getting master ID")
	describeInput := &ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name: aws.String("tag:Name"),
				Values: []*string{
					aws.String(fmt.Sprintf("%s-master", os.Getenv("CLUSTER_NAME"))),
				},
			},
			&ec2.Filter{
				Name: aws.String("instance-state-name"),
				Values: []*string{
					aws.String("running"),
				},
			},
		},
	}
	res, err := c.EC2.DescribeInstances(describeInput)
	if err != nil {
		t.Errorf("unexpected error getting master id %v", err)
	}
	if len(res.Reservations) != 1 {
		t.Errorf("unexpected number of reservations %d", len(res.Reservations))
	}
	if len(res.Reservations[0].Instances) != 1 {
		t.Errorf("unexpected number of instances %d", len(res.Reservations[0].Instances))
	}
	masterID := res.Reservations[0].Instances[0].InstanceId

	log.Println("rebooting master")
	rebootInput := &ec2.RebootInstancesInput{
		InstanceIds: []*string{
			masterID,
		},
	}
	_, err = c.EC2.RebootInstances(rebootInput)
	if err != nil {
		t.Errorf("unexpected error rebooting  master %v", err)
	}

	if err := f.WaitForAPIDown(); err != nil {
		t.Errorf("unexpected error waiting for master shutting down %v", err)
	}

	if err := f.WaitForGuestReady(); err != nil {
		t.Errorf("unexpected error waiting for guest cluster ready, %v", err)
	}
}

func TestWorkersScaling(t *testing.T) {
	currentWorkers, err := numberOfWorkers(os.Getenv("CLUSTER_NAME"))
	if err != nil {
		t.Fatalf("unexpected error getting number of workers %v", err)
	}
	currentMasters, err := numberOfMasters(os.Getenv("CLUSTER_NAME"))
	if err != nil {
		t.Fatalf("unexpected error getting number of masters %v", err)
	}

	// increase number of workers
	expectedWorkers := currentWorkers + 1
	log.Printf("Increasing the number of workers to %d", expectedWorkers)
	err = addWorker(os.Getenv("CLUSTER_NAME"))
	if err != nil {
		t.Fatalf("unexpected error setting number of workers to %d, %v", expectedWorkers, err)
	}

	if err := f.WaitForNodesUp(currentMasters + expectedWorkers); err != nil {
		t.Fatalf("unexpected error waiting for %d nodes up, %v", expectedWorkers, err)
	}
	log.Printf("%d worker nodes ready", expectedWorkers)

	// decrease number of workers
	expectedWorkers--
	log.Printf("Decreasing the number of workers to %d", expectedWorkers)
	err = removeWorker(os.Getenv("CLUSTER_NAME"))
	if err != nil {
		t.Fatalf("unexpected error setting number of workers to %d, %v", expectedWorkers, err)
	}

	if err := f.WaitForNodesUp(currentMasters + expectedWorkers); err != nil {
		t.Fatalf("unexpected error waiting for %d nodes up, %v", expectedWorkers, err)
	}
	log.Printf("%d worker nodes ready", expectedWorkers)
}

func operatorSetup() error {
	if err := f.InstallCertOperator(); err != nil {
		return microerror.Mask(err)
	}

	if err := f.InstallCertResource(); err != nil {
		return microerror.Mask(err)
	}

	if err := f.InstallAwsOperator(); err != nil {
		return microerror.Mask(err)
	}

	if err := writeAWSResourceValues(); err != nil {
		return microerror.Maskf(err, "unexpected error writing aws-resource-lab values file")
	}

	if err := runCmd("helm registry install quay.io/giantswarm/aws-resource-lab-chart:stable -- -n aws-resource-lab --values " + awsOperatorValuesFile); err != nil {
		return microerror.Maskf(err, "unexpected error installing aws-resource-lab chart: %v")
	}

	logEntry := "cluster '${CLUSTER_NAME}' processed"
	if os.Getenv("VERSION_BUNDLE_VERSION") == "0.2.0" {
		logEntry = "creating AWS cloudformation stack: created"
	}

	operatorPodName, err := f.PodName("giantswarm", "app=aws-operator")
	if err != nil {
		return microerror.Maskf(err, "unexpected error getting operator pod name: %v")
	}

	if err := f.WaitForPodLog("giantswarm", logEntry, operatorPodName); err != nil {
		return microerror.Maskf(err, "unexpected error waiting for guest cluster installed: %v")
	}

	if err := f.WaitForGuestReady(); err != nil {
		return microerror.Maskf(err, "unexpected error waiting for guest cluster ready")
	}

	return nil
}

func operatorTearDown() {
	runCmd("helm delete cert-resource-lab --purge")
	runCmd("helm delete cert-operator --purge")
	runCmd("helm delete aws-resource-lab --purge")
	runCmd("helm delete aws-operator --purge")
}

func writeAWSResourceValues() error {
	awsResourceChartValuesEnv := os.ExpandEnv(awsResourceChartValues)

	tmpl, err := template.New("awsResource").Parse(awsResourceChartValuesEnv)
	if err != nil {
		return microerror.Mask(err)
	}

	f, err := os.Create(awsResourceValuesFile)
	if err != nil {
		return microerror.Mask(err)
	}
	defer f.Close()

	vpc, err := newAWSVPCBlock(c)
	if err != nil {
		return microerror.Mask(err)
	}

	err = tmpl.Execute(f, vpc)
	if err != nil {
		return microerror.Mask(err)
	}

	return nil
}

func numberOfWorkers(clusterName string) (int, error) {
	cluster, err := f.AWSCluster(clusterName)
	if err != nil {
		return 0, microerror.Mask(err)
	}

	return keyv2.WorkerCount(*cluster), nil
}

func numberOfMasters(clusterName string) (int, error) {
	cluster, err := f.AWSCluster(clusterName)
	if err != nil {
		return 0, microerror.Mask(err)
	}

	return keyv2.MasterCount(*cluster), nil
}

func addWorker(clusterName string) error {
	cluster, err := f.AWSCluster(clusterName)
	if err != nil {
		return microerror.Mask(err)
	}

	newWorker := cluster.Spec.AWS.Workers[0]

	patch := make([]PatchSpec, 1)
	patch[0].Op = "add"
	patch[0].Path = "/spec/aws/workers/-"
	patch[0].Value = newWorker

	return f.ApplyAWSConfigPatch(patch, clusterName)
}

func removeWorker(clusterName string) error {
	patch := make([]PatchSpec, 1)
	patch[0].Op = "remove"
	patch[0].Path = "/spec/aws/workers/1"

	return f.ApplyAWSConfigPatch(patch, clusterName)
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/elazarl/goproxy"

	"flag"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"
	"time"
)

// Digest auth. operation type
const (
	validateUser int = iota
	getNonce     int = iota
	maintPing    int = iota
)

// Digest auth. resp status
const (
	authOk     int = iota
	authFailed int = iota
	nonceOk    int = iota
	maintOk    int = iota
)

const proxyForwardedForHeader = "X-Forwarded-For"
const tcpKeepAliveInterval = 1 * time.Minute

type basicAuthRequest struct {
	data        *basicAuthData
	respChannel chan *basicAuthResponse
}

type basicAuthResponse struct {
	status bool
}

type digestAuthRequest struct {
	data        *digestAuthData
	op          int
	respChannel chan *digestAuthResponse
}

type digestAuthResponse struct {
	data   string
	status int
}

func makeBasicAuthValidator(auth *basicAuth) basicAuthFunc {
	channel := make(chan *basicAuthRequest)
	validator := func() {
		for e := range channel {
			status := auth.validate(e.data)
			e.respChannel <- &basicAuthResponse{status}
		}
	}

	go validator()

	return func(authData *basicAuthData) *basicAuthResponse {
		request := &basicAuthRequest{authData, make(chan *basicAuthResponse)}
		channel <- request
		return <-request.respChannel
	}
}

func makeDigestAuthValidator(auth *digestAuth) digestAuthFunc {
	channel := make(chan *digestAuthRequest)

	processor := func() {
		for e := range channel {
			var response *digestAuthResponse
			switch e.op {
			case validateUser:
				status := auth.validate(e.data)
				if status {
					response = &digestAuthResponse{status: authOk}
				} else {
					response = &digestAuthResponse{status: authFailed}
				}
			case getNonce:
				nonce := auth.newNonce()
				response = &digestAuthResponse{status: nonceOk, data: nonce}
			case maintPing:
				auth.expireNonces()
				response = &digestAuthResponse{status: maintOk}
			default:
				panic("unexpected operation type")
			}
			e.respChannel <- response
		}
	}

	maintPinger := func() {
		for {
			request := &digestAuthRequest{op: maintPing, respChannel: make(chan *digestAuthResponse)}
			channel <- request
			response := <-request.respChannel
			if response.status != maintOk {
				log.Fatal("unexpected status")
			}
			time.Sleep(30 * time.Minute)
		}
	}

	go processor()
	go maintPinger()

	authFunc := func(authData *digestAuthData, op int) *digestAuthResponse {
		request := &digestAuthRequest{data: authData, op: op, respChannel: make(chan *digestAuthResponse)}
		channel <- request
		return <-request.respChannel
	}

	return authFunc
}

func setAllowedNetworksHandler(conf *configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.AllowedNetworks != nil && len(conf.AllowedNetworks) > 0 {
		proxy.OnRequest(goproxy.Not(sourceIPMatches(conf.AllowedNetworks))).HandleConnect(goproxy.AlwaysReject)
		proxy.OnRequest(goproxy.Not(sourceIPMatches(conf.AllowedNetworks))).DoFunc(
			func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
				return req, goproxy.NewResponse(req, goproxy.ContentTypeHtml, http.StatusForbidden, "Access denied")
			})
	}

	if conf.DisallowedNetworks != nil && len(conf.DisallowedNetworks) > 0 {
		proxy.OnRequest(sourceIPMatches(conf.DisallowedNetworks)).HandleConnect(goproxy.AlwaysReject)
		proxy.OnRequest(sourceIPMatches(conf.DisallowedNetworks)).DoFunc(
			func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
				return req, goproxy.NewResponse(req, goproxy.ContentTypeHtml, http.StatusForbidden, "Access denied")
			})
	}
}

func sourceIPMatches(networks []string) goproxy.ReqConditionFunc {
	cidrs := make([](*net.IPNet), len(networks))
	for idx, network := range networks {
		_, cidrnet, _ := net.ParseCIDR(network)
		cidrs[idx] = cidrnet
	}

	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		ip, _, err := net.SplitHostPort(req.RemoteAddr)
		if err != nil {
			ctx.Warnf("coudn't parse remote address %v: %v", req.RemoteAddr, err)
			return false
		}
		addr := net.ParseIP(ip)
		for _, network := range cidrs {
			if network.Contains(addr) {
				return true
			}
		}
		return false
	}
}

func setAllowedConnectPortsHandler(conf *configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.AllowedConnectPorts != nil && len(conf.AllowedConnectPorts) > 0 {
		ports := make([]string, len(conf.AllowedConnectPorts))
		for i, v := range conf.AllowedConnectPorts {
			ports[i] = ":" + v
		}
		rx := "(" + strings.Join(ports, "|") + ")$"
		proxy.OnRequest(goproxy.Not(goproxy.ReqHostMatches(regexp.MustCompile(rx)))).HandleConnect(goproxy.AlwaysReject)
	}
}

func setForwardedForHeaderHandler(conf *configuration, proxy *goproxy.ProxyHttpServer) {
	f := func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		ip, _, err := net.SplitHostPort(req.RemoteAddr)
		if err != nil {
			ctx.Warnf("coudn't parse remote address %v: %v", req.RemoteAddr, err)
			return req, nil
		}

		switch conf.ForwardedFor {
		case "on":
			header := req.Header.Get(proxyForwardedForHeader)
			if header == "" {
				req.Header.Add(proxyForwardedForHeader, ip)
			} else {
				header = header + ", " + ip
				req.Header.Del(proxyForwardedForHeader)
				req.Header.Add(proxyForwardedForHeader, header)
			}
		case "delete":
			req.Header.Del(proxyForwardedForHeader)
		case "truncate":
			req.Header.Del(proxyForwardedForHeader)
			req.Header.Add(proxyForwardedForHeader, ip)
		}

		return req, nil
	}

	proxy.OnRequest().DoFunc(f)
}

func makeCustomDial(localAddr *net.TCPAddr) func(string, string) (net.Conn, error) {
	return func(network, addr string) (net.Conn, error) {
		remoteAddr, err := net.ResolveTCPAddr(network, addr)
		if err != nil {
			return nil, err
		}

		conn, err := net.DialTCP(network, localAddr, remoteAddr)
		if err == nil {
			conn.SetKeepAlive(true)
			conn.SetKeepAlivePeriod(tcpKeepAliveInterval)
		}

		c := timedConn{
			Conn:         conn,
			readTimeout:  defaultReadTimeout,
			writeTimeout: defaultWriteTimeout,
		}

		return c, err
	}
}

func createProxy(conf *configuration) *goproxy.ProxyHttpServer {
	proxy := goproxy.NewProxyHttpServer()
	setActivityLog(conf, proxy)

	var laddr string

	addressOk := true

	if conf.BindIP != "" {
		if ip := net.ParseIP(conf.BindIP); ip != nil {
			if ip4 := ip.To4(); ip4 != nil {
				laddr = conf.BindIP + ":0"
			} else if ip16 := ip.To16(); ip16 != nil {
				laddr = "[" + conf.BindIP + "]:0"
			} else {
				proxy.Logger.Printf("WARN: couldn't use \"%v\" as outgoing request address.\n", conf.BindIP)
				addressOk = false
			}
		}
	}

	if addressOk {
		if laddr != "" {
			if addr, err := net.ResolveTCPAddr("tcp", laddr); err == nil {
				proxy.Tr.Dial = makeCustomDial(addr)
			} else {
				proxy.Logger.Printf("WARN: couldn't use \"%v\" as outgoing request address. %v\n", conf.BindIP, err)
			}
		} else {
			proxy.Tr.Dial = makeCustomDial(nil)
		}
	}

	return proxy
}

func setActivityLog(conf *configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.ActivityLog != "" {
		fh, err := os.OpenFile(conf.ActivityLog, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			log.Fatalf("couldn't open activity log file %v: %v", conf.ActivityLog, err)
		}
		proxy.Logger = log.New(fh, "", log.LstdFlags)
	}
}

func setSignalHandler(conf *configuration, proxy *goproxy.ProxyHttpServer, logger *proxyLogger) {
	signalChannel := make(chan os.Signal)
	signal.Notify(signalChannel, os.Interrupt, syscall.SIGUSR1)

	go func() {
		for {
			sig := <-signalChannel
			switch sig {
			case os.Interrupt:
				proxy.Logger.Println("got interrupt signal, exiting")
				logger.close()
				os.Exit(0)
			case syscall.SIGUSR1:
				proxy.Logger.Println("got USR1 signal, reopening logs")
				// reopen access log
				logger.reopen()
				// reopen activity log
				setActivityLog(conf, proxy)
			}
		}
	}()
}

func setAuthenticationHandler(conf *configuration, proxy *goproxy.ProxyHttpServer, logger *proxyLogger) {
	if conf.AuthFile != "" {
		if conf.AuthType == "basic" {
			auth, err := newBasicAuthFromFile(conf.AuthFile)
			if err != nil {
				proxy.Logger.Fatalf("couldn't create basic auth structure: %v\n", err)
			}
			setProxyBasicAuth(proxy, conf.AuthRealm, makeBasicAuthValidator(auth), logger)
		} else {
			auth, err := newDigestAuthFromFile(conf.AuthFile)
			if err != nil {
				proxy.Logger.Fatalf("couldn't create digest auth structure: %v\n", err)
			}
			setProxyDigestAuth(proxy, conf.AuthRealm, makeDigestAuthValidator(auth), logger)
		}
	} else {
		// If there is neither Digest no Basic authentication we still need to setup
		// handler to log HTTPS CONNECT requests
		setHTTPSLoggingHandler(proxy, logger)
	}
}

func setHTTPSLoggingHandler(proxy *goproxy.ProxyHttpServer, logger *proxyLogger) {
	proxy.OnRequest().HandleConnectFunc(
		func(host string, ctx *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) {
			if ctx.Req == nil {
				ctx.Req = emptyReq
			}

			if logger != nil {
				logger.log(ctx)
			}

			return goproxy.OkConnect, host
		})
}

func setHTTPLoggingHandler(proxy *goproxy.ProxyHttpServer, logger *proxyLogger) {
	proxy.OnResponse().DoFunc(
		func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
			logger.logResponse(resp, ctx)
			return resp
		})
}

func main() {
	config := flag.String("config", "microproxy.json", "proxy configuration file")
	verbose := flag.Bool("v", false, "enable verbose debug mode")

	flag.Parse()

	conf := newConfigurationFromFile(*config)

	proxy := createProxy(conf)
	proxy.Verbose = *verbose

	logger := newProxyLogger(conf)

	setHTTPLoggingHandler(proxy, logger)
	setAllowedConnectPortsHandler(conf, proxy)
	setAllowedNetworksHandler(conf, proxy)
	setForwardedForHeaderHandler(conf, proxy)
	setSignalHandler(conf, proxy, logger)

	// To be called first while processing handlers' stack,
	// has to be placed last in the source code.
	setAuthenticationHandler(conf, proxy, logger)

	proxy.Logger.Println("starting proxy")

	log.Fatal(http.ListenAndServe(conf.Listen, proxy))
}
<commit_msg>catch TERM signal, log that event and exit<commit_after>package main

import (
	"github.com/elazarl/goproxy"

	"flag"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"
	"time"
)

// Digest auth. operation type
const (
	validateUser int = iota
	getNonce     int = iota
	maintPing    int = iota
)

// Digest auth. resp status
const (
	authOk     int = iota
	authFailed int = iota
	nonceOk    int = iota
	maintOk    int = iota
)

const proxyForwardedForHeader = "X-Forwarded-For"
const tcpKeepAliveInterval = 1 * time.Minute

type basicAuthRequest struct {
	data        *basicAuthData
	respChannel chan *basicAuthResponse
}

type basicAuthResponse struct {
	status bool
}

type digestAuthRequest struct {
	data        *digestAuthData
	op          int
	respChannel chan *digestAuthResponse
}

type digestAuthResponse struct {
	data   string
	status int
}

func makeBasicAuthValidator(auth *basicAuth) basicAuthFunc {
	channel := make(chan *basicAuthRequest)
	validator := func() {
		for e := range channel {
			status := auth.validate(e.data)
			e.respChannel <- &basicAuthResponse{status}
		}
	}

	go validator()

	return func(authData *basicAuthData) *basicAuthResponse {
		request := &basicAuthRequest{authData, make(chan *basicAuthResponse)}
		channel <- request
		return <-request.respChannel
	}
}

func makeDigestAuthValidator(auth *digestAuth) digestAuthFunc {
	channel := make(chan *digestAuthRequest)

	processor := func() {
		for e := range channel {
			var response *digestAuthResponse
			switch e.op {
			case validateUser:
				status := auth.validate(e.data)
				if status {
					response = &digestAuthResponse{status: authOk}
				} else {
					response = &digestAuthResponse{status: authFailed}
				}
			case getNonce:
				nonce := auth.newNonce()
				response = &digestAuthResponse{status: nonceOk, data: nonce}
			case maintPing:
				auth.expireNonces()
				response = &digestAuthResponse{status: maintOk}
			default:
				panic("unexpected operation type")
			}
			e.respChannel <- response
		}
	}

	maintPinger := func() {
		for {
			request := &digestAuthRequest{op: maintPing, respChannel: make(chan *digestAuthResponse)}
			channel <- request
			response := <-request.respChannel
			if response.status != maintOk {
				log.Fatal("unexpected status")
			}
			time.Sleep(30 * time.Minute)
		}
	}

	go processor()
	go maintPinger()

	authFunc := func(authData *digestAuthData, op int) *digestAuthResponse {
		request := &digestAuthRequest{data: authData, op: op, respChannel: make(chan *digestAuthResponse)}
		channel <- request
		return <-request.respChannel
	}

	return authFunc
}

func setAllowedNetworksHandler(conf *configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.AllowedNetworks != nil && len(conf.AllowedNetworks) > 0 {
		proxy.OnRequest(goproxy.Not(sourceIPMatches(conf.AllowedNetworks))).HandleConnect(goproxy.AlwaysReject)
		proxy.OnRequest(goproxy.Not(sourceIPMatches(conf.AllowedNetworks))).DoFunc(
			func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
				return req, goproxy.NewResponse(req, goproxy.ContentTypeHtml, http.StatusForbidden, "Access denied")
			})
	}

	if conf.DisallowedNetworks != nil && len(conf.DisallowedNetworks) > 0 {
		proxy.OnRequest(sourceIPMatches(conf.DisallowedNetworks)).HandleConnect(goproxy.AlwaysReject)
		proxy.OnRequest(sourceIPMatches(conf.DisallowedNetworks)).DoFunc(
			func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
				return req, goproxy.NewResponse(req, goproxy.ContentTypeHtml, http.StatusForbidden, "Access denied")
			})
	}
}

func sourceIPMatches(networks []string) goproxy.ReqConditionFunc {
	cidrs := make([](*net.IPNet), len(networks))
	for idx, network := range networks {
		_, cidrnet, _ := net.ParseCIDR(network)
		cidrs[idx] = cidrnet
	}

	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		ip, _, err := net.SplitHostPort(req.RemoteAddr)
		if err != nil {
			ctx.Warnf("coudn't parse remote address %v: %v", req.RemoteAddr, err)
			return false
		}
		addr := net.ParseIP(ip)
		for _, network := range cidrs {
			if network.Contains(addr) {
				return true
			}
		}
		return false
	}
}

func setAllowedConnectPortsHandler(conf *configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.AllowedConnectPorts != nil && len(conf.AllowedConnectPorts) > 0 {
		ports := make([]string, len(conf.AllowedConnectPorts))
		for i, v := range conf.AllowedConnectPorts {
			ports[i] = ":" + v
		}
		rx := "(" + strings.Join(ports, "|") + ")$"
		proxy.OnRequest(goproxy.Not(goproxy.ReqHostMatches(regexp.MustCompile(rx)))).HandleConnect(goproxy.AlwaysReject)
	}
}

func setForwardedForHeaderHandler(conf *configuration, proxy *goproxy.ProxyHttpServer) {
	f := func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		ip, _, err := net.SplitHostPort(req.RemoteAddr)
		if err != nil {
			ctx.Warnf("coudn't parse remote address %v: %v", req.RemoteAddr, err)
			return req, nil
		}

		switch conf.ForwardedFor {
		case "on":
			header := req.Header.Get(proxyForwardedForHeader)
			if header == "" {
				req.Header.Add(proxyForwardedForHeader, ip)
			} else {
				header = header + ", " + ip
				req.Header.Del(proxyForwardedForHeader)
				req.Header.Add(proxyForwardedForHeader, header)
			}
		case "delete":
			req.Header.Del(proxyForwardedForHeader)
		case "truncate":
			req.Header.Del(proxyForwardedForHeader)
			req.Header.Add(proxyForwardedForHeader, ip)
		}

		return req, nil
	}

	proxy.OnRequest().DoFunc(f)
}

func makeCustomDial(localAddr *net.TCPAddr) func(string, string) (net.Conn, error) {
	return func(network, addr string) (net.Conn, error) {
		remoteAddr, err := net.ResolveTCPAddr(network, addr)
		if err != nil {
			return nil, err
		}

		conn, err := net.DialTCP(network, localAddr, remoteAddr)
		if err == nil {
			conn.SetKeepAlive(true)
			conn.SetKeepAlivePeriod(tcpKeepAliveInterval)
		}

		c := timedConn{
			Conn:         conn,
			readTimeout:  defaultReadTimeout,
			writeTimeout: defaultWriteTimeout,
		}

		return c, err
	}
}

func createProxy(conf *configuration) *goproxy.ProxyHttpServer {
	proxy := goproxy.NewProxyHttpServer()
	setActivityLog(conf, proxy)

	var laddr string

	addressOk := true

	if conf.BindIP != "" {
		if ip := net.ParseIP(conf.BindIP); ip != nil {
			if ip4 := ip.To4(); ip4 != nil {
				laddr = conf.BindIP + ":0"
			} else if ip16 := ip.To16(); ip16 != nil {
				laddr = "[" + conf.BindIP + "]:0"
			} else {
				proxy.Logger.Printf("WARN: couldn't use \"%v\" as outgoing request address.\n", conf.BindIP)
				addressOk = false
			}
		}
	}

	if addressOk {
		if laddr != "" {
			if addr, err := net.ResolveTCPAddr("tcp", laddr); err == nil {
				proxy.Tr.Dial = makeCustomDial(addr)
			} else {
				proxy.Logger.Printf("WARN: couldn't use \"%v\" as outgoing request address. %v\n", conf.BindIP, err)
			}
		} else {
			proxy.Tr.Dial = makeCustomDial(nil)
		}
	}

	return proxy
}

func setActivityLog(conf *configuration, proxy *goproxy.ProxyHttpServer) {
	if conf.ActivityLog != "" {
		fh, err := os.OpenFile(conf.ActivityLog, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			log.Fatalf("couldn't open activity log file %v: %v", conf.ActivityLog, err)
		}
		proxy.Logger = log.New(fh, "", log.LstdFlags)
	}
}

func setSignalHandler(conf *configuration, proxy *goproxy.ProxyHttpServer, logger *proxyLogger) {
	signalChannel := make(chan os.Signal)
	signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM, syscall.SIGUSR1)

	signalHandler := func() {
		for sig := range signalChannel {
			switch sig {
			case os.Interrupt, syscall.SIGTERM:
				proxy.Logger.Println("got interrupt signal, exiting")
				logger.close()
				os.Exit(0)
			case syscall.SIGUSR1:
				proxy.Logger.Println("got USR1 signal, reopening logs")
				// reopen access log
				logger.reopen()
				// reopen activity log
				setActivityLog(conf, proxy)
			}
		}
	}

	go signalHandler()
}

func setAuthenticationHandler(conf *configuration, proxy *goproxy.ProxyHttpServer, logger *proxyLogger) {
	if conf.AuthFile != "" {
		if conf.AuthType == "basic" {
			auth, err := newBasicAuthFromFile(conf.AuthFile)
			if err != nil {
				proxy.Logger.Fatalf("couldn't create basic auth structure: %v\n", err)
			}
			setProxyBasicAuth(proxy, conf.AuthRealm, makeBasicAuthValidator(auth), logger)
		} else {
			auth, err := newDigestAuthFromFile(conf.AuthFile)
			if err != nil {
				proxy.Logger.Fatalf("couldn't create digest auth structure: %v\n", err)
			}
			setProxyDigestAuth(proxy, conf.AuthRealm, makeDigestAuthValidator(auth), logger)
		}
	} else {
		// If there is neither Digest no Basic authentication we still need to setup
		// handler to log HTTPS CONNECT requests
		setHTTPSLoggingHandler(proxy, logger)
	}
}

func setHTTPSLoggingHandler(proxy *goproxy.ProxyHttpServer, logger *proxyLogger) {
	proxy.OnRequest().HandleConnectFunc(
		func(host string, ctx *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) {
			if ctx.Req == nil {
				ctx.Req = emptyReq
			}

			if logger != nil {
				logger.log(ctx)
			}

			return goproxy.OkConnect, host
		})
}

func setHTTPLoggingHandler(proxy *goproxy.ProxyHttpServer, logger *proxyLogger) {
	proxy.OnResponse().DoFunc(
		func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
			logger.logResponse(resp, ctx)
			return resp
		})
}

func main() {
	config := flag.String("config", "microproxy.json", "proxy configuration file")
	verbose := flag.Bool("v", false, "enable verbose debug mode")

	flag.Parse()

	conf := newConfigurationFromFile(*config)

	proxy := createProxy(conf)
	proxy.Verbose = *verbose

	logger := newProxyLogger(conf)

	setHTTPLoggingHandler(proxy, logger)
	setAllowedConnectPortsHandler(conf, proxy)
	setAllowedNetworksHandler(conf, proxy)
	setForwardedForHeaderHandler(conf, proxy)
	setSignalHandler(conf, proxy, logger)

	// To be called first while processing handlers' stack,
	// has to be placed last in the source code.
	setAuthenticationHandler(conf, proxy, logger)

	proxy.Logger.Println("starting proxy")

	log.Fatal(http.ListenAndServe(conf.Listen, proxy))
}
<|endoftext|>


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

package fetch

import (
	"archive/zip"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/source"
)

// Version and commit time are pre specified when fetching a local module, as these
// fields are normally obtained from a proxy.
var (
	LocalVersion    = "v0.0.0"
	LocalCommitTime = time.Time{}
)

// A directoryModuleGetter is a ModuleGetter whose source is a directory in the file system that contains
// a module's files.
type directoryModuleGetter struct {
	dir string // the directory containing the module's files
}

// NewDirectoryModuleGetter returns a ModuleGetter for reading a module from a directory.
func NewDirectoryModuleGetter(dir string) ModuleGetter {
	return &directoryModuleGetter{dir: dir}
}

// Info returns basic information about the module.
func (g *directoryModuleGetter) Info(ctx context.Context, path, version string) (*proxy.VersionInfo, error) {
	return &proxy.VersionInfo{
		Version: LocalVersion,
		Time:    LocalCommitTime,
	}, nil
}

// Mod returns the contents of the module's go.mod file.
// If the file does not exist, it returns a synthesized one.
func (g *directoryModuleGetter) Mod(ctx context.Context, path, version string) ([]byte, error) {
	data, err := ioutil.ReadFile(filepath.Join(g.dir, "go.mod"))
	if errors.Is(err, os.ErrNotExist) {
		if path == "" {
			return nil, fmt.Errorf("no module path: %w", derrors.BadModule)
		}
		return []byte(fmt.Sprintf("module %s\n", path)), nil
	}
	return data, err
}

// Zip returns a reader for the module's zip file.
func (g *directoryModuleGetter) Zip(ctx context.Context, path, version string) (*zip.Reader, error) {
	return createZipReader(g.dir, path, LocalVersion)
}

// ZipSize returns the approximate size of the zip file in bytes.
func (g *directoryModuleGetter) ZipSize(ctx context.Context, path, version string) (int64, error) {
	return 0, errors.New("directoryModuleGetter.ZipSize unimplemented")
}

// FetchLocalModule fetches a module from a local directory and process its contents
// to return an internal.Module and other related information. modulePath is not necessary
// if the module has a go.mod file, but if both exist, then they must match.
// FetchResult.Error should be checked to verify that the fetch succeeded. Even if the
// error is non-nil the result may contain useful data.
func FetchLocalModule(ctx context.Context, modulePath, localPath string, sourceClient *source.Client) *FetchResult {
	g := NewDirectoryModuleGetter(localPath)
	fr := FetchModule(ctx, modulePath, LocalVersion, g, sourceClient)
	if fr.Error != nil {
		fr.Error = fmt.Errorf("FetchLocalModule(%q, %q): %w", modulePath, localPath, fr.Error)
	}
	return fr
}

// createZipReader creates a zip file from a directory given a local path and
// returns a zip.Reader to be passed to processZipFile. The purpose of the
// function is to transform a local go module into a zip file to be processed by
// existing functions.
func createZipReader(localPath, modulePath, version string) (*zip.Reader, error) {
	buf := new(bytes.Buffer)
	w := zip.NewWriter(buf)
	err := filepath.Walk(localPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}

		readFrom, err := os.Open(path)
		if err != nil {
			return err
		}
		defer readFrom.Close()

		writeTo, err := w.Create(filepath.Join(moduleVersionDir(modulePath, version), strings.TrimPrefix(path, localPath)))
		if err != nil {
			return err
		}

		_, err = io.Copy(writeTo, readFrom)
		return err
	})
	if err != nil {
		return nil, err
	}
	if err := w.Close(); err != nil {
		return nil, err
	}

	reader := bytes.NewReader(buf.Bytes())
	return zip.NewReader(reader, reader.Size())
}
<commit_msg>internal/fetch: a directoryModuleGetter knows its module path<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package fetch

import (
	"archive/zip"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/mod/modfile"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/source"
)

// Version and commit time are pre specified when fetching a local module, as these
// fields are normally obtained from a proxy.
var (
	LocalVersion    = "v0.0.0"
	LocalCommitTime = time.Time{}
)

// A directoryModuleGetter is a ModuleGetter whose source is a directory in the file system that contains
// a module's files.
type directoryModuleGetter struct {
	modulePath string
	dir        string
}

// NewDirectoryModuleGetter returns a ModuleGetter for reading a module from a directory.
func NewDirectoryModuleGetter(modulePath, dir string) (*directoryModuleGetter, error) {
	if modulePath == "" {
		goModBytes, err := ioutil.ReadFile(filepath.Join(dir, "go.mod"))
		if err != nil {
			return nil, fmt.Errorf("cannot obtain module path for %q (%v): %w", dir, err, derrors.BadModule)
		}
		modulePath = modfile.ModulePath(goModBytes)
		if modulePath == "" {
			return nil, fmt.Errorf("go.mod in %q has no module path: %w", dir, derrors.BadModule)
		}
	}
	return &directoryModuleGetter{
		dir:        dir,
		modulePath: modulePath,
	}, nil
}

func (g *directoryModuleGetter) checkPath(path string) error {
	if path != g.modulePath {
		return fmt.Errorf("given module path %q does not match %q for directory %q: %w",
			path, g.modulePath, g.dir, derrors.NotFound)
	}
	return nil
}

// Info returns basic information about the module.
func (g *directoryModuleGetter) Info(ctx context.Context, path, version string) (*proxy.VersionInfo, error) {
	if err := g.checkPath(path); err != nil {
		return nil, err
	}
	return &proxy.VersionInfo{
		Version: LocalVersion,
		Time:    LocalCommitTime,
	}, nil
}

// Mod returns the contents of the module's go.mod file.
// If the file does not exist, it returns a synthesized one.
func (g *directoryModuleGetter) Mod(ctx context.Context, path, version string) ([]byte, error) {
	if err := g.checkPath(path); err != nil {
		return nil, err
	}
	data, err := ioutil.ReadFile(filepath.Join(g.dir, "go.mod"))
	if errors.Is(err, os.ErrNotExist) {
		return []byte(fmt.Sprintf("module %s\n", g.modulePath)), nil
	}
	return data, err
}

// Zip returns a reader for the module's zip file.
func (g *directoryModuleGetter) Zip(ctx context.Context, path, version string) (*zip.Reader, error) {
	if err := g.checkPath(path); err != nil {
		return nil, err
	}
	return createZipReader(g.dir, path, LocalVersion)
}

// ZipSize returns the approximate size of the zip file in bytes.
func (g *directoryModuleGetter) ZipSize(ctx context.Context, path, version string) (int64, error) {
	return 0, errors.New("directoryModuleGetter.ZipSize unimplemented")
}

// FetchLocalModule fetches a module from a local directory and process its contents
// to return an internal.Module and other related information. modulePath is not necessary
// if the module has a go.mod file, but if both exist, then they must match.
// FetchResult.Error should be checked to verify that the fetch succeeded. Even if the
// error is non-nil the result may contain useful data.
func FetchLocalModule(ctx context.Context, modulePath, localPath string, sourceClient *source.Client) *FetchResult {
	g, err := NewDirectoryModuleGetter(modulePath, localPath)
	if err != nil {
		return &FetchResult{
			ModulePath: modulePath,
			Error:      err,
		}
	}
	if modulePath == "" {
		modulePath = g.modulePath
	}
	fr := FetchModule(ctx, modulePath, LocalVersion, g, sourceClient)
	if fr.Error != nil {
		fr.Error = fmt.Errorf("FetchLocalModule(%q, %q): %w", modulePath, localPath, fr.Error)
	}
	return fr
}

// createZipReader creates a zip file from a directory given a local path and
// returns a zip.Reader to be passed to processZipFile. The purpose of the
// function is to transform a local go module into a zip file to be processed by
// existing functions.
func createZipReader(localPath, modulePath, version string) (*zip.Reader, error) {
	buf := new(bytes.Buffer)
	w := zip.NewWriter(buf)
	err := filepath.Walk(localPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}

		readFrom, err := os.Open(path)
		if err != nil {
			return err
		}
		defer readFrom.Close()

		writeTo, err := w.Create(filepath.Join(moduleVersionDir(modulePath, version), strings.TrimPrefix(path, localPath)))
		if err != nil {
			return err
		}

		_, err = io.Copy(writeTo, readFrom)
		return err
	})
	if err != nil {
		return nil, err
	}
	if err := w.Close(); err != nil {
		return nil, err
	}

	reader := bytes.NewReader(buf.Bytes())
	return zip.NewReader(reader, reader.Size())
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/csv"
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"
)

// TODO: Write tests!
func main() {
	// Test 02
	//textColumn := []string{"One", "Two", "Three"}
	//intColumn := []int{1, 2, 3}
	//c1 := Column{}
	//c1.fillColumn(textColumn)
	//c2 := Column{}
	//c2.fillColumn(intColumn)
	//df := DataFrame{
	//columns:  []Column{c1, c2},
	//nCols:    2,
	//nRows:    3,
	//colnames: []string{"Text", "Ints"},
	//}
	//fmt.Println(df)

	// Test 01
	in := `A,B,C,D
1,2,3,4
5,6,7,8`
	df := DataFrame{}
	r := csv.NewReader(strings.NewReader(in))
	records, err := r.ReadAll()
	if err != nil {
		panic(err)
	}
	err = df.loadAndParse(records, []string{"string", "int", "string", "int"})
	if err != nil {
		panic(err)
	}

	for _, v := range df.columns {
		fmt.Println(v)
		fmt.Println(v.colType)
	}
	fmt.Println(df)
}

// DataFrame Definition
// ====================
type DataFrame struct {
	columns  []Column
	colnames []string
	nCols    int
	nRows    int
}

// DataFrame Methods
// =================
func (df *DataFrame) loadAndParse(records [][]string, types []string) error {
	err := df.loadData(records)
	if err != nil {
		return err
	}
	if df.nCols != len(types) {
		return errors.New("Number of columns different from number of types")
	}
	for k, v := range df.columns {
		v.parseType(types[k])
		df.columns[k].colType = types[k]
	}
	return nil
}

func (df *DataFrame) loadData(records [][]string) error {
	// TODO: Check if empty records
	// TODO: More error checking

	// Get DataFrame dimensions
	nRows := len(records) - 1
	if nRows == 0 {
		return errors.New("Empty dataframe")
	}
	nCols := len(records[0])

	// Generate a virtual df to store the temporary values
	newDf := DataFrame{
		columns:  []Column{},
		colnames: records[0],
		nRows:    nRows,
		nCols:    nCols,
	}

	for j := 0; j < nCols; j++ {
		col := []string{}
		for i := 1; i < nRows+1; i++ {
			col = append(col, records[i][j])
		}
		column := Column{}
		column.fillColumn(col)
		newDf.columns = append(newDf.columns, column)
	}
	*df = newDf
	return nil
}

func (df DataFrame) String() string {
	str := ""
	if len(df.colnames) != 0 {
		str += "\t"
		for _, v := range df.colnames {
			str += v
			str += "\t"
		}
		str += "\n"
		str += "\n"
	}
	for i := 0; i < df.nRows; i++ {
		str += strconv.Itoa(i+1) + ":\t"
		for j := 0; j < df.nCols; j++ {
			str += fmt.Sprint(df.columns[j].row[i])
			str += "\t"
		}
		str += "\n"
	}
	return str
}

// Column Definition
// =================
type Column struct {
	row     []interface{}
	colType string
}

// Column Methods
// ==============
func (c Column) String() string {
	return fmt.Sprint(c.row)
}

func (c *Column) parseType(t string) error {
	var newRows interface{}
	switch t {
	case "int":
		newRows = []int{}
	case "float":
		newRows = []float64{}
	case "string":
		newRows = []string{}
	}
	for _, v := range c.row {
		r := fmt.Sprint(v)
		switch t {
		case "int":
			i, err := strconv.Atoi(r)
			if err != nil {
				return err
			}
			newRows = append(newRows.([]int), i)
		case "float":
			i, err := strconv.ParseFloat(r, 64)
			if err != nil {
				return err
			}
			newRows = append(newRows.([]float64), i)
		case "string":
			newRows = append(newRows.([]string), r)
		}
	}
	c.fillColumn(newRows)
	return nil
}

// TODO: Should this return an error?
func (c *Column) fillColumn(values interface{}) {
	switch reflect.TypeOf(values).Kind() {
	case reflect.Slice:
		s := reflect.ValueOf(values)
		for i := 0; i < s.Len(); i++ {
			c.row = append(c.row, s.Index(i).Interface())
			c.colType = fmt.Sprint(s.Index(i).Type())
		}
	}
}
<commit_msg>Move colnames inside the individual columns<commit_after>package main

import (
	"encoding/csv"
	"errors"
	"fmt"
	"os"
	"reflect"
	"strconv"
)

// TODO: Write tests!
func main() {
	// Test 02
	//textColumn := []string{"One", "Two", "Three"}
	//intColumn := []int{1, 2, 3}
	//c1 := Column{}
	//c1.fillColumn(textColumn)
	//c2 := Column{}
	//c2.fillColumn(intColumn)
	//df := DataFrame{
	//columns:  []Column{c1, c2},
	//nCols:    2,
	//nRows:    3,
	//colnames: []string{"Text", "Ints"},
	//}
	//fmt.Println(df)

	// Test 01
	//in := `A,B,C,D
	//1,2,3,4
	//5,6,7,8`
	//df := DataFrame{}
	//r := csv.NewReader(strings.NewReader(in))
	//records, err := r.ReadAll()
	//if err != nil {
	//panic(err)
	//}
	//err = df.loadAndParse(records, []string{"string", "int", "string", "int"})
	//if err != nil {
	//panic(err)
	//}

	//for _, v := range df.columns {
	//fmt.Println(v)
	//}
	//fmt.Println(df)

	// Test 03
	df := DataFrame{}
	csvfile, err := os.Open("example.csv")
	if err != nil {
		fmt.Println(err)
		return
	}
	r := csv.NewReader(csvfile)
	records, err := r.ReadAll()
	if err != nil {
		panic(err)
	}
	df.loadData(records)
	fmt.Println(df)
	for _, v := range df.columns {
		fmt.Println(v.colType)
		fmt.Println(v.maxCharLength)
	}
}

// DataFrame Definition
// ====================
type DataFrame struct {
	columns []Column
	nCols   int
	nRows   int
}

// DataFrame Methods
// =================
func (df *DataFrame) loadAndParse(records [][]string, types []string) error {
	err := df.loadData(records)
	if err != nil {
		return err
	}
	if df.nCols != len(types) {
		return errors.New("Number of columns different from number of types")
	}
	for k, v := range df.columns {
		v.parseType(types[k])
		df.columns[k].colType = types[k]
	}
	return nil
}

func (df *DataFrame) loadData(records [][]string) error {
	// Get DataFrame dimensions
	nRows := len(records) - 1
	if nRows <= 0 {
		return errors.New("Empty dataframe")
	}
	nCols := len(records[0])

	// Generate a virtual df to store the temporary values
	newDf := DataFrame{
		columns: []Column{},
		nRows:   nRows,
		nCols:   nCols,
	}

	for j := 0; j < nCols; j++ {
		col := []string{}
		for i := 1; i < nRows+1; i++ {
			col = append(col, records[i][j])
		}
		column := Column{}
		column.colName = records[0][j]
		column.maxCharLength = len(column.colName)
		column.fillColumn(col)
		newDf.columns = append(newDf.columns, column)
	}
	*df = newDf
	return nil
}

func (df DataFrame) colnames() (colnames []string) {
	for _, v := range df.columns {
		colnames = append(colnames, v.colName)
	}
	return
}

// TODO: Truncate output for the same tabular format?
func (df DataFrame) String() (str string) {
	if len(df.colnames()) != 0 {
		str += "\t"
		for _, v := range df.colnames() {
			str += v
			str += "\t"
		}
		str += "\n"
		str += "\n"
	}
	for i := 0; i < df.nRows; i++ {
		str += strconv.Itoa(i+1) + ":\t"
		for j := 0; j < df.nCols; j++ {
			str += fmt.Sprint(df.columns[j].row[i])
			str += "\t"
		}
		str += "\n"
	}
	return str
}

// Column Definition
// =================
type Column struct {
	row           []interface{}
	colType       string
	colName       string
	maxCharLength int
}

// Column Methods
// ==============
func (c Column) String() string {
	return fmt.Sprint(c.row)
}

func (c *Column) parseType(t string) error {
	var newRows interface{}
	switch t {
	case "int":
		newRows = []int{}
	case "float":
		newRows = []float64{}
	case "string":
		newRows = []string{}
	}
	for _, v := range c.row {
		r := fmt.Sprint(v)
		switch t {
		case "int":
			i, err := strconv.Atoi(r)
			if err != nil {
				return err
			}
			newRows = append(newRows.([]int), i)
		case "float":
			i, err := strconv.ParseFloat(r, 64)
			if err != nil {
				return err
			}
			newRows = append(newRows.([]float64), i)
		case "string":
			newRows = append(newRows.([]string), r)
		}
	}
	c.fillColumn(newRows)
	return nil
}

// TODO: Should this return an error?
func (c *Column) fillColumn(values interface{}) {
	switch reflect.TypeOf(values).Kind() {
	case reflect.Slice:
		s := reflect.ValueOf(values)
		for i := 0; i < s.Len(); i++ {
			c.row = append(c.row, s.Index(i).Interface())
			c.colType = fmt.Sprint(s.Index(i).Type())
			rowStr := fmt.Sprint(s.Index(i).Interface())
			if len(rowStr) > c.maxCharLength {
				c.maxCharLength = len(rowStr)
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2015-2017 Jeevanandam M (jeeva@myjeeva.com), All rights reserved.
// resty source code and usage is governed by a MIT style
// license that can be found in the LICENSE file.

package resty

import (
	"bytes"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"reflect"
	"strings"
)

//
// Request Middleware(s)
//

func parseRequestURL(c *Client, r *Request) error {
	// Parsing request URL
	reqURL, err := url.Parse(r.URL)
	if err != nil {
		return err
	}

	// If Request.Url is relative path then added c.HostUrl into
	// the request URL otherwise Request.Url will be used as-is
	if !reqURL.IsAbs() {
		if !strings.HasPrefix(r.URL, "/") {
			r.URL = "/" + r.URL
		}

		reqURL, err = url.Parse(c.HostURL + r.URL)
		if err != nil {
			return err
		}
	}

	// Adding Query Param
	query := reqURL.Query()
	for k, v := range c.QueryParam {
		for _, iv := range v {
			query.Add(k, iv)
		}
	}

	for k, v := range r.QueryParam {
		// remove query param from client level by key
		// since overrides happens for that key in the request
		query.Del(k)

		for _, iv := range v {
			query.Add(k, iv)
		}
	}

	reqURL.RawQuery = query.Encode()
	r.URL = reqURL.String()

	return nil
}

func parseRequestHeader(c *Client, r *Request) error {
	hdr := http.Header{}
	for k := range c.Header {
		hdr.Set(k, c.Header.Get(k))
	}
	for k := range r.Header {
		hdr.Set(k, r.Header.Get(k))
	}

	if IsStringEmpty(hdr.Get(hdrUserAgentKey)) {
		hdr.Set(hdrUserAgentKey, fmt.Sprintf(hdrUserAgentValue, Version))
	}

	if IsStringEmpty(hdr.Get(hdrAcceptKey)) && !IsStringEmpty(hdr.Get(hdrContentTypeKey)) {
		hdr.Set(hdrAcceptKey, hdr.Get(hdrContentTypeKey))
	}

	r.Header = hdr

	return nil
}

func parseRequestBody(c *Client, r *Request) (err error) {
	if isPayloadSupported(r.Method, c.AllowGetMethodPayload) {
		// Handling Multipart
		if r.isMultiPart && !(r.Method == MethodPatch) {
			if err = handleMultipart(c, r); err != nil {
				return
			}

			goto CL
		}

		// Handling Form Data
		if len(c.FormData) > 0 || len(r.FormData) > 0 {
			handleFormData(c, r)

			goto CL
		}

		// Handling Request body
		if r.Body != nil {
			handleContentType(c, r)

			if err = handleRequestBody(c, r); err != nil {
				return
			}
		}
	} else {
		r.Header.Del(hdrContentTypeKey)
	}

CL:
	// by default resty won't set content length, you can if you want to :)
	if (c.setContentLength || r.setContentLength) && r.bodyBuf != nil {
		r.Header.Set(hdrContentLengthKey, fmt.Sprintf("%d", r.bodyBuf.Len()))
	}

	return
}

func createHTTPRequest(c *Client, r *Request) (err error) {
	if r.bodyBuf == nil {
		r.RawRequest, err = http.NewRequest(r.Method, r.URL, nil)
	} else {
		r.RawRequest, err = http.NewRequest(r.Method, r.URL, r.bodyBuf)
	}

	if err != nil {
		return
	}

	// Assign close connection option
	r.RawRequest.Close = c.closeConnection

	// Add headers into http request
	r.RawRequest.Header = r.Header

	// Add cookies into http request
	for _, cookie := range c.Cookies {
		r.RawRequest.AddCookie(cookie)
	}

	// it's for non-http scheme option
	if r.RawRequest.URL != nil && r.RawRequest.URL.Scheme == "" {
		r.RawRequest.URL.Scheme = c.scheme
		r.RawRequest.URL.Host = r.URL
	}

	// Use context if it was specified
	r.addContextIfAvailable()

	return
}

func addCredentials(c *Client, r *Request) error {
	var isBasicAuth bool
	// Basic Auth
	if r.UserInfo != nil { // takes precedence
		r.RawRequest.SetBasicAuth(r.UserInfo.Username, r.UserInfo.Password)
		isBasicAuth = true
	} else if c.UserInfo != nil {
		r.RawRequest.SetBasicAuth(c.UserInfo.Username, c.UserInfo.Password)
		isBasicAuth = true
	}

	if !c.DisableWarn {
		if isBasicAuth && !strings.HasPrefix(r.URL, "https") {
			c.Log.Println("WARNING - Using Basic Auth in HTTP mode is not secure.")
		}
	}

	// Token Auth
	if !IsStringEmpty(r.Token) { // takes precedence
		r.RawRequest.Header.Set(hdrAuthorizationKey, "Bearer "+r.Token)
	} else if !IsStringEmpty(c.Token) {
		r.RawRequest.Header.Set(hdrAuthorizationKey, "Bearer "+c.Token)
	}

	return nil
}

func requestLogger(c *Client, r *Request) error {
	if c.Debug {
		rr := r.RawRequest
		c.Log.Println()
		c.disableLogPrefix()
		c.Log.Println("---------------------- REQUEST LOG -----------------------")
		c.Log.Printf("%s  %s  %s\n", r.Method, rr.URL.RequestURI(), rr.Proto)
		c.Log.Printf("HOST   : %s", rr.URL.Host)
		c.Log.Println("HEADERS:")
		for h, v := range rr.Header {
			c.Log.Printf("%25s: %v", h, strings.Join(v, ", "))
		}
		c.Log.Printf("BODY   :\n%v", r.fmtBodyString())
		c.Log.Println("----------------------------------------------------------")
		c.enableLogPrefix()
	}

	return nil
}

//‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾
// Response Middleware(s)
//___________________________________

func responseLogger(c *Client, res *Response) error {
	if c.Debug {
		c.Log.Println()
		c.disableLogPrefix()
		c.Log.Println("---------------------- RESPONSE LOG -----------------------")
		c.Log.Printf("STATUS 		: %s", res.Status())
		c.Log.Printf("RECEIVED AT	: %v", res.ReceivedAt())
		c.Log.Printf("RESPONSE TIME	: %v", res.Time())
		c.Log.Println("HEADERS:")
		for h, v := range res.Header() {
			c.Log.Printf("%30s: %v", h, strings.Join(v, ", "))
		}
		if res.Request.isSaveResponse {
			c.Log.Printf("BODY   :\n***** RESPONSE WRITTEN INTO FILE *****")
		} else {
			c.Log.Printf("BODY   :\n%v", res.fmtBodyString())
		}
		c.Log.Println("----------------------------------------------------------")
		c.enableLogPrefix()
	}

	return nil
}

func parseResponseBody(c *Client, res *Response) (err error) {
	// Handles only JSON or XML content type
	ct := firstNonEmpty(res.Header().Get(hdrContentTypeKey), res.Request.fallbackContentType)
	if IsJSONType(ct) || IsXMLType(ct) {
		// Considered as Result
		if res.StatusCode() > 199 && res.StatusCode() < 300 {
			if res.Request.Result != nil {
				err = Unmarshalc(c, ct, res.body, res.Request.Result)
				return
			}
		}

		// Considered as Error
		if res.StatusCode() > 399 {
			// global error interface
			if res.Request.Error == nil && c.Error != nil {
				res.Request.Error = reflect.New(c.Error).Interface()
			}

			if res.Request.Error != nil {
				err = Unmarshalc(c, ct, res.body, res.Request.Error)
			}
		}
	}

	return
}

func handleMultipart(c *Client, r *Request) (err error) {
	r.bodyBuf = acquireBuffer()
	w := multipart.NewWriter(r.bodyBuf)

	for k, v := range c.FormData {
		for _, iv := range v {
			if err = w.WriteField(k, iv); err != nil {
				return err
			}
		}
	}

	for k, v := range r.FormData {
		for _, iv := range v {
			if strings.HasPrefix(k, "@") { // file
				err = addFile(w, k[1:], iv)
				if err != nil {
					return
				}
			} else { // form value
				if err = w.WriteField(k, iv); err != nil {
					return err
				}
			}
		}
	}

	// #21 - adding io.Reader support
	if len(r.multipartFiles) > 0 {
		for _, f := range r.multipartFiles {
			err = addFileReader(w, f)
			if err != nil {
				return
			}
		}
	}

	r.Header.Set(hdrContentTypeKey, w.FormDataContentType())
	err = w.Close()

	return
}

func handleFormData(c *Client, r *Request) {
	formData := url.Values{}

	for k, v := range c.FormData {
		for _, iv := range v {
			formData.Add(k, iv)
		}
	}

	for k, v := range r.FormData {
		// remove form data field from client level by key
		// since overrides happens for that key in the request
		formData.Del(k)

		for _, iv := range v {
			formData.Add(k, iv)
		}
	}

	r.bodyBuf = bytes.NewBuffer([]byte(formData.Encode()))
	r.Header.Set(hdrContentTypeKey, formContentType)
	r.isFormData = true
}

func handleContentType(c *Client, r *Request) {
	contentType := r.Header.Get(hdrContentTypeKey)
	if IsStringEmpty(contentType) {
		contentType = DetectContentType(r.Body)
		r.Header.Set(hdrContentTypeKey, contentType)
	}
}

func handleRequestBody(c *Client, r *Request) (err error) {
	var bodyBytes []byte
	contentType := r.Header.Get(hdrContentTypeKey)
	kind := kindOf(r.Body)
	r.bodyBuf = nil

	if reader, ok := r.Body.(io.Reader); ok {
		r.bodyBuf = acquireBuffer()
		_, err = r.bodyBuf.ReadFrom(reader)
	} else if b, ok := r.Body.([]byte); ok {
		bodyBytes = b
	} else if s, ok := r.Body.(string); ok {
		bodyBytes = []byte(s)
	} else if IsJSONType(contentType) &&
		(kind == reflect.Struct || kind == reflect.Map || kind == reflect.Slice) {
		bodyBytes, err = c.JSONMarshal(r.Body)
	} else if IsXMLType(contentType) && (kind == reflect.Struct) {
		bodyBytes, err = xml.Marshal(r.Body)
	}

	if bodyBytes == nil && r.bodyBuf == nil {
		err = errors.New("Unsupported 'Body' type/value")
	}

	// if any errors during body bytes handling, return it
	if err != nil {
		return
	}

	// []byte into Buffer
	if bodyBytes != nil && r.bodyBuf == nil {
		r.bodyBuf = acquireBuffer()
		_, _ = r.bodyBuf.Write(bodyBytes)
	}

	return
}

func saveResponseIntoFile(c *Client, res *Response) error {
	if res.Request.isSaveResponse {
		file := ""

		if len(c.outputDirectory) > 0 && !filepath.IsAbs(res.Request.outputFile) {
			file += c.outputDirectory + string(filepath.Separator)
		}

		file = filepath.Clean(file + res.Request.outputFile)
		if err := createDirectory(filepath.Dir(file)); err != nil {
			return err
		}

		outFile, err := os.Create(file)
		if err != nil {
			return err
		}
		defer func() {
			_ = outFile.Close()
		}()

		// io.Copy reads maximum 32kb size, it is perfect for large file download too
		defer func() {
			_ = res.RawResponse.Body.Close()
		}()
		written, err := io.Copy(outFile, res.RawResponse.Body)
		if err != nil {
			return err
		}

		res.size = written
	}

	return nil
}
<commit_msg>#97 removed the line which cleans up the content-type header when payload is not present<commit_after>// Copyright (c) 2015-2017 Jeevanandam M (jeeva@myjeeva.com), All rights reserved.
// resty source code and usage is governed by a MIT style
// license that can be found in the LICENSE file.

package resty

import (
	"bytes"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"reflect"
	"strings"
)

//
// Request Middleware(s)
//

func parseRequestURL(c *Client, r *Request) error {
	// Parsing request URL
	reqURL, err := url.Parse(r.URL)
	if err != nil {
		return err
	}

	// If Request.Url is relative path then added c.HostUrl into
	// the request URL otherwise Request.Url will be used as-is
	if !reqURL.IsAbs() {
		if !strings.HasPrefix(r.URL, "/") {
			r.URL = "/" + r.URL
		}

		reqURL, err = url.Parse(c.HostURL + r.URL)
		if err != nil {
			return err
		}
	}

	// Adding Query Param
	query := reqURL.Query()
	for k, v := range c.QueryParam {
		for _, iv := range v {
			query.Add(k, iv)
		}
	}

	for k, v := range r.QueryParam {
		// remove query param from client level by key
		// since overrides happens for that key in the request
		query.Del(k)

		for _, iv := range v {
			query.Add(k, iv)
		}
	}

	reqURL.RawQuery = query.Encode()
	r.URL = reqURL.String()

	return nil
}

func parseRequestHeader(c *Client, r *Request) error {
	hdr := http.Header{}
	for k := range c.Header {
		hdr.Set(k, c.Header.Get(k))
	}
	for k := range r.Header {
		hdr.Set(k, r.Header.Get(k))
	}

	if IsStringEmpty(hdr.Get(hdrUserAgentKey)) {
		hdr.Set(hdrUserAgentKey, fmt.Sprintf(hdrUserAgentValue, Version))
	}

	if IsStringEmpty(hdr.Get(hdrAcceptKey)) && !IsStringEmpty(hdr.Get(hdrContentTypeKey)) {
		hdr.Set(hdrAcceptKey, hdr.Get(hdrContentTypeKey))
	}

	r.Header = hdr

	return nil
}

func parseRequestBody(c *Client, r *Request) (err error) {
	if isPayloadSupported(r.Method, c.AllowGetMethodPayload) {
		// Handling Multipart
		if r.isMultiPart && !(r.Method == MethodPatch) {
			if err = handleMultipart(c, r); err != nil {
				return
			}

			goto CL
		}

		// Handling Form Data
		if len(c.FormData) > 0 || len(r.FormData) > 0 {
			handleFormData(c, r)

			goto CL
		}

		// Handling Request body
		if r.Body != nil {
			handleContentType(c, r)

			if err = handleRequestBody(c, r); err != nil {
				return
			}
		}
	}

CL:
	// by default resty won't set content length, you can if you want to :)
	if (c.setContentLength || r.setContentLength) && r.bodyBuf != nil {
		r.Header.Set(hdrContentLengthKey, fmt.Sprintf("%d", r.bodyBuf.Len()))
	}

	return
}

func createHTTPRequest(c *Client, r *Request) (err error) {
	if r.bodyBuf == nil {
		r.RawRequest, err = http.NewRequest(r.Method, r.URL, nil)
	} else {
		r.RawRequest, err = http.NewRequest(r.Method, r.URL, r.bodyBuf)
	}

	if err != nil {
		return
	}

	// Assign close connection option
	r.RawRequest.Close = c.closeConnection

	// Add headers into http request
	r.RawRequest.Header = r.Header

	// Add cookies into http request
	for _, cookie := range c.Cookies {
		r.RawRequest.AddCookie(cookie)
	}

	// it's for non-http scheme option
	if r.RawRequest.URL != nil && r.RawRequest.URL.Scheme == "" {
		r.RawRequest.URL.Scheme = c.scheme
		r.RawRequest.URL.Host = r.URL
	}

	// Use context if it was specified
	r.addContextIfAvailable()

	return
}

func addCredentials(c *Client, r *Request) error {
	var isBasicAuth bool
	// Basic Auth
	if r.UserInfo != nil { // takes precedence
		r.RawRequest.SetBasicAuth(r.UserInfo.Username, r.UserInfo.Password)
		isBasicAuth = true
	} else if c.UserInfo != nil {
		r.RawRequest.SetBasicAuth(c.UserInfo.Username, c.UserInfo.Password)
		isBasicAuth = true
	}

	if !c.DisableWarn {
		if isBasicAuth && !strings.HasPrefix(r.URL, "https") {
			c.Log.Println("WARNING - Using Basic Auth in HTTP mode is not secure.")
		}
	}

	// Token Auth
	if !IsStringEmpty(r.Token) { // takes precedence
		r.RawRequest.Header.Set(hdrAuthorizationKey, "Bearer "+r.Token)
	} else if !IsStringEmpty(c.Token) {
		r.RawRequest.Header.Set(hdrAuthorizationKey, "Bearer "+c.Token)
	}

	return nil
}

func requestLogger(c *Client, r *Request) error {
	if c.Debug {
		rr := r.RawRequest
		c.Log.Println()
		c.disableLogPrefix()
		c.Log.Println("---------------------- REQUEST LOG -----------------------")
		c.Log.Printf("%s  %s  %s\n", r.Method, rr.URL.RequestURI(), rr.Proto)
		c.Log.Printf("HOST   : %s", rr.URL.Host)
		c.Log.Println("HEADERS:")
		for h, v := range rr.Header {
			c.Log.Printf("%25s: %v", h, strings.Join(v, ", "))
		}
		c.Log.Printf("BODY   :\n%v", r.fmtBodyString())
		c.Log.Println("----------------------------------------------------------")
		c.enableLogPrefix()
	}

	return nil
}

//‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾
// Response Middleware(s)
//___________________________________

func responseLogger(c *Client, res *Response) error {
	if c.Debug {
		c.Log.Println()
		c.disableLogPrefix()
		c.Log.Println("---------------------- RESPONSE LOG -----------------------")
		c.Log.Printf("STATUS 		: %s", res.Status())
		c.Log.Printf("RECEIVED AT	: %v", res.ReceivedAt())
		c.Log.Printf("RESPONSE TIME	: %v", res.Time())
		c.Log.Println("HEADERS:")
		for h, v := range res.Header() {
			c.Log.Printf("%30s: %v", h, strings.Join(v, ", "))
		}
		if res.Request.isSaveResponse {
			c.Log.Printf("BODY   :\n***** RESPONSE WRITTEN INTO FILE *****")
		} else {
			c.Log.Printf("BODY   :\n%v", res.fmtBodyString())
		}
		c.Log.Println("----------------------------------------------------------")
		c.enableLogPrefix()
	}

	return nil
}

func parseResponseBody(c *Client, res *Response) (err error) {
	// Handles only JSON or XML content type
	ct := firstNonEmpty(res.Header().Get(hdrContentTypeKey), res.Request.fallbackContentType)
	if IsJSONType(ct) || IsXMLType(ct) {
		// Considered as Result
		if res.StatusCode() > 199 && res.StatusCode() < 300 {
			if res.Request.Result != nil {
				err = Unmarshalc(c, ct, res.body, res.Request.Result)
				return
			}
		}

		// Considered as Error
		if res.StatusCode() > 399 {
			// global error interface
			if res.Request.Error == nil && c.Error != nil {
				res.Request.Error = reflect.New(c.Error).Interface()
			}

			if res.Request.Error != nil {
				err = Unmarshalc(c, ct, res.body, res.Request.Error)
			}
		}
	}

	return
}

func handleMultipart(c *Client, r *Request) (err error) {
	r.bodyBuf = acquireBuffer()
	w := multipart.NewWriter(r.bodyBuf)

	for k, v := range c.FormData {
		for _, iv := range v {
			if err = w.WriteField(k, iv); err != nil {
				return err
			}
		}
	}

	for k, v := range r.FormData {
		for _, iv := range v {
			if strings.HasPrefix(k, "@") { // file
				err = addFile(w, k[1:], iv)
				if err != nil {
					return
				}
			} else { // form value
				if err = w.WriteField(k, iv); err != nil {
					return err
				}
			}
		}
	}

	// #21 - adding io.Reader support
	if len(r.multipartFiles) > 0 {
		for _, f := range r.multipartFiles {
			err = addFileReader(w, f)
			if err != nil {
				return
			}
		}
	}

	r.Header.Set(hdrContentTypeKey, w.FormDataContentType())
	err = w.Close()

	return
}

func handleFormData(c *Client, r *Request) {
	formData := url.Values{}

	for k, v := range c.FormData {
		for _, iv := range v {
			formData.Add(k, iv)
		}
	}

	for k, v := range r.FormData {
		// remove form data field from client level by key
		// since overrides happens for that key in the request
		formData.Del(k)

		for _, iv := range v {
			formData.Add(k, iv)
		}
	}

	r.bodyBuf = bytes.NewBuffer([]byte(formData.Encode()))
	r.Header.Set(hdrContentTypeKey, formContentType)
	r.isFormData = true
}

func handleContentType(c *Client, r *Request) {
	contentType := r.Header.Get(hdrContentTypeKey)
	if IsStringEmpty(contentType) {
		contentType = DetectContentType(r.Body)
		r.Header.Set(hdrContentTypeKey, contentType)
	}
}

func handleRequestBody(c *Client, r *Request) (err error) {
	var bodyBytes []byte
	contentType := r.Header.Get(hdrContentTypeKey)
	kind := kindOf(r.Body)
	r.bodyBuf = nil

	if reader, ok := r.Body.(io.Reader); ok {
		r.bodyBuf = acquireBuffer()
		_, err = r.bodyBuf.ReadFrom(reader)
	} else if b, ok := r.Body.([]byte); ok {
		bodyBytes = b
	} else if s, ok := r.Body.(string); ok {
		bodyBytes = []byte(s)
	} else if IsJSONType(contentType) &&
		(kind == reflect.Struct || kind == reflect.Map || kind == reflect.Slice) {
		bodyBytes, err = c.JSONMarshal(r.Body)
	} else if IsXMLType(contentType) && (kind == reflect.Struct) {
		bodyBytes, err = xml.Marshal(r.Body)
	}

	if bodyBytes == nil && r.bodyBuf == nil {
		err = errors.New("Unsupported 'Body' type/value")
	}

	// if any errors during body bytes handling, return it
	if err != nil {
		return
	}

	// []byte into Buffer
	if bodyBytes != nil && r.bodyBuf == nil {
		r.bodyBuf = acquireBuffer()
		_, _ = r.bodyBuf.Write(bodyBytes)
	}

	return
}

func saveResponseIntoFile(c *Client, res *Response) error {
	if res.Request.isSaveResponse {
		file := ""

		if len(c.outputDirectory) > 0 && !filepath.IsAbs(res.Request.outputFile) {
			file += c.outputDirectory + string(filepath.Separator)
		}

		file = filepath.Clean(file + res.Request.outputFile)
		if err := createDirectory(filepath.Dir(file)); err != nil {
			return err
		}

		outFile, err := os.Create(file)
		if err != nil {
			return err
		}
		defer func() {
			_ = outFile.Close()
		}()

		// io.Copy reads maximum 32kb size, it is perfect for large file download too
		defer func() {
			_ = res.RawResponse.Body.Close()
		}()
		written, err := io.Copy(outFile, res.RawResponse.Body)
		if err != nil {
			return err
		}

		res.size = written
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/nindalf/gotop/daemon"
)

func main() {
	fmt.Println(daemon.UpSince())
}
<commit_msg>Add CPU, Memory and Uptime to command line<commit_after>package main

import (
	"fmt"
	"time"

	ui "github.com/gizak/termui"
	"github.com/nindalf/gotop/daemon"
)

func updateMemCPU(gcpu, gmem *ui.Gauge) {
	done := make(chan struct{})
	memInfoChan, errc := daemon.TotalMemory(done, daemon.Delay)
	cpuInfoChan, errc := daemon.TotalCPU(done, daemon.Delay)
	timeout := time.After(2 * daemon.Delay)
	defer func() {
		close(done)
		// Necessary to read from error channel to prevent sending goroutine going into deadlock
		<-errc
	}()
	for {
		select {
		case memInfo := <-memInfoChan:
			pc := int(100 * (float64(memInfo.MemTotal-memInfo.MemFree) / float64(memInfo.MemTotal)))
			gmem.Percent = pc
			ui.Render(ui.Body)
		case cpuInfo := <-cpuInfoChan:
			pc := int(cpuInfo.AverageUtilization)
			gcpu.Percent = pc
			ui.Render(ui.Body)
		case err := <-errc:
			if err != nil {
			}
			return
		case <-timeout:
		}
	}
}

func updateUptime(p *ui.Par) {
	for {
		<-time.After(2 * daemon.Delay)
		ut, _ := daemon.Uptime()
		p.Text = ut.String()
		ui.Render(ui.Body)
	}
}

func main() {
	err := ui.Init()
	fmt.Println(daemon.UpSince())
	if err != nil {
		fmt.Println("Could not initialise UI")
	}
	defer ui.Close()

	ut, _ := daemon.Uptime()
	p := ui.NewPar(ut.String())
	p.Height = 3
	p.Width = 50
	p.TextFgColor = ui.ColorWhite
	p.Border.Label = "Uptime"
	p.Border.FgColor = ui.ColorCyan

	g0 := ui.NewGauge()
	g0.Percent = 40
	g0.Width = 50
	g0.Height = 3
	g0.Border.Label = "Memory"
	g0.BarColor = ui.ColorRed
	g0.Border.FgColor = ui.ColorWhite
	g0.Border.LabelFgColor = ui.ColorCyan

	g2 := ui.NewGauge()
	g2.Percent = 60
	g2.Width = 50
	g2.Height = 3
	g2.PercentColor = ui.ColorBlue
	g2.Y = 3
	g2.Border.Label = "CPU"
	g2.BarColor = ui.ColorYellow
	g2.Border.FgColor = ui.ColorWhite

	ui.Body.AddRows(ui.NewRow(ui.NewCol(6, 0, g0), ui.NewCol(6, 0, p)),
		ui.NewRow(ui.NewCol(6, 0, g2)))
	ui.Body.Align()
	ui.Render(ui.Body)
	go updateMemCPU(g2, g0)
	go updateUptime(p)
	<-ui.EventCh()
}
<|endoftext|>


<commit_before>/*
Middleware for Goji that reports HTTP requests to Apinalytics

Add it to your Goji mux m as follows.

    m.Use(BuildMiddleWare(myAppId, nil))

To add your own data to the events reported add a callback.

    callback := func(c *web.C, data map[string]interface{}, r *http.Request) {
        data["my_parameter"] = c.Env["important_info"].(string)
    }

    m.Use(BuildMiddleWare(myAppId, callback))

*/
package apinalytics_client

import (
	"net/http"
	"time"

	"github.com/zenazn/goji/web"
)

/*
Build a middleware function that reports HTTP requests to Apinalytics.

The event the middleware sends is as follows.

    {
        "url":         r.RequestURI,
        "path":        r.URL.Path,
        "method":      r.Method,
        "status_code": w.Status,  // Http status code from response
        "duration_ns": time.Since(start).Nanoseconds(),
        "user_agent":  r.UserAgent(),
        "header":      r.Header,
    }

The callback function allows you to add your own data to the event recorded.  Use it as follows to add events to the "apievents" collection.

    callback := func(c *web.C, data map[string]interface{}, r *http.Request) {
        data["my_parameter"] = c.Env["important_info"].(string)
    }

    m.Use(BuildMiddleWare(myApplicationId, callback))

*/
func BuildMiddleWare(applicationId string,
	callback func(c *web.C, event *AnalyticsEvent, r *http.Request),
) func(c *web.C, h http.Handler) http.Handler {
	sender := NewSender(applicationId)

	// Return the middleware that references the analytics queue we just made
	return func(c *web.C, h http.Handler) http.Handler {
		handler := func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()
			ww := &StatusTrackingResponseWriter{w, http.StatusOK}

			h.ServeHTTP(ww, r)

			event := &AnalyticsEvent{
				Timestamp:  time.Now(),
				Method:     r.Method,
				Url:        r.RequestURI,
				ResponseUS: int(time.Since(start).Nanoseconds() / 1000),
				StatusCode: ww.Status,
			}
			// "path":        r.URL.Path,
			// "user_agent":  r.UserAgent(),
			// "header":      r.Header,
			// Get more data for the analytics event
			if callback != nil {
				callback(c, event, r)
			}

			sender.Queue(event)
		}
		return http.HandlerFunc(handler)
	}
}
<commit_msg>Get function name from c.Env<commit_after>/*
Middleware for Goji that reports HTTP requests to Apinalytics

Add it to your Goji mux m as follows.

    m.Use(BuildMiddleWare(myAppId, nil))

To add your own data to the events reported add a callback.

    callback := func(c *web.C, data map[string]interface{}, r *http.Request) {
        data["my_parameter"] = c.Env["important_info"].(string)
    }

    m.Use(BuildMiddleWare(myAppId, callback))

*/
package apinalytics_client

import (
	"net/http"
	"time"

	"github.com/zenazn/goji/web"
)

/*
Build a middleware function that reports HTTP requests to Apinalytics.

The event the middleware sends is as follows.

    {
        "url":         r.RequestURI,
        "path":        r.URL.Path,
        "method":      r.Method,
        "status_code": w.Status,  // Http status code from response
        "duration_ns": time.Since(start).Nanoseconds(),
        "user_agent":  r.UserAgent(),
        "header":      r.Header,
    }

The callback function allows you to add your own data to the event recorded.  Use it as follows to add events to the "apievents" collection.

    callback := func(c *web.C, data map[string]interface{}, r *http.Request) {
        data["my_parameter"] = c.Env["important_info"].(string)
    }

    m.Use(BuildMiddleWare(myApplicationId, callback))

*/
func BuildMiddleWare(applicationId string,
	callback func(c *web.C, event *AnalyticsEvent, r *http.Request),
) func(c *web.C, h http.Handler) http.Handler {
	sender := NewSender(applicationId)

	// Return the middleware that references the analytics queue we just made
	return func(c *web.C, h http.Handler) http.Handler {
		handler := func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()
			ww := &StatusTrackingResponseWriter{w, http.StatusOK}

			h.ServeHTTP(ww, r)

			function := c.Env["function"].(string)
			if function == "" {
				function = "unknown"
			}
			event := &AnalyticsEvent{
				Timestamp:  time.Now(),
				Method:     r.Method,
				Url:        r.RequestURI,
				Function:   function,
				ResponseUS: int(time.Since(start).Nanoseconds() / 1000),
				StatusCode: ww.Status,
			}
			// "path":        r.URL.Path,
			// "user_agent":  r.UserAgent(),
			// "header":      r.Header,
			// Get more data for the analytics event
			if callback != nil {
				callback(c, event, r)
			}

			sender.Queue(event)
		}
		return http.HandlerFunc(handler)
	}
}
<|endoftext|>


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

package main

import (
	"bytes"
	"encoding/json"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"net/http"
)

func init() {
	http.HandleFunc("/fmt", fmtHandler)
}

type fmtResponse struct {
	Body  string
	Error string
}

func fmtHandler(w http.ResponseWriter, r *http.Request) {
	resp := new(fmtResponse)
	body, err := gofmt(r.FormValue("body"))
	if err != nil {
		resp.Error = err.Error()
	} else {
		resp.Body = body
	}
	json.NewEncoder(w).Encode(resp)
}

func gofmt(body string) (string, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "prog.go", body, parser.ParseComments)
	if err != nil {
		return "", err
	}
	ast.SortImports(fset, f)
	var buf bytes.Buffer
	err = printer.Fprint(&buf, fset, f)
	if err != nil {
		return "", err
	}
	return buf.String(), nil
}
<commit_msg>go-tour: fix format function<commit_after>// Copyright 2012 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/json"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"net/http"
)

func init() {
	http.HandleFunc("/fmt", fmtHandler)
}

type fmtResponse struct {
	Body  string
	Error string
}

func fmtHandler(w http.ResponseWriter, r *http.Request) {
	resp := new(fmtResponse)
	body, err := gofmt(r.FormValue("body"))
	if err != nil {
		resp.Error = err.Error()
	} else {
		resp.Body = body
	}
	json.NewEncoder(w).Encode(resp)
}

func gofmt(body string) (string, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "prog.go", body, parser.ParseComments)
	if err != nil {
		return "", err
	}
	ast.SortImports(fset, f)
	var buf bytes.Buffer
	config := &printer.Config{Mode: printer.UseSpaces | printer.TabIndent, Tabwidth: 8}
	err = config.Fprint(&buf, fset, f)
	if err != nil {
		return "", err
	}
	return buf.String(), nil
}
<|endoftext|>