<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package cloudup

import (
	"fmt"
	"strings"
	"testing"

	kopsapi "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/util/pkg/architectures"
)

func buildMinimalNodeInstanceGroup(subnets ...string) *kopsapi.InstanceGroup {
	g := &kopsapi.InstanceGroup{}
	g.ObjectMeta.Name = "nodes"
	g.Spec.Role = kopsapi.InstanceGroupRoleNode
	g.Spec.MinSize = fi.Int32(1)
	g.Spec.MaxSize = fi.Int32(1)
	g.Spec.Image = "my-image"
	g.Spec.Subnets = subnets

	return g
}

func buildMinimalMasterInstanceGroup(subnet string) *kopsapi.InstanceGroup {
	g := &kopsapi.InstanceGroup{}
	g.ObjectMeta.Name = "master-" + subnet
	g.Spec.Role = kopsapi.InstanceGroupRoleMaster
	g.Spec.MinSize = fi.Int32(1)
	g.Spec.MaxSize = fi.Int32(1)
	g.Spec.Image = "my-image"
	g.Spec.Subnets = []string{subnet}

	return g
}

func TestPopulateInstanceGroup_Name_Required(t *testing.T) {
	_, cluster := buildMinimalCluster()
	g := buildMinimalNodeInstanceGroup()
	g.ObjectMeta.Name = ""

	channel := &kopsapi.Channel{}

	expectErrorFromPopulateInstanceGroup(t, cluster, g, channel, "objectMeta.name")
}

func TestPopulateInstanceGroup_Role_Required(t *testing.T) {
	_, cluster := buildMinimalCluster()
	g := buildMinimalNodeInstanceGroup()
	g.Spec.Role = ""

	channel := &kopsapi.Channel{}

	expectErrorFromPopulateInstanceGroup(t, cluster, g, channel, "spec.role")
}

func TestPopulateInstanceGroup_AddTaintsCollision(t *testing.T) {
	_, cluster := buildMinimalCluster()
	input := buildMinimalNodeInstanceGroup()
	input.Spec.Taints = []string{"nvidia.com/gpu:NoSchedule"}
	input.Spec.MachineType = "g4dn.xlarge"
	cluster.Spec.Containerd.NvidiaGPU = &kopsapi.NvidiaGPUConfig{Enabled: fi.Bool(true)}

	channel := &kopsapi.Channel{}

	cloud, err := BuildCloud(cluster)
	if err != nil {
		t.Fatalf("error from BuildCloud: %v", err)
	}
	output, err := PopulateInstanceGroupSpec(cluster, input, cloud, channel)
	if err != nil {
		t.Fatalf("error from PopulateInstanceGroupSpec: %v", err)
	}
	if len(output.Spec.Taints) != 1 {
		t.Errorf("Expected only 1 taint, got %d", len(output.Spec.Taints))
	}
}

func TestPopulateInstanceGroup_AddTaints(t *testing.T) {
	_, cluster := buildMinimalCluster()
	input := buildMinimalNodeInstanceGroup()
	input.Spec.MachineType = "g4dn.xlarge"
	cluster.Spec.Containerd.NvidiaGPU = &kopsapi.NvidiaGPUConfig{Enabled: fi.Bool(true)}

	channel := &kopsapi.Channel{}

	cloud, err := BuildCloud(cluster)
	if err != nil {
		t.Fatalf("error from BuildCloud: %v", err)
	}
	output, err := PopulateInstanceGroupSpec(cluster, input, cloud, channel)
	if err != nil {
		t.Fatalf("error from PopulateInstanceGroupSpec: %v", err)
	}
	if len(output.Spec.Taints) != 1 {
		t.Errorf("Expected only 1 taint, got %d", len(output.Spec.Taints))
	}
}

func expectErrorFromPopulateInstanceGroup(t *testing.T, cluster *kopsapi.Cluster, g *kopsapi.InstanceGroup, channel *kopsapi.Channel, message string) {
	cloud, err := BuildCloud(cluster)
	if err != nil {
		t.Fatalf("error from BuildCloud: %v", err)
	}

	_, err = PopulateInstanceGroupSpec(cluster, g, cloud, channel)
	if err == nil {
		t.Fatalf("Expected error from PopulateInstanceGroup")
	}
	actualMessage := fmt.Sprintf("%v", err)
	if !strings.Contains(actualMessage, message) {
		t.Fatalf("Expected error %q, got %q", message, actualMessage)
	}
}

func TestMachineArchitecture(t *testing.T) {
	tests := []struct {
		machineType string
		arch        architectures.Architecture
		err         error
	}{
		{
			machineType: "t2.micro",
			arch:        architectures.ArchitectureAmd64,
			err:         nil,
		},
		{
			machineType: "t3.micro",
			arch:        architectures.ArchitectureAmd64,
			err:         nil,
		},
		{
			machineType: "a1.large",
			arch:        architectures.ArchitectureArm64,
			err:         nil,
		},
	}
	for _, test := range tests {
		t.Run(fmt.Sprintf("%s-%s", test.machineType, test.arch), func(t *testing.T) {
			_, cluster := buildMinimalCluster()
			cloud, err := BuildCloud(cluster)
			if err != nil {
				t.Fatalf("error from BuildCloud: %v", err)
			}

			arch, err := MachineArchitecture(cloud, test.machineType)
			if err != test.err {
				t.Errorf("actual error %q differs from expected error %q", err, test.err)
				return
			}

			if arch != test.arch {
				t.Errorf("actual architecture %q differs from expected architecture %q", arch, test.arch)
				return
			}
		})
	}
}
<commit_msg>Add test for ensuring taints are merged correctly<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package cloudup

import (
	"fmt"
	"strings"
	"testing"

	kopsapi "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/util/pkg/architectures"
)

func buildMinimalNodeInstanceGroup(subnets ...string) *kopsapi.InstanceGroup {
	g := &kopsapi.InstanceGroup{}
	g.ObjectMeta.Name = "nodes"
	g.Spec.Role = kopsapi.InstanceGroupRoleNode
	g.Spec.MinSize = fi.Int32(1)
	g.Spec.MaxSize = fi.Int32(1)
	g.Spec.Image = "my-image"
	g.Spec.Subnets = subnets

	return g
}

func buildMinimalMasterInstanceGroup(subnet string) *kopsapi.InstanceGroup {
	g := &kopsapi.InstanceGroup{}
	g.ObjectMeta.Name = "master-" + subnet
	g.Spec.Role = kopsapi.InstanceGroupRoleMaster
	g.Spec.MinSize = fi.Int32(1)
	g.Spec.MaxSize = fi.Int32(1)
	g.Spec.Image = "my-image"
	g.Spec.Subnets = []string{subnet}

	return g
}

func TestPopulateInstanceGroup_Name_Required(t *testing.T) {
	_, cluster := buildMinimalCluster()
	g := buildMinimalNodeInstanceGroup()
	g.ObjectMeta.Name = ""

	channel := &kopsapi.Channel{}

	expectErrorFromPopulateInstanceGroup(t, cluster, g, channel, "objectMeta.name")
}

func TestPopulateInstanceGroup_Role_Required(t *testing.T) {
	_, cluster := buildMinimalCluster()
	g := buildMinimalNodeInstanceGroup()
	g.Spec.Role = ""

	channel := &kopsapi.Channel{}

	expectErrorFromPopulateInstanceGroup(t, cluster, g, channel, "spec.role")
}

// TestPopulateInstanceGroup_AddTaintsCollision ensures we handle IGs with a user configured taint that kOps also adds by default
func TestPopulateInstanceGroup_AddTaintsCollision(t *testing.T) {
	_, cluster := buildMinimalCluster()
	input := buildMinimalNodeInstanceGroup()
	input.Spec.Taints = []string{"nvidia.com/gpu:NoSchedule"}
	input.Spec.MachineType = "g4dn.xlarge"
	cluster.Spec.Containerd.NvidiaGPU = &kopsapi.NvidiaGPUConfig{Enabled: fi.Bool(true)}

	channel := &kopsapi.Channel{}

	cloud, err := BuildCloud(cluster)
	if err != nil {
		t.Fatalf("error from BuildCloud: %v", err)
	}
	output, err := PopulateInstanceGroupSpec(cluster, input, cloud, channel)
	if err != nil {
		t.Fatalf("error from PopulateInstanceGroupSpec: %v", err)
	}
	if len(output.Spec.Kubelet.Taints) != 1 {
		t.Errorf("Expected only 1 taint, got %d", len(output.Spec.Taints))
	}
}

// TestPopulateInstanceGroup_AddTaintsCollision2 ensures we handle taints that are configured in multiple parts of the spec and multiple resources.
// This one also adds a second taint that we should see in the final result
func TestPopulateInstanceGroup_AddTaintsCollision3(t *testing.T) {
	taint := "e2etest:NoSchedule"
	taint2 := "e2etest:NoExecute"
	_, cluster := buildMinimalCluster()
	cluster.Spec.Kubelet = &kopsapi.KubeletConfigSpec{
		Taints: []string{taint, taint2},
	}
	input := buildMinimalNodeInstanceGroup()
	input.Spec.Taints = []string{taint}
	input.Spec.Kubelet = &kopsapi.KubeletConfigSpec{
		Taints: []string{taint},
	}

	channel := &kopsapi.Channel{}

	cloud, err := BuildCloud(cluster)
	if err != nil {
		t.Fatalf("error from BuildCloud: %v", err)
	}
	output, err := PopulateInstanceGroupSpec(cluster, input, cloud, channel)
	if err != nil {
		t.Fatalf("error from PopulateInstanceGroupSpec: %v", err)
	}
	if len(output.Spec.Kubelet.Taints) != 2 {
		t.Errorf("Expected only 2 taints, got %d", len(output.Spec.Kubelet.Taints))
	}
}

func TestPopulateInstanceGroup_AddTaints(t *testing.T) {
	_, cluster := buildMinimalCluster()
	input := buildMinimalNodeInstanceGroup()
	input.Spec.MachineType = "g4dn.xlarge"
	cluster.Spec.Containerd.NvidiaGPU = &kopsapi.NvidiaGPUConfig{Enabled: fi.Bool(true)}

	channel := &kopsapi.Channel{}

	cloud, err := BuildCloud(cluster)
	if err != nil {
		t.Fatalf("error from BuildCloud: %v", err)
	}
	output, err := PopulateInstanceGroupSpec(cluster, input, cloud, channel)
	if err != nil {
		t.Fatalf("error from PopulateInstanceGroupSpec: %v", err)
	}
	if len(output.Spec.Taints) != 1 {
		t.Errorf("Expected only 1 taint, got %d", len(output.Spec.Taints))
	}
}

func expectErrorFromPopulateInstanceGroup(t *testing.T, cluster *kopsapi.Cluster, g *kopsapi.InstanceGroup, channel *kopsapi.Channel, message string) {
	cloud, err := BuildCloud(cluster)
	if err != nil {
		t.Fatalf("error from BuildCloud: %v", err)
	}

	_, err = PopulateInstanceGroupSpec(cluster, g, cloud, channel)
	if err == nil {
		t.Fatalf("Expected error from PopulateInstanceGroup")
	}
	actualMessage := fmt.Sprintf("%v", err)
	if !strings.Contains(actualMessage, message) {
		t.Fatalf("Expected error %q, got %q", message, actualMessage)
	}
}

func TestMachineArchitecture(t *testing.T) {
	tests := []struct {
		machineType string
		arch        architectures.Architecture
		err         error
	}{
		{
			machineType: "t2.micro",
			arch:        architectures.ArchitectureAmd64,
			err:         nil,
		},
		{
			machineType: "t3.micro",
			arch:        architectures.ArchitectureAmd64,
			err:         nil,
		},
		{
			machineType: "a1.large",
			arch:        architectures.ArchitectureArm64,
			err:         nil,
		},
	}
	for _, test := range tests {
		t.Run(fmt.Sprintf("%s-%s", test.machineType, test.arch), func(t *testing.T) {
			_, cluster := buildMinimalCluster()
			cloud, err := BuildCloud(cluster)
			if err != nil {
				t.Fatalf("error from BuildCloud: %v", err)
			}

			arch, err := MachineArchitecture(cloud, test.machineType)
			if err != test.err {
				t.Errorf("actual error %q differs from expected error %q", err, test.err)
				return
			}

			if arch != test.arch {
				t.Errorf("actual architecture %q differs from expected architecture %q", arch, test.arch)
				return
			}
		})
	}
}
<|endoftext|>


<commit_before>package hoverfly

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	// static assets
	_ "github.com/SpectoLabs/hoverfly/statik"
	"github.com/rakyll/statik/fs"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/go-zoo/bone"
	"github.com/gorilla/websocket"
	"github.com/meatballhat/negroni-logrus"
)

// recordedRequests struct encapsulates payload data
type recordedRequests struct {
	Data []Payload `json:"data"`
}

type recordsCount struct {
	Count int `json:"count"`
}

type statsResponse struct {
	Stats        Stats `json:"stats"`
	RecordsCount int   `json:"recordsCount"`
}

type stateRequest struct {
	Mode        string `json:"mode"`
	Destination string `json:"destination"`
}

type messageResponse struct {
	Message string `json:"message"`
}

// StartAdminInterface - starts admin interface web server
func (d *DBClient) StartAdminInterface() {
	go func() {
		// starting admin interface
		mux := getBoneRouter(*d)
		n := negroni.Classic()

		logLevel := log.ErrorLevel

		if d.Cfg.Verbose {
			logLevel = log.DebugLevel
		}

		n.Use(negronilogrus.NewCustomMiddleware(logLevel, &log.JSONFormatter{}, "admin"))
		n.UseHandler(mux)

		// admin interface starting message
		log.WithFields(log.Fields{
			"AdminPort": d.Cfg.AdminPort,
		}).Info("Admin interface is starting...")

		n.Run(fmt.Sprintf(":%s", d.Cfg.AdminPort))
	}()
}

// getBoneRouter returns mux for admin interface
func getBoneRouter(d DBClient) *bone.Mux {
	mux := bone.New()

	mux.Get("/records", http.HandlerFunc(d.AllRecordsHandler))
	mux.Delete("/records", http.HandlerFunc(d.DeleteAllRecordsHandler))
	mux.Post("/records", http.HandlerFunc(d.ImportRecordsHandler))

	mux.Get("/count", http.HandlerFunc(d.RecordsCount))
	mux.Get("/stats", http.HandlerFunc(d.StatsHandler))
	mux.Get("/statsws", http.HandlerFunc(d.StatsWSHandler))

	mux.Get("/state", http.HandlerFunc(d.CurrentStateHandler))
	mux.Post("/state", http.HandlerFunc(d.StateHandler))

	if d.Cfg.Development {
		// since hoverfly is not started from cmd/hoverfly/hoverfly
		// we have to target to that directory
		log.Warn("Hoverfly is serving files from /static/dist instead of statik binary!")
		mux.Handle("/*", http.FileServer(http.Dir("../../static/dist")))
	} else {
		// preparing static assets for embedded admin
		statikFS, err := fs.New()

		if err != nil {
			log.WithFields(log.Fields{
				"Error": err.Error(),
			}).Error("Failed to load statikFS, admin UI might not work :(")
		}

		mux.Handle("/*", http.FileServer(statikFS))
	}

	return mux
}

// AllRecordsHandler returns JSON content type http response
func (d *DBClient) AllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	records, err := d.Cache.GetAllRequests()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordedRequests
		response.Data = records
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// RecordsCount returns number of captured requests as a JSON payload
func (d *DBClient) RecordsCount(w http.ResponseWriter, req *http.Request) {
	count, err := d.Cache.RecordsCount()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordsCount
		response.Count = count
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// StatsHandler - returns current stats about Hoverfly (request counts, record count)
func (d *DBClient) StatsHandler(w http.ResponseWriter, req *http.Request) {
	stats := d.Counter.Flush()

	count, err := d.Cache.RecordsCount()

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	var sr statsResponse
	sr.Stats = stats
	sr.RecordsCount = count

	w.Header().Set("Content-Type", "application/json")

	b, err := json.Marshal(sr)

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Write(b)
		return
	}

}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// StatsWSHandler - returns current stats about Hoverfly (request counts, record count) through the websocket
func (d *DBClient) StatsWSHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			return
		}
		log.WithFields(log.Fields{
			"message": string(p),
		}).Info("Got message...")

		for _ = range time.Tick(1 * time.Second) {

			count, err := d.Cache.RecordsCount()

			if err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Error("got error while trying to get records count")
				return
			}

			stats := d.Counter.Flush()

			var sr statsResponse
			sr.Stats = stats
			sr.RecordsCount = count

			b, err := json.Marshal(sr)

			if err = conn.WriteMessage(messageType, b); err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Debug("Got error when writing message...")
				return
			}
		}

	}

}

// ImportRecordsHandler - accepts JSON payload and saves it to cache
func (d *DBClient) ImportRecordsHandler(w http.ResponseWriter, req *http.Request) {

	var requests recordedRequests

	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	var response messageResponse

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		response.Message = "Bad request. Nothing to import!"
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &requests)

	if err != nil {
		w.WriteHeader(422) // can't process this entity
		return
	}

	err = d.ImportPayloads(requests.Data)

	if err != nil {
		response.Message = err.Error()
		w.WriteHeader(400)
	} else {
		response.Message = fmt.Sprintf("%d payloads import complete.", len(requests.Data))
	}

	b, err := json.Marshal(response)
	w.Write(b)

}

// DeleteAllRecordsHandler - deletes all captured requests
func (d *DBClient) DeleteAllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	err := d.Cache.DeleteData()

	var en Entry
	en.ActionType = ActionTypeWipeDB
	en.Message = "wipe"
	en.Time = time.Now()

	if err := d.Hooks.Fire(ActionTypeWipeDB, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeWipeDB,
		}).Error("failed to fire hook")
	}

	w.Header().Set("Content-Type", "application/json")

	var response messageResponse
	if err != nil {
		if err.Error() == "bucket not found" {
			response.Message = fmt.Sprintf("No records found")
			w.WriteHeader(200)
		} else {
			response.Message = fmt.Sprintf("Something went wrong: %s", err.Error())
			w.WriteHeader(500)
		}
	} else {
		response.Message = "Proxy cache deleted successfuly"
		w.WriteHeader(200)
	}
	b, err := json.Marshal(response)

	w.Write(b)
	return
}

// CurrentStateHandler returns current state
func (d *DBClient) CurrentStateHandler(w http.ResponseWriter, req *http.Request) {
	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination

	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)
}

// StateHandler handles current proxy state
func (d *DBClient) StateHandler(w http.ResponseWriter, r *http.Request) {
	var sr stateRequest

	// this is mainly for testing, since when you create
	if r.Body == nil {
		r.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("")))
	}

	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &sr)

	if err != nil {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(400) // can't process this entity
		return
	}

	availableModes := map[string]bool{
		"virtualize": true,
		"capture":    true,
		"modify":     true,
		"synthesize": true,
	}

	if !availableModes[sr.Mode] {
		log.WithFields(log.Fields{
			"suppliedMode": sr.Mode,
		}).Error("Wrong mode found, can't change state")
		http.Error(w, "Bad mode supplied, available modes: virtualize, capture, modify, synthesize.", 400)
		return
	}

	log.WithFields(log.Fields{
		"newState": sr.Mode,
		"body":     string(body),
	}).Info("Handling state change request!")

	// setting new state
	d.Cfg.SetMode(sr.Mode)

	var en Entry
	en.ActionType = ActionTypeConfigurationChanged
	en.Message = "changed"
	en.Time = time.Now()
	en.Data = []byte("sr.Mode")

	if err := d.Hooks.Fire(ActionTypeConfigurationChanged, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeConfigurationChanged,
		}).Error("failed to fire hook")
	}

	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination
	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)

}
<commit_msg>getting token and logging out<commit_after>package hoverfly

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	// static assets
	_ "github.com/SpectoLabs/hoverfly/statik"
	"github.com/rakyll/statik/fs"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/go-zoo/bone"
	"github.com/gorilla/websocket"
	"github.com/meatballhat/negroni-logrus"

	// auth
	"github.com/SpectoLabs/hoverfly/authentication"
	"github.com/SpectoLabs/hoverfly/authentication/controllers"
)

// recordedRequests struct encapsulates payload data
type recordedRequests struct {
	Data []Payload `json:"data"`
}

type recordsCount struct {
	Count int `json:"count"`
}

type statsResponse struct {
	Stats        Stats `json:"stats"`
	RecordsCount int   `json:"recordsCount"`
}

type stateRequest struct {
	Mode        string `json:"mode"`
	Destination string `json:"destination"`
}

type messageResponse struct {
	Message string `json:"message"`
}

// StartAdminInterface - starts admin interface web server
func (d *DBClient) StartAdminInterface() {
	go func() {
		// starting admin interface
		mux := getBoneRouter(*d)
		n := negroni.Classic()

		logLevel := log.ErrorLevel

		if d.Cfg.Verbose {
			logLevel = log.DebugLevel
		}

		n.Use(negronilogrus.NewCustomMiddleware(logLevel, &log.JSONFormatter{}, "admin"))
		n.UseHandler(mux)

		// admin interface starting message
		log.WithFields(log.Fields{
			"AdminPort": d.Cfg.AdminPort,
		}).Info("Admin interface is starting...")

		n.Run(fmt.Sprintf(":%s", d.Cfg.AdminPort))
	}()
}

// getBoneRouter returns mux for admin interface
func getBoneRouter(d DBClient) *bone.Mux {
	mux := bone.New()

	// getting auth controllers and middleware
	ac := controllers.GetNewAuthenticationController(d.AB)
	am := authentication.GetNewAuthenticationMiddleware(d.AB)

	mux.Post("/token-auth", http.HandlerFunc(ac.Login))
	mux.Get("/logout", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(ac.Logout),
	))

	mux.Get("/records", http.HandlerFunc(d.AllRecordsHandler))
	mux.Delete("/records", http.HandlerFunc(d.DeleteAllRecordsHandler))
	mux.Post("/records", http.HandlerFunc(d.ImportRecordsHandler))

	mux.Get("/count", http.HandlerFunc(d.RecordsCount))
	mux.Get("/stats", http.HandlerFunc(d.StatsHandler))
	mux.Get("/statsws", http.HandlerFunc(d.StatsWSHandler))

	mux.Get("/state", http.HandlerFunc(d.CurrentStateHandler))
	mux.Post("/state", http.HandlerFunc(d.StateHandler))

	if d.Cfg.Development {
		// since hoverfly is not started from cmd/hoverfly/hoverfly
		// we have to target to that directory
		log.Warn("Hoverfly is serving files from /static/dist instead of statik binary!")
		mux.Handle("/*", http.FileServer(http.Dir("../../static/dist")))
	} else {
		// preparing static assets for embedded admin
		statikFS, err := fs.New()

		if err != nil {
			log.WithFields(log.Fields{
				"Error": err.Error(),
			}).Error("Failed to load statikFS, admin UI might not work :(")
		}

		mux.Handle("/*", http.FileServer(statikFS))
	}

	return mux
}

// AllRecordsHandler returns JSON content type http response
func (d *DBClient) AllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	records, err := d.Cache.GetAllRequests()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordedRequests
		response.Data = records
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// RecordsCount returns number of captured requests as a JSON payload
func (d *DBClient) RecordsCount(w http.ResponseWriter, req *http.Request) {
	count, err := d.Cache.RecordsCount()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordsCount
		response.Count = count
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// StatsHandler - returns current stats about Hoverfly (request counts, record count)
func (d *DBClient) StatsHandler(w http.ResponseWriter, req *http.Request) {
	stats := d.Counter.Flush()

	count, err := d.Cache.RecordsCount()

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	var sr statsResponse
	sr.Stats = stats
	sr.RecordsCount = count

	w.Header().Set("Content-Type", "application/json")

	b, err := json.Marshal(sr)

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Write(b)
		return
	}

}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// StatsWSHandler - returns current stats about Hoverfly (request counts, record count) through the websocket
func (d *DBClient) StatsWSHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			return
		}
		log.WithFields(log.Fields{
			"message": string(p),
		}).Info("Got message...")

		for _ = range time.Tick(1 * time.Second) {

			count, err := d.Cache.RecordsCount()

			if err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Error("got error while trying to get records count")
				return
			}

			stats := d.Counter.Flush()

			var sr statsResponse
			sr.Stats = stats
			sr.RecordsCount = count

			b, err := json.Marshal(sr)

			if err = conn.WriteMessage(messageType, b); err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Debug("Got error when writing message...")
				return
			}
		}

	}

}

// ImportRecordsHandler - accepts JSON payload and saves it to cache
func (d *DBClient) ImportRecordsHandler(w http.ResponseWriter, req *http.Request) {

	var requests recordedRequests

	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	var response messageResponse

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		response.Message = "Bad request. Nothing to import!"
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &requests)

	if err != nil {
		w.WriteHeader(422) // can't process this entity
		return
	}

	err = d.ImportPayloads(requests.Data)

	if err != nil {
		response.Message = err.Error()
		w.WriteHeader(400)
	} else {
		response.Message = fmt.Sprintf("%d payloads import complete.", len(requests.Data))
	}

	b, err := json.Marshal(response)
	w.Write(b)

}

// DeleteAllRecordsHandler - deletes all captured requests
func (d *DBClient) DeleteAllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	err := d.Cache.DeleteData()

	var en Entry
	en.ActionType = ActionTypeWipeDB
	en.Message = "wipe"
	en.Time = time.Now()

	if err := d.Hooks.Fire(ActionTypeWipeDB, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeWipeDB,
		}).Error("failed to fire hook")
	}

	w.Header().Set("Content-Type", "application/json")

	var response messageResponse
	if err != nil {
		if err.Error() == "bucket not found" {
			response.Message = fmt.Sprintf("No records found")
			w.WriteHeader(200)
		} else {
			response.Message = fmt.Sprintf("Something went wrong: %s", err.Error())
			w.WriteHeader(500)
		}
	} else {
		response.Message = "Proxy cache deleted successfuly"
		w.WriteHeader(200)
	}
	b, err := json.Marshal(response)

	w.Write(b)
	return
}

// CurrentStateHandler returns current state
func (d *DBClient) CurrentStateHandler(w http.ResponseWriter, req *http.Request) {
	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination

	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)
}

// StateHandler handles current proxy state
func (d *DBClient) StateHandler(w http.ResponseWriter, r *http.Request) {
	var sr stateRequest

	// this is mainly for testing, since when you create
	if r.Body == nil {
		r.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("")))
	}

	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &sr)

	if err != nil {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(400) // can't process this entity
		return
	}

	availableModes := map[string]bool{
		"virtualize": true,
		"capture":    true,
		"modify":     true,
		"synthesize": true,
	}

	if !availableModes[sr.Mode] {
		log.WithFields(log.Fields{
			"suppliedMode": sr.Mode,
		}).Error("Wrong mode found, can't change state")
		http.Error(w, "Bad mode supplied, available modes: virtualize, capture, modify, synthesize.", 400)
		return
	}

	log.WithFields(log.Fields{
		"newState": sr.Mode,
		"body":     string(body),
	}).Info("Handling state change request!")

	// setting new state
	d.Cfg.SetMode(sr.Mode)

	var en Entry
	en.ActionType = ActionTypeConfigurationChanged
	en.Message = "changed"
	en.Time = time.Now()
	en.Data = []byte("sr.Mode")

	if err := d.Hooks.Fire(ActionTypeConfigurationChanged, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeConfigurationChanged,
		}).Error("failed to fire hook")
	}

	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination
	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)

}
<|endoftext|>


<commit_before>package distro

import (
	"bytes"
	"compress/gzip"
	"compress/zlib"
	"encoding/base64"
)

type gzipTransformer struct {
	writer *gzip.Writer
}

func (gzt gzipTransformer) Transform(data []byte) []byte {
	var buf bytes.Buffer

	if gzt.writer == nil {
		gzt.writer = gzip.NewWriter(&buf)
	} else {
		gzt.writer.Reset(&buf)
	}

	gzt.writer.Write(data)
	gzt.writer.Close()

	return bytesBufferToBase64(buf)
}

type zlibTransformer struct {
	writer *zlib.Writer
}

func (zlt zlibTransformer) Transform(data []byte) []byte {
	var buf bytes.Buffer

	if zlt.writer == nil {
		zlt.writer = zlib.NewWriter(&buf)
	} else {
		zlt.writer.Reset(&buf)
	}

	zlt.writer.Write(data)
	zlt.writer.Close()

	return bytesBufferToBase64(buf)
}

func bytesBufferToBase64(buf bytes.Buffer) []byte {
	dst := make([]byte, 0, 0)

	base64.StdEncoding.Encode(dst, buf.Bytes())
	return dst
}
<commit_msg>Fix base64 encoding<commit_after>package distro

import (
	"bytes"
	"compress/gzip"
	"compress/zlib"
	"encoding/base64"
)

type gzipTransformer struct {
	writer *gzip.Writer
}

func (gzt gzipTransformer) Transform(data []byte) []byte {
	var buf bytes.Buffer

	if gzt.writer == nil {
		gzt.writer = gzip.NewWriter(&buf)
	} else {
		gzt.writer.Reset(&buf)
	}

	gzt.writer.Write(data)
	gzt.writer.Close()

	return bytesBufferToBase64(buf)
}

type zlibTransformer struct {
	writer *zlib.Writer
}

func (zlt zlibTransformer) Transform(data []byte) []byte {
	var buf bytes.Buffer

	if zlt.writer == nil {
		zlt.writer = zlib.NewWriter(&buf)
	} else {
		zlt.writer.Reset(&buf)
	}

	zlt.writer.Write(data)
	zlt.writer.Close()

	return bytesBufferToBase64(buf)
}

func bytesBufferToBase64(buf bytes.Buffer) []byte {
	dst := make([]byte, base64.StdEncoding.EncodedLen(buf.Len()))
	base64.StdEncoding.Encode(dst, buf.Bytes())
	return dst
}
<|endoftext|>


<commit_before>package imguploader

import (
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/util"
)

type S3Uploader struct {
	region    string
	bucket    string
	acl       string
	secretKey string
	accessKey string
	log       log.Logger
}

func NewS3Uploader(region, bucket, acl, accessKey, secretKey string) *S3Uploader {
	return &S3Uploader{
		region:    region,
		bucket:    bucket,
		acl:       acl,
		accessKey: accessKey,
		secretKey: secretKey,
		log:       log.New("s3uploader"),
	}
}

func (u *S3Uploader) Upload(imageDiskPath string) (string, error) {
	sess, err := session.NewSession()
	if err != nil {
		return "", err
	}
	creds := credentials.NewChainCredentials(
		[]credentials.Provider{
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     u.accessKey,
				SecretAccessKey: u.secretKey,
			}},
			&credentials.EnvProvider{},
			&ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(sess), ExpiryWindow: 5 * time.Minute},
		})
	cfg := &aws.Config{
		Region:      aws.String(u.region),
		Credentials: creds,
	}

	key := util.GetRandomString(20) + ".png"
	log.Debug("Uploading image to s3", "bucket = ", u.bucket, ", key = ", key)

	file, err := os.Open(imageDiskPath)
	if err != nil {
		return "", err
	}

	sess, err = session.NewSession(cfg)
	if err != nil {
		return "", err
	}
	svc := s3.New(sess, cfg)
	params := &s3.PutObjectInput{
		Bucket:      aws.String(u.bucket),
		Key:         aws.String(key),
		ACL:         aws.String(u.acl),
		Body:        file,
		ContentType: aws.String("image/png"),
	}
	_, err = svc.PutObject(params)
	if err != nil {
		return "", err
	}

	return "https://" + u.bucket + ".s3.amazonaws.com/" + key, nil
}
<commit_msg>Specify region for s3 (#8251)<commit_after>package imguploader

import (
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/util"
)

type S3Uploader struct {
	region    string
	bucket    string
	acl       string
	secretKey string
	accessKey string
	log       log.Logger
}

func NewS3Uploader(region, bucket, acl, accessKey, secretKey string) *S3Uploader {
	return &S3Uploader{
		region:    region,
		bucket:    bucket,
		acl:       acl,
		accessKey: accessKey,
		secretKey: secretKey,
		log:       log.New("s3uploader"),
	}
}

func (u *S3Uploader) Upload(imageDiskPath string) (string, error) {
	sess, err := session.NewSession()
	if err != nil {
		return "", err
	}
	creds := credentials.NewChainCredentials(
		[]credentials.Provider{
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     u.accessKey,
				SecretAccessKey: u.secretKey,
			}},
			&credentials.EnvProvider{},
			&ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(sess), ExpiryWindow: 5 * time.Minute},
		})
	cfg := &aws.Config{
		Region:      aws.String(u.region),
		Credentials: creds,
	}

	key := util.GetRandomString(20) + ".png"
	log.Debug("Uploading image to s3", "bucket = ", u.bucket, ", key = ", key)

	file, err := os.Open(imageDiskPath)
	if err != nil {
		return "", err
	}

	sess, err = session.NewSession(cfg)
	if err != nil {
		return "", err
	}
	svc := s3.New(sess, cfg)
	params := &s3.PutObjectInput{
		Bucket:      aws.String(u.bucket),
		Key:         aws.String(key),
		ACL:         aws.String(u.acl),
		Body:        file,
		ContentType: aws.String("image/png"),
	}
	_, err = svc.PutObject(params)
	if err != nil {
		return "", err
	}

	return "https://" + u.bucket + ".s3-" + u.region + ".amazonaws.com/" + key, nil
}
<|endoftext|>


<commit_before>package gorelic

import (
	"errors"
	"fmt"
	"log"
	"net/http"

	metrics "github.com/yvasiyarov/go-metrics"
	"github.com/yvasiyarov/newrelic_platform_go"
)

const (
	// DefaultNewRelicPollInterval - how often we will report metrics to NewRelic.
	// Recommended values is 60 seconds
	DefaultNewRelicPollInterval = 60

	// DefaultGcPollIntervalInSeconds - how often we will get garbage collector run statistic
	// Default value is - every 10 seconds
	// During GC stat pooling - mheap will be locked, so be carefull changing this value
	DefaultGcPollIntervalInSeconds = 10

	// DefaultMemoryAllocatorPollIntervalInSeconds - how often we will get memory allocator statistic.
	// Default value is - every 60 seconds
	// During this process stoptheword() is called, so be carefull changing this value
	DefaultMemoryAllocatorPollIntervalInSeconds = 60

	//DefaultAgentGuid is plugin ID in NewRelic.
	//You should not change it unless you want to create your own plugin.
	DefaultAgentGuid = "com.github.yvasiyarov.GoRelic"

	//CurrentAgentVersion is plugin version
	CurrentAgentVersion = "0.0.6"

	//DefaultAgentName in NewRelic GUI. You can change it.
	DefaultAgentName = "Go daemon"
)

//Agent - is NewRelic agent implementation.
//Agent start separate go routine which will report data to NewRelic
type Agent struct {
	NewrelicName                string
	NewrelicLicense             string
	NewrelicPollInterval        int
	Verbose                     bool
	CollectGcStat               bool
	CollectMemoryStat           bool
	CollectHTTPStat             bool
	CollectHTTPStatuses         bool
	GCPollInterval              int
	MemoryAllocatorPollInterval int
	AgentGUID                   string
	AgentVersion                string
	plugin                      *newrelic_platform_go.NewrelicPlugin
	HTTPTimer                   metrics.Timer
	HTTPStatusCounters          map[int]metrics.Counter
	Tracer                      *Tracer
	CustomMetrics               []newrelic_platform_go.IMetrica

	// All HTTP requests will be done using this client. Change it if you need
	// to use a proxy.
	Client http.Client
}

// NewAgent builds new Agent objects.
func NewAgent() *Agent {
	agent := &Agent{
		NewrelicName:                DefaultAgentName,
		NewrelicPollInterval:        DefaultNewRelicPollInterval,
		Verbose:                     false,
		CollectGcStat:               true,
		CollectMemoryStat:           true,
		GCPollInterval:              DefaultGcPollIntervalInSeconds,
		MemoryAllocatorPollInterval: DefaultMemoryAllocatorPollIntervalInSeconds,
		AgentGUID:                   DefaultAgentGuid,
		AgentVersion:                CurrentAgentVersion,
		Tracer:                      nil,
		CustomMetrics:               make([]newrelic_platform_go.IMetrica, 0),
	}
	return agent
}

// our custom component
type resettableComponent struct {
	newrelic_platform_go.IComponent
	counters map[int]metrics.Counter
}

// newrelic_platform_go.IComponent interface implementation
func (c resettableComponent) ClearSentData() {
	c.IComponent.ClearSentData()
	for _, counter := range c.counters {
		counter.Clear()
	}
}

//WrapHTTPHandlerFunc  instrument HTTP handler functions to collect HTTP metrics
func (agent *Agent) WrapHTTPHandlerFunc(h tHTTPHandlerFunc) tHTTPHandlerFunc {
	agent.CollectHTTPStat = true
	agent.initTimer()
	return func(w http.ResponseWriter, req *http.Request) {
		proxy := newHTTPHandlerFunc(h)
		proxy.timer = agent.HTTPTimer
		proxy.ServeHTTP(w, req)
	}
}

//WrapHTTPHandler  instrument HTTP handler object to collect HTTP metrics
func (agent *Agent) WrapHTTPHandler(h http.Handler) http.Handler {
	agent.CollectHTTPStat = true
	agent.initTimer()

	proxy := newHTTPHandler(h)
	proxy.timer = agent.HTTPTimer
	return proxy
}

//AddCustomMetric adds metric to be collected periodically with NewrelicPollInterval interval
func (agent *Agent) AddCustomMetric(metric newrelic_platform_go.IMetrica) {
	agent.CustomMetrics = append(agent.CustomMetrics, metric)
}

//Run initialize Agent instance and start harvest go routine
func (agent *Agent) Run() error {
	if agent.NewrelicLicense == "" {
		return errors.New("please, pass a valid newrelic license key")
	}

	agent.plugin = newrelic_platform_go.NewNewrelicPlugin(agent.AgentVersion, agent.NewrelicLicense, agent.NewrelicPollInterval)
	agent.plugin.Client = agent.Client
	component := newrelic_platform_go.NewPluginComponent(agent.NewrelicName, agent.AgentGUID)
	agent.plugin.AddComponent(component)

	addRuntimeMericsToComponent(component)
	agent.Tracer = newTracer(component)

	// Check agent flags and add relevant metrics.
	if agent.CollectGcStat {
		addGCMericsToComponent(component, agent.GCPollInterval)
		agent.debug(fmt.Sprintf("Init GC metrics collection. Poll interval %d seconds.", agent.GCPollInterval))
	}

	if agent.CollectMemoryStat {
		addMemoryMericsToComponent(component, agent.MemoryAllocatorPollInterval)
		agent.debug(fmt.Sprintf("Init memory allocator metrics collection. Poll interval %d seconds.", agent.MemoryAllocatorPollInterval))
	}

	if agent.CollectHTTPStat {
		agent.initTimer()
		addHTTPMericsToComponent(component, agent.HTTPTimer)
		agent.debug(fmt.Sprintf("Init HTTP metrics collection."))
	}

	for _, metric := range agent.CustomMetrics {
		component.AddMetrica(metric)
		agent.debug(fmt.Sprintf("Init %s metric collection.", metric.GetName()))
	}

	if agent.CollectHTTPStatuses {
		agent.initStatusCounters()
		component := &resettableComponent{component, agent.HTTPStatusCounters}
		addHTTPStatusMetricsToComponent(component, agent.HTTPStatusCounters)
		agent.debug(fmt.Sprintf("Init HTTP status metrics collection."))
	}

	// Init newrelic reporting plugin.
	agent.plugin = newrelic_platform_go.NewNewrelicPlugin(agent.AgentVersion, agent.NewrelicLicense, agent.NewrelicPollInterval)
	agent.plugin.Verbose = agent.Verbose

	// Add our metrics component to the plugin.
	agent.plugin.AddComponent(component)

	// Start reporting!
	go agent.plugin.Run()
	return nil
}

//Initialize global metrics.Timer object, used to collect HTTP metrics
func (agent *Agent) initTimer() {
	if agent.HTTPTimer == nil {
		agent.HTTPTimer = metrics.NewTimer()
	}
}

//Initialize metrics.Counters objects, used to collect HTTP statuses
func (agent *Agent) initStatusCounters() {
	httpStatuses := []int{
		http.StatusContinue, http.StatusSwitchingProtocols,

		http.StatusOK, http.StatusCreated, http.StatusAccepted, http.StatusNonAuthoritativeInfo,
		http.StatusNoContent, http.StatusResetContent, http.StatusPartialContent,

		http.StatusMultipleChoices, http.StatusMovedPermanently, http.StatusFound, http.StatusSeeOther,
		http.StatusNotModified, http.StatusUseProxy, http.StatusTemporaryRedirect,

		http.StatusBadRequest, http.StatusUnauthorized, http.StatusPaymentRequired, http.StatusForbidden,
		http.StatusNotFound, http.StatusMethodNotAllowed, http.StatusNotAcceptable, http.StatusProxyAuthRequired,
		http.StatusRequestTimeout, http.StatusConflict, http.StatusGone, http.StatusLengthRequired,
		http.StatusPreconditionFailed, http.StatusRequestEntityTooLarge, http.StatusRequestURITooLong, http.StatusUnsupportedMediaType,
		http.StatusRequestedRangeNotSatisfiable, http.StatusExpectationFailed, http.StatusTeapot,

		http.StatusInternalServerError, http.StatusNotImplemented, http.StatusBadGateway,
		http.StatusServiceUnavailable, http.StatusGatewayTimeout, http.StatusHTTPVersionNotSupported,
	}

	agent.HTTPStatusCounters = make(map[int]metrics.Counter, len(httpStatuses))
	for _, statusCode := range httpStatuses {
		agent.HTTPStatusCounters[statusCode] = metrics.NewCounter()
	}
}

//Print debug messages
func (agent *Agent) debug(msg string) {
	if agent.Verbose {
		log.Println(msg)
	}
}
<commit_msg>Fixed merge problems.<commit_after>package gorelic

import (
	"errors"
	"fmt"
	"log"
	"net/http"

	metrics "github.com/yvasiyarov/go-metrics"
	"github.com/yvasiyarov/newrelic_platform_go"
)

const (
	// DefaultNewRelicPollInterval - how often we will report metrics to NewRelic.
	// Recommended values is 60 seconds
	DefaultNewRelicPollInterval = 60

	// DefaultGcPollIntervalInSeconds - how often we will get garbage collector run statistic
	// Default value is - every 10 seconds
	// During GC stat pooling - mheap will be locked, so be carefull changing this value
	DefaultGcPollIntervalInSeconds = 10

	// DefaultMemoryAllocatorPollIntervalInSeconds - how often we will get memory allocator statistic.
	// Default value is - every 60 seconds
	// During this process stoptheword() is called, so be carefull changing this value
	DefaultMemoryAllocatorPollIntervalInSeconds = 60

	//DefaultAgentGuid is plugin ID in NewRelic.
	//You should not change it unless you want to create your own plugin.
	DefaultAgentGuid = "com.github.yvasiyarov.GoRelic"

	//CurrentAgentVersion is plugin version
	CurrentAgentVersion = "0.0.6"

	//DefaultAgentName in NewRelic GUI. You can change it.
	DefaultAgentName = "Go daemon"
)

//Agent - is NewRelic agent implementation.
//Agent start separate go routine which will report data to NewRelic
type Agent struct {
	NewrelicName                string
	NewrelicLicense             string
	NewrelicPollInterval        int
	Verbose                     bool
	CollectGcStat               bool
	CollectMemoryStat           bool
	CollectHTTPStat             bool
	CollectHTTPStatuses         bool
	GCPollInterval              int
	MemoryAllocatorPollInterval int
	AgentGUID                   string
	AgentVersion                string
	plugin                      *newrelic_platform_go.NewrelicPlugin
	HTTPTimer                   metrics.Timer
	HTTPStatusCounters          map[int]metrics.Counter
	Tracer                      *Tracer
	CustomMetrics               []newrelic_platform_go.IMetrica

	// All HTTP requests will be done using this client. Change it if you need
	// to use a proxy.
	Client http.Client
}

// NewAgent builds new Agent objects.
func NewAgent() *Agent {
	agent := &Agent{
		NewrelicName:                DefaultAgentName,
		NewrelicPollInterval:        DefaultNewRelicPollInterval,
		Verbose:                     false,
		CollectGcStat:               true,
		CollectMemoryStat:           true,
		GCPollInterval:              DefaultGcPollIntervalInSeconds,
		MemoryAllocatorPollInterval: DefaultMemoryAllocatorPollIntervalInSeconds,
		AgentGUID:                   DefaultAgentGuid,
		AgentVersion:                CurrentAgentVersion,
		Tracer:                      nil,
		CustomMetrics:               make([]newrelic_platform_go.IMetrica, 0),
	}
	return agent
}

// our custom component
type resettableComponent struct {
	newrelic_platform_go.IComponent
	counters map[int]metrics.Counter
}

// newrelic_platform_go.IComponent interface implementation
func (c resettableComponent) ClearSentData() {
	c.IComponent.ClearSentData()
	for _, counter := range c.counters {
		counter.Clear()
	}
}

//WrapHTTPHandlerFunc  instrument HTTP handler functions to collect HTTP metrics
func (agent *Agent) WrapHTTPHandlerFunc(h tHTTPHandlerFunc) tHTTPHandlerFunc {
	agent.CollectHTTPStat = true
	agent.initTimer()
	return func(w http.ResponseWriter, req *http.Request) {
		proxy := newHTTPHandlerFunc(h)
		proxy.timer = agent.HTTPTimer
		proxy.ServeHTTP(w, req)
	}
}

//WrapHTTPHandler  instrument HTTP handler object to collect HTTP metrics
func (agent *Agent) WrapHTTPHandler(h http.Handler) http.Handler {
	agent.CollectHTTPStat = true
	agent.initTimer()

	proxy := newHTTPHandler(h)
	proxy.timer = agent.HTTPTimer
	return proxy
}

//AddCustomMetric adds metric to be collected periodically with NewrelicPollInterval interval
func (agent *Agent) AddCustomMetric(metric newrelic_platform_go.IMetrica) {
	agent.CustomMetrics = append(agent.CustomMetrics, metric)
}

//Run initialize Agent instance and start harvest go routine
func (agent *Agent) Run() error {
	if agent.NewrelicLicense == "" {
		return errors.New("please, pass a valid newrelic license key")
	}

	agent.plugin = newrelic_platform_go.NewNewrelicPlugin(agent.AgentVersion, agent.NewrelicLicense, agent.NewrelicPollInterval)
	agent.plugin.Client = agent.Client

	var component newrelic_platform_go.IComponent
	component = newrelic_platform_go.NewPluginComponent(agent.NewrelicName, agent.AgentGUID)

	// Add default metrics and tracer.
	addRuntimeMericsToComponent(component)
	agent.Tracer = newTracer(component)

	// Check agent flags and add relevant metrics.
	if agent.CollectGcStat {
		addGCMericsToComponent(component, agent.GCPollInterval)
		agent.debug(fmt.Sprintf("Init GC metrics collection. Poll interval %d seconds.", agent.GCPollInterval))
	}

	if agent.CollectMemoryStat {
		addMemoryMericsToComponent(component, agent.MemoryAllocatorPollInterval)
		agent.debug(fmt.Sprintf("Init memory allocator metrics collection. Poll interval %d seconds.", agent.MemoryAllocatorPollInterval))
	}

	if agent.CollectHTTPStat {
		agent.initTimer()
		addHTTPMericsToComponent(component, agent.HTTPTimer)
		agent.debug(fmt.Sprintf("Init HTTP metrics collection."))
	}

	for _, metric := range agent.CustomMetrics {
		component.AddMetrica(metric)
		agent.debug(fmt.Sprintf("Init %s metric collection.", metric.GetName()))
	}

	if agent.CollectHTTPStatuses {
		agent.initStatusCounters()
		component = &resettableComponent{component, agent.HTTPStatusCounters}
		addHTTPStatusMetricsToComponent(component, agent.HTTPStatusCounters)
		agent.debug(fmt.Sprintf("Init HTTP status metrics collection."))
	}

	// Init newrelic reporting plugin.
	agent.plugin = newrelic_platform_go.NewNewrelicPlugin(agent.AgentVersion, agent.NewrelicLicense, agent.NewrelicPollInterval)
	agent.plugin.Verbose = agent.Verbose

	// Add our metrics component to the plugin.
	agent.plugin.AddComponent(component)

	// Start reporting!
	go agent.plugin.Run()
	return nil
}

//Initialize global metrics.Timer object, used to collect HTTP metrics
func (agent *Agent) initTimer() {
	if agent.HTTPTimer == nil {
		agent.HTTPTimer = metrics.NewTimer()
	}
}

//Initialize metrics.Counters objects, used to collect HTTP statuses
func (agent *Agent) initStatusCounters() {
	httpStatuses := []int{
		http.StatusContinue, http.StatusSwitchingProtocols,

		http.StatusOK, http.StatusCreated, http.StatusAccepted, http.StatusNonAuthoritativeInfo,
		http.StatusNoContent, http.StatusResetContent, http.StatusPartialContent,

		http.StatusMultipleChoices, http.StatusMovedPermanently, http.StatusFound, http.StatusSeeOther,
		http.StatusNotModified, http.StatusUseProxy, http.StatusTemporaryRedirect,

		http.StatusBadRequest, http.StatusUnauthorized, http.StatusPaymentRequired, http.StatusForbidden,
		http.StatusNotFound, http.StatusMethodNotAllowed, http.StatusNotAcceptable, http.StatusProxyAuthRequired,
		http.StatusRequestTimeout, http.StatusConflict, http.StatusGone, http.StatusLengthRequired,
		http.StatusPreconditionFailed, http.StatusRequestEntityTooLarge, http.StatusRequestURITooLong, http.StatusUnsupportedMediaType,
		http.StatusRequestedRangeNotSatisfiable, http.StatusExpectationFailed, http.StatusTeapot,

		http.StatusInternalServerError, http.StatusNotImplemented, http.StatusBadGateway,
		http.StatusServiceUnavailable, http.StatusGatewayTimeout, http.StatusHTTPVersionNotSupported,
	}

	agent.HTTPStatusCounters = make(map[int]metrics.Counter, len(httpStatuses))
	for _, statusCode := range httpStatuses {
		agent.HTTPStatusCounters[statusCode] = metrics.NewCounter()
	}
}

//Print debug messages
func (agent *Agent) debug(msg string) {
	if agent.Verbose {
		log.Println(msg)
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2018 Johannes Kohnen <wjkohnen@users.noreply.github.com>
 *
 * 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 airac

import (
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"
)

const (
	format        = "2006-01-02"
	cycleDuration = 24192e11 // 4 weeks
)

var (
	epoch = time.Date(1901, time.January, 10, 0, 0, 0, 0, time.UTC)
)

// AIRAC represents an Aeronautical Information Regulation And Control (AIRAC) cycle.
type AIRAC uint16

type Airac = AIRAC

// Effective returns the effective date of this AIRAC cycle.
func (a AIRAC) Effective() time.Time {
	return epoch.Add(time.Duration(a) * cycleDuration)
}

// Year returns the year for this AIRAC cycle's identifier.
func (a AIRAC) Year() int {
	return a.Effective().Year()
}

// Ordinal returns the ordinal for this AIRAC cycle's identifier.
func (a AIRAC) Ordinal() int {
	return (a.Effective().YearDay()-1)/28 + 1
}

// FromDate returns the AIRAC cycle that occurred at date. A date before the
// internal epoch (1901-01-10) may return wrong data. The upper limit is
// year 2192.
func FromDate(date time.Time) AIRAC {
	a := date.Sub(epoch) / cycleDuration
	return AIRAC(a)
}

// FromString returns an AIRAC cycle that matches the identifier <yyoo>,
// i.e. the last two digits of the year and the ordinal, each with leading
// zeros. This works for years between 1964 and 2063. Identifiers between
// "6401" and "9913" are interpreted as AIRAC cycles between the years 1964
// and 1999 inclusive. AIRAC cycles between "0001" and "6313" are
// interpreted as AIRAC cycles between the years 2000 and 2063 inclusive.
func FromString(yyoo string) (AIRAC, error) {
	year, ordinal, err := parseIdentifier(yyoo)
	if err != nil {
		return 0, err
	}

	lastAiracOfPreviousYear := FromDate(time.Date(year-1, time.December, 31, 0, 0, 0, 0, time.UTC))
	airac := lastAiracOfPreviousYear + AIRAC(ordinal)

	if airac.Year() != year {
		return 0, fmt.Errorf("illegal AIRAC id \"%s\"", yyoo)
	}

	return airac, nil
}

func parseIdentifier(yyoo string) (year, ordinal int, err error) {
	yyoo = strings.TrimSpace(yyoo)
	if len(yyoo) != 4 {
		return 0, 0, fmt.Errorf("illegal AIRAC id \"%s\"", yyoo)
	}
	yyooInt, err := strconv.Atoi(yyoo)
	if err != nil {
		return 0, 0, fmt.Errorf("illegal AIRAC id \"%s\"", yyoo)
	}

	year, ordinal = (yyooInt/100)+1900, yyooInt%100
	if year <= 1963 {
		year += 100
	}
	return year, ordinal, nil
}

// FromStringMust returns an AIRAC cycle that matches the identifier <yyoo>
// like FromString, but does not return an error. If there is an error it will
// panic instead.
func FromStringMust(yyoo string) AIRAC {
	airac, err := FromString(yyoo)
	if err != nil {
		panic(err)
	}
	return airac
}

// String returns a short representation of this AIRAC cycle. "YYOO"
func (a AIRAC) String() string {
	return fmt.Sprintf("%02d%02d", a.Year()%100, a.Ordinal())
}

// LongString returns a verbose representation of this AIRAC cycle.
// "YYOO (effective: YYYY-MM-DD; expires: YYYY-MM-DD)"
func (a AIRAC) LongString() string {
	n := a + 1
	return fmt.Sprintf("%02d%02d (effective: %s; expires: %s)",
		a.Year()%100,
		a.Ordinal(),
		a.Effective().Format(format),
		n.Effective().Add(-1).Format(format),
	)
}

// ByChrono is an []AIRAC wrapper, that satisfies sort.Interface and can be used
// to chronologically sort AIRAC instances.
type ByChrono []AIRAC

// Len ist the number of elements in the collection.
func (c ByChrono) Len() int { return len(c) }

// Less reports whether the element with
// index i should sort before the element with index j.
func (c ByChrono) Less(i, j int) bool { return c[i] < c[j] }

// Swap swaps the elements with indexes i and j.
func (c ByChrono) Swap(i, j int) { c[i], c[j] = c[j], c[i] }

// static assert
var _ sort.Interface = ByChrono{}
<commit_msg>String formated quoted strings with %q<commit_after>/*
 * Copyright (c) 2018 Johannes Kohnen <wjkohnen@users.noreply.github.com>
 *
 * 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 airac

import (
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"
)

const (
	format        = "2006-01-02"
	cycleDuration = 24192e11 // 4 weeks
)

var (
	epoch = time.Date(1901, time.January, 10, 0, 0, 0, 0, time.UTC)
)

// AIRAC represents an Aeronautical Information Regulation And Control (AIRAC) cycle.
type AIRAC uint16

type Airac = AIRAC

// Effective returns the effective date of this AIRAC cycle.
func (a AIRAC) Effective() time.Time {
	return epoch.Add(time.Duration(a) * cycleDuration)
}

// Year returns the year for this AIRAC cycle's identifier.
func (a AIRAC) Year() int {
	return a.Effective().Year()
}

// Ordinal returns the ordinal for this AIRAC cycle's identifier.
func (a AIRAC) Ordinal() int {
	return (a.Effective().YearDay()-1)/28 + 1
}

// FromDate returns the AIRAC cycle that occurred at date. A date before the
// internal epoch (1901-01-10) may return wrong data. The upper limit is
// year 2192.
func FromDate(date time.Time) AIRAC {
	a := date.Sub(epoch) / cycleDuration
	return AIRAC(a)
}

// FromString returns an AIRAC cycle that matches the identifier <yyoo>,
// i.e. the last two digits of the year and the ordinal, each with leading
// zeros. This works for years between 1964 and 2063. Identifiers between
// "6401" and "9913" are interpreted as AIRAC cycles between the years 1964
// and 1999 inclusive. AIRAC cycles between "0001" and "6313" are
// interpreted as AIRAC cycles between the years 2000 and 2063 inclusive.
func FromString(yyoo string) (AIRAC, error) {
	year, ordinal, err := parseIdentifier(yyoo)
	if err != nil {
		return 0, err
	}

	lastAiracOfPreviousYear := FromDate(time.Date(year-1, time.December, 31, 0, 0, 0, 0, time.UTC))
	airac := lastAiracOfPreviousYear + AIRAC(ordinal)

	if airac.Year() != year {
		return 0, fmt.Errorf("illegal AIRAC id %q", yyoo)
	}

	return airac, nil
}

func parseIdentifier(yyoo string) (year, ordinal int, err error) {
	yyoo = strings.TrimSpace(yyoo)
	if len(yyoo) != 4 {
		return 0, 0, fmt.Errorf("illegal AIRAC id %q", yyoo)
	}
	yyooInt, err := strconv.Atoi(yyoo)
	if err != nil {
		return 0, 0, fmt.Errorf("illegal AIRAC id %q", yyoo)
	}

	year, ordinal = (yyooInt/100)+1900, yyooInt%100
	if year <= 1963 {
		year += 100
	}
	return year, ordinal, nil
}

// FromStringMust returns an AIRAC cycle that matches the identifier <yyoo>
// like FromString, but does not return an error. If there is an error it will
// panic instead.
func FromStringMust(yyoo string) AIRAC {
	airac, err := FromString(yyoo)
	if err != nil {
		panic(err)
	}
	return airac
}

// String returns a short representation of this AIRAC cycle. "YYOO"
func (a AIRAC) String() string {
	return fmt.Sprintf("%02d%02d", a.Year()%100, a.Ordinal())
}

// LongString returns a verbose representation of this AIRAC cycle.
// "YYOO (effective: YYYY-MM-DD; expires: YYYY-MM-DD)"
func (a AIRAC) LongString() string {
	n := a + 1
	return fmt.Sprintf("%02d%02d (effective: %s; expires: %s)",
		a.Year()%100,
		a.Ordinal(),
		a.Effective().Format(format),
		n.Effective().Add(-1).Format(format),
	)
}

// ByChrono is an []AIRAC wrapper, that satisfies sort.Interface and can be used
// to chronologically sort AIRAC instances.
type ByChrono []AIRAC

// Len ist the number of elements in the collection.
func (c ByChrono) Len() int { return len(c) }

// Less reports whether the element with
// index i should sort before the element with index j.
func (c ByChrono) Less(i, j int) bool { return c[i] < c[j] }

// Swap swaps the elements with indexes i and j.
func (c ByChrono) Swap(i, j int) { c[i], c[j] = c[j], c[i] }

// static assert
var _ sort.Interface = ByChrono{}
<|endoftext|>


<commit_before>package imgur

import (
	"encoding/json"
	"errors"
	"strconv"
	"strings"
)

type albumInfoDataWrapper struct {
	Ai      *AlbumInfo `json:"data"`
	Success bool       `json:"success"`
	Status  int        `json:"status"`
}

// AlbumInfo contains all album information provided by imgur
type AlbumInfo struct {
	// The ID for the album
	ID string `json:"id"`
	// The title of the album in the gallery
	Title string `json:"title"`
	// The description of the album in the gallery
	Description string `json:"description"`
	// Time inserted into the gallery, epoch time
	DateTime int `json:"datetime"`
	// The ID of the album cover image
	Cover string `json:"cover"`
	// The width, in pixels, of the album cover image
	CoverWidth int `json:"cover_width"`
	// The height, in pixels, of the album cover image
	CoverHeight int `json:"cover_height"`
	// The account username or null if it's anonymous.
	AccountURL string `json:"account_url"`
	// The account ID or null if it's anonymous.
	AccountID int `json:"account_id"`
	// The privacy level of the album, you can only view public if not logged in as album owner
	Privacy string `json:"privacy"`
	// The view layout of the album.
	Layout string `json:"layout"`
	// The number of album views
	Views int `json:"views"`
	// The URL link to the album
	Link string `json:"link"`
	// Indicates if the current user favorited the image. Defaults to false if not signed in.
	Favorite bool `json:"favorite"`
	// Indicates if the image has been marked as nsfw or not. Defaults to null if information is not available.
	Nsfw bool `json:"nsfw"`
	// If the image has been categorized by our backend then this will contain the section the image belongs in. (funny, cats, adviceanimals, wtf, etc)
	Section string `json:"secion"`
	// Order number of the album on the user's album page (defaults to 0 if their albums haven't been reordered)
	Order int `json:"order"`
	// OPTIONAL, the deletehash, if you're logged in as the album owner
	Deletehash string `json:"deletehash,omitempty"`
	// The total number of images in the album
	ImagesCount int `json:"images_count"`
	// An array of all the images in the album (only available when requesting the direct album)
	Images []ImageInfo `json:"images"`
	// Current rate limit
	Limit *RateLimit
}

// GetAlbumInfo queries imgur for information on a album
// returns album info, status code of the request, error
func (client *Client) GetAlbumInfo(id string) (*AlbumInfo, int, error) {
	body, rl, err := client.getURL("album/" + id)
	if err != nil {
		return nil, -1, errors.New("Problem getting URL for album info ID " + id + " - " + err.Error())
	}
	//client.Log.Debugf("%v\n", body)

	dec := json.NewDecoder(strings.NewReader(body))
	var alb albumInfoDataWrapper
	if err := dec.Decode(&alb); err != nil {
		return nil, -1, errors.New("Problem decoding json for albumID " + id + " - " + err.Error())
	}

	if !alb.Success {
		return nil, alb.Status, errors.New("Request to imgur failed for albumID " + id + " - " + strconv.Itoa(alb.Status))
	}

	alb.Ai.Limit = rl
	return alb.Ai, alb.Status, nil
}
<commit_msg>fix typo in json tag name<commit_after>package imgur

import (
	"encoding/json"
	"errors"
	"strconv"
	"strings"
)

type albumInfoDataWrapper struct {
	Ai      *AlbumInfo `json:"data"`
	Success bool       `json:"success"`
	Status  int        `json:"status"`
}

// AlbumInfo contains all album information provided by imgur
type AlbumInfo struct {
	// The ID for the album
	ID string `json:"id"`
	// The title of the album in the gallery
	Title string `json:"title"`
	// The description of the album in the gallery
	Description string `json:"description"`
	// Time inserted into the gallery, epoch time
	DateTime int `json:"datetime"`
	// The ID of the album cover image
	Cover string `json:"cover"`
	// The width, in pixels, of the album cover image
	CoverWidth int `json:"cover_width"`
	// The height, in pixels, of the album cover image
	CoverHeight int `json:"cover_height"`
	// The account username or null if it's anonymous.
	AccountURL string `json:"account_url"`
	// The account ID or null if it's anonymous.
	AccountID int `json:"account_id"`
	// The privacy level of the album, you can only view public if not logged in as album owner
	Privacy string `json:"privacy"`
	// The view layout of the album.
	Layout string `json:"layout"`
	// The number of album views
	Views int `json:"views"`
	// The URL link to the album
	Link string `json:"link"`
	// Indicates if the current user favorited the image. Defaults to false if not signed in.
	Favorite bool `json:"favorite"`
	// Indicates if the image has been marked as nsfw or not. Defaults to null if information is not available.
	Nsfw bool `json:"nsfw"`
	// If the image has been categorized by our backend then this will contain the section the image belongs in. (funny, cats, adviceanimals, wtf, etc)
	Section string `json:"section"`
	// Order number of the album on the user's album page (defaults to 0 if their albums haven't been reordered)
	Order int `json:"order"`
	// OPTIONAL, the deletehash, if you're logged in as the album owner
	Deletehash string `json:"deletehash,omitempty"`
	// The total number of images in the album
	ImagesCount int `json:"images_count"`
	// An array of all the images in the album (only available when requesting the direct album)
	Images []ImageInfo `json:"images"`
	// Current rate limit
	Limit *RateLimit
}

// GetAlbumInfo queries imgur for information on a album
// returns album info, status code of the request, error
func (client *Client) GetAlbumInfo(id string) (*AlbumInfo, int, error) {
	body, rl, err := client.getURL("album/" + id)
	if err != nil {
		return nil, -1, errors.New("Problem getting URL for album info ID " + id + " - " + err.Error())
	}
	//client.Log.Debugf("%v\n", body)

	dec := json.NewDecoder(strings.NewReader(body))
	var alb albumInfoDataWrapper
	if err := dec.Decode(&alb); err != nil {
		return nil, -1, errors.New("Problem decoding json for albumID " + id + " - " + err.Error())
	}

	if !alb.Success {
		return nil, alb.Status, errors.New("Request to imgur failed for albumID " + id + " - " + strconv.Itoa(alb.Status))
	}

	alb.Ai.Limit = rl
	return alb.Ai, alb.Status, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Daniel Pupius

package dbps

import (
	"bytes"
	"errors"
	"expvar"
	"fmt"
	"io/ioutil"
	"log"
	"path"
	"sort"
	"sync"
	"time"

	"github.com/dpup/dbps/internal/dropbox"
	"github.com/dpup/dbps/internal/goexif/exif"
	"github.com/dpup/rcache"
)

// Album queries dropbox and keeps a list of photos in date order.
type Album struct {
	folder  string
	dropbox *dropbox.Dropbox
	cache   rcache.Cache

	lastHash  string
	photoList photoList
	photoMap  map[string]Photo
	loading   bool
	mu        sync.RWMutex
}

func NewAlbum(folder string, dropbox *dropbox.Dropbox) *Album {
	a := &Album{folder: folder, dropbox: dropbox, cache: rcache.New(folder)}
	a.cache.RegisterFetcher(a.fetchOriginal)
	a.cache.RegisterFetcher(a.fetchThumbnail)

	expvar.Publish(fmt.Sprintf("photos (%s)", folder), expvar.Func(func() interface{} {
		return a.photoMap
	}))

	return a
}

// Monitor starts a go routine which calls Load() every interval to pick up new
// changes
func (a *Album) Monitor(interval time.Duration) {
	c := interval
	go func() {
		for {
			time.Sleep(c)
			err := a.Load()
			if err != nil {
				log.Printf("album: failed to refresh after %s: %s", c, err)
				c = c * 2
			} else {
				c = interval
			}
		}
	}()
}

// Load fetches metadata about the photos in a folder. If the folder hasn't
// changed since Load was last called then no work wil be done.
func (a *Album) Load() error {
	a.mu.Lock()
	if a.loading {
		a.mu.Unlock()
		return errors.New("album: load already in progress")
	}
	a.loading = true
	defer func() { a.loading = false }()
	a.mu.Unlock()

	entry, err := a.dropbox.Metadata(a.folder, true, false, a.lastHash, "", 5000)

	if dbError, ok := err.(*dropbox.Error); ok && dbError.StatusCode == 304 {
		log.Println("album: no metadata changes detected")
		return nil
	} else if err != nil {
		return fmt.Errorf("album: failed to get metadata: %s", err)
	}

	if !entry.IsDir {
		return errors.New("album: provided path was not a directory")
	}

	log.Println("album: loading image metadata")

	var wg sync.WaitGroup

	photos := make(photoList, len(entry.Contents))
	for i, e := range entry.Contents {
		name := path.Base(e.Path)
		clientModified := time.Time(e.ClientMtime)
		dropboxModified := time.Time(e.Modified)

		// e.Hash is empty so use own approximation.
		hash := fmt.Sprintf("%d:%d:%d", e.Bytes, clientModified.Unix(), dropboxModified.Unix())

		// If no entry exists, or the entry is stale, then load the photo to get its
		// exif data. Loads are done in parallel.
		if old, ok := a.photoMap[name]; !ok || old.Hash != hash {
			photos[i] = Photo{
				Filename:        name,
				MimeType:        e.MimeType,
				Size:            e.Bytes,
				Hash:            hash,
				DropboxModified: dropboxModified,
				ExifCreated:     clientModified, // Default to the last modified time.
			}

			wg.Add(1)
			a.cache.Invalidate(originalCacheKey{name}, true)
			go a.loadExifInfo(&photos[i], &wg)

		} else {
			photos[i] = old
		}
	}

	log.Printf("album: waiting for new images to load")
	wg.Wait()
	sort.Sort(photos)

	// TODO(dan): Currently we are not clearing the cache of deleted images, for
	// the existing usecase that is a rare scenario. Can easily be added by
	// asking for deleted items and checking entry.IsDeleted

	a.mu.Lock()
	a.lastHash = entry.Hash
	a.photoList = photos
	a.photoMap = make(map[string]Photo)
	for _, p := range photos {
		a.photoMap[p.Filename] = p
	}
	a.mu.Unlock()

	log.Println("album: metadata load complete")

	return nil
}

// FirstPhoto returns the ... first photo.
func (a *Album) FirstPhoto() Photo {
	return a.photoList[0]
}

// Photo returns the metadata for a photo and the image data, or an error if it doesn't exist.
func (a *Album) Photo(name string) (Photo, []byte, error) {
	if photo, ok := a.photoMap[name]; ok {
		data, err := a.cache.Get(originalCacheKey{name})
		return photo, data, err
	} else {
		return Photo{}, nil, fmt.Errorf("album: no photo with name: %s", name)
	}
}

// Thumbnail returns the metadata for a photo and a thumbnail, or an error if it doesn't exist.
func (a *Album) Thumbnail(name string, width, height uint) (Photo, []byte, error) {
	if photo, ok := a.photoMap[name]; ok {
		data, err := a.cache.Get(thumbCacheKey{name, width, height})
		return photo, data, err
	} else {
		return Photo{}, nil, fmt.Errorf("album: no photo with name: %s", name)
	}
}

// Photos returns a copy of the PhotoList.
func (a *Album) Photos() []Photo {
	a.mu.RLock()
	defer a.mu.RUnlock()
	c := make(photoList, len(a.photoList))
	copy(c, a.photoList)
	return c
}

func (a *Album) loadExifInfo(p *Photo, wg *sync.WaitGroup) {
	defer func() { wg.Done() }()

	data, err := a.cache.Get(originalCacheKey{p.Filename})
	if err != nil {
		log.Printf("album: error renewing cache for %s: %s", p, err)
		return
	}

	x, err := exif.Decode(bytes.NewReader(data))
	if err != nil {
		log.Printf("album: error reading exif for %s: %s", p, err)
		return
	}

	t, err := x.DateTime()
	if err != nil {
		log.Printf("album: error reading exif datetime for %s: %s", p, err)
		return
	}

	p.ExifCreated = t
}

func (a *Album) fetchOriginal(key originalCacheKey) ([]byte, error) {
	// TODO(dan): Add timeout, Download gets stuck.
	filename := key.Filename
	log.Printf("album: fetching %s", filename)
	reader, _, err := a.dropbox.Download(path.Join(a.folder, filename), "", 0)
	if err != nil {
		return []byte{}, err
	}
	return ioutil.ReadAll(reader)
}

func (a *Album) fetchThumbnail(key thumbCacheKey) ([]byte, error) {
	data, err := a.cache.Get(originalCacheKey{key.Filename})
	if err != nil {
		return []byte{}, err
	}
	log.Printf("album: resizing %s", key.Filename)
	return Resize(data, key.Width, key.Height)
}

type originalCacheKey struct {
	Filename string
}

func (t originalCacheKey) Dependencies() []rcache.CacheKey {
	return rcache.NoDeps
}

func (o originalCacheKey) String() string {
	return o.Filename
}

type thumbCacheKey struct {
	Filename string
	Width    uint
	Height   uint
}

func (t thumbCacheKey) Dependencies() []rcache.CacheKey {
	return []rcache.CacheKey{originalCacheKey{t.Filename}}
}

func (t thumbCacheKey) String() string {
	return fmt.Sprintf("%s@%dx%d", t.Filename, t.Width, t.Height)
}
<commit_msg>Use new rcache CacheKey interface<commit_after>// Copyright 2015 Daniel Pupius

package dbps

import (
	"bytes"
	"errors"
	"expvar"
	"fmt"
	"io/ioutil"
	"log"
	"path"
	"sort"
	"sync"
	"time"

	"github.com/dpup/dbps/internal/dropbox"
	"github.com/dpup/dbps/internal/goexif/exif"
	"github.com/dpup/rcache"
)

// Album queries dropbox and keeps a list of photos in date order.
type Album struct {
	folder  string
	dropbox *dropbox.Dropbox
	cache   rcache.Cache

	lastHash  string
	photoList photoList
	photoMap  map[string]Photo
	loading   bool
	mu        sync.RWMutex
}

func NewAlbum(folder string, dropbox *dropbox.Dropbox) *Album {
	a := &Album{folder: folder, dropbox: dropbox, cache: rcache.New(folder)}
	a.cache.RegisterFetcher(a.fetchOriginal)
	a.cache.RegisterFetcher(a.fetchThumbnail)

	expvar.Publish(fmt.Sprintf("photos (%s)", folder), expvar.Func(func() interface{} {
		return a.photoMap
	}))

	return a
}

// Monitor starts a go routine which calls Load() every interval to pick up new
// changes
func (a *Album) Monitor(interval time.Duration) {
	c := interval
	go func() {
		for {
			time.Sleep(c)
			err := a.Load()
			if err != nil {
				log.Printf("album: failed to refresh after %s: %s", c, err)
				c = c * 2
			} else {
				c = interval
			}
		}
	}()
}

// Load fetches metadata about the photos in a folder. If the folder hasn't
// changed since Load was last called then no work wil be done.
func (a *Album) Load() error {
	a.mu.Lock()
	if a.loading {
		a.mu.Unlock()
		return errors.New("album: load already in progress")
	}
	a.loading = true
	defer func() { a.loading = false }()
	a.mu.Unlock()

	entry, err := a.dropbox.Metadata(a.folder, true, false, a.lastHash, "", 5000)

	if dbError, ok := err.(*dropbox.Error); ok && dbError.StatusCode == 304 {
		log.Println("album: no metadata changes detected")
		return nil
	} else if err != nil {
		return fmt.Errorf("album: failed to get metadata: %s", err)
	}

	if !entry.IsDir {
		return errors.New("album: provided path was not a directory")
	}

	log.Println("album: loading image metadata")

	var wg sync.WaitGroup

	photos := make(photoList, len(entry.Contents))
	for i, e := range entry.Contents {
		name := path.Base(e.Path)
		clientModified := time.Time(e.ClientMtime)
		dropboxModified := time.Time(e.Modified)

		// e.Hash is empty so use own approximation.
		hash := fmt.Sprintf("%d:%d:%d", e.Bytes, clientModified.Unix(), dropboxModified.Unix())

		// If no entry exists, or the entry is stale, then load the photo to get its
		// exif data. Loads are done in parallel.
		if old, ok := a.photoMap[name]; !ok || old.Hash != hash {
			photos[i] = Photo{
				Filename:        name,
				MimeType:        e.MimeType,
				Size:            e.Bytes,
				Hash:            hash,
				DropboxModified: dropboxModified,
				ExifCreated:     clientModified, // Default to the last modified time.
			}

			wg.Add(1)
			a.cache.Invalidate(originalCacheKey{name}, true)
			go a.loadExifInfo(&photos[i], &wg)

		} else {
			photos[i] = old
		}
	}

	log.Printf("album: waiting for new images to load")
	wg.Wait()
	sort.Sort(photos)

	// TODO(dan): Currently we are not clearing the cache of deleted images, for
	// the existing usecase that is a rare scenario. Can easily be added by
	// asking for deleted items and checking entry.IsDeleted

	a.mu.Lock()
	a.lastHash = entry.Hash
	a.photoList = photos
	a.photoMap = make(map[string]Photo)
	for _, p := range photos {
		a.photoMap[p.Filename] = p
	}
	a.mu.Unlock()

	log.Println("album: metadata load complete")

	return nil
}

// FirstPhoto returns the ... first photo.
func (a *Album) FirstPhoto() Photo {
	return a.photoList[0]
}

// Photo returns the metadata for a photo and the image data, or an error if it doesn't exist.
func (a *Album) Photo(name string) (Photo, []byte, error) {
	if photo, ok := a.photoMap[name]; ok {
		data, err := a.cache.Get(originalCacheKey{name})
		return photo, data, err
	} else {
		return Photo{}, nil, fmt.Errorf("album: no photo with name: %s", name)
	}
}

// Thumbnail returns the metadata for a photo and a thumbnail, or an error if it doesn't exist.
func (a *Album) Thumbnail(name string, width, height uint) (Photo, []byte, error) {
	if photo, ok := a.photoMap[name]; ok {
		data, err := a.cache.Get(thumbCacheKey{name, width, height})
		return photo, data, err
	} else {
		return Photo{}, nil, fmt.Errorf("album: no photo with name: %s", name)
	}
}

// Photos returns a copy of the PhotoList.
func (a *Album) Photos() []Photo {
	a.mu.RLock()
	defer a.mu.RUnlock()
	c := make(photoList, len(a.photoList))
	copy(c, a.photoList)
	return c
}

func (a *Album) loadExifInfo(p *Photo, wg *sync.WaitGroup) {
	defer func() { wg.Done() }()

	data, err := a.cache.Get(originalCacheKey{p.Filename})
	if err != nil {
		log.Printf("album: error renewing cache for %s: %s", p, err)
		return
	}

	x, err := exif.Decode(bytes.NewReader(data))
	if err != nil {
		log.Printf("album: error reading exif for %s: %s", p, err)
		return
	}

	t, err := x.DateTime()
	if err != nil {
		log.Printf("album: error reading exif datetime for %s: %s", p, err)
		return
	}

	p.ExifCreated = t
}

func (a *Album) fetchOriginal(key originalCacheKey) ([]byte, error) {
	// TODO(dan): Add timeout, Download gets stuck.
	filename := key.Filename
	log.Printf("album: fetching %s", filename)
	reader, _, err := a.dropbox.Download(path.Join(a.folder, filename), "", 0)
	if err != nil {
		return []byte{}, err
	}
	return ioutil.ReadAll(reader)
}

func (a *Album) fetchThumbnail(key thumbCacheKey) ([]byte, error) {
	data, err := a.cache.Get(originalCacheKey{key.Filename})
	if err != nil {
		return []byte{}, err
	}
	log.Printf("album: resizing %s", key.Filename)
	return Resize(data, key.Width, key.Height)
}

type originalCacheKey struct {
	Filename string
}

func (o originalCacheKey) String() string {
	return o.Filename
}

type thumbCacheKey struct {
	Filename string
	Width    uint
	Height   uint
}

func (t thumbCacheKey) Dependencies() []interface{} {
	return []interface{}{originalCacheKey{t.Filename}}
}

func (t thumbCacheKey) String() string {
	return fmt.Sprintf("%s@%dx%d", t.Filename, t.Width, t.Height)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"
	"net"
	"fmt"
);

func main() {
	var mirrorDomain string = "current.cvd.clamav.net";
	mirrorTxtRecord, err := net.LookupTXT(mirrorDomain);

	if (err != nil) {
		msg := fmt.Sprintf("Unable to resolve TXT record for %v", mirrorDomain);
		log.Fatal(msg, err);
		os.Exit(1);
	}

	fmt.Printf("%v", mirrorTxtRecord);
}
<commit_msg>Broke out parsing into separate functions.<commit_after>package main

import (
	"log"
	"os"
	"net"
	"fmt"
	"strings"
);

func main() {
	logger := log.New(os.Stdout, "", log.LstdFlags);
	var mirrorDomain string = "current.cvd.clamav.net";
	var mirrorTxtRecord string = pullTxtRecord(mirrorDomain);
	logger.Printf("TXT record for [%v]: %v", mirrorDomain, mirrorTxtRecord);

	clamav, mainv, dailyv, x, y, z, safebrowsingv, bytecodev := parseTxtRecord(mirrorTxtRecord);
	logger.Printf("TXT record values parsed: " +
		"[clamav=%v,mainv=%v,dailyv=%v,x=%v,y=%v,z=%v,safebrowsingv=%v,bytecodev=%v",
		clamav, mainv, dailyv, x, y, z, safebrowsingv, bytecodev);
}

func pullTxtRecord(mirrorDomain string) (string) {
	mirrorTxtRecords, err := net.LookupTXT(mirrorDomain);

	if (err != nil) {
		msg := fmt.Sprintf("Unable to resolve TXT record for %v", mirrorDomain);
		log.Fatal(msg, err);
		os.Exit(1);
	}

	if (len(mirrorTxtRecords) < 1) {
		msg := fmt.Sprintf("No TXT records returned for %v", mirrorDomain);
		log.Fatal(msg);
		os.Exit(1);
	}

	return mirrorTxtRecords[0];
}

func parseTxtRecord(mirrorTxtRecord string) (string, string, string, string, string, string, string, string) {
	s := strings.SplitN(mirrorTxtRecord, ":", 8);
	return s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7];
}<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/getlantern/fronted"
	"github.com/getlantern/go-igdman/igdman"
	"github.com/getlantern/golog"
	"github.com/getlantern/nattywad"
	"github.com/getlantern/waddell"

	"github.com/getlantern/flashlight/globals"
	"github.com/getlantern/flashlight/nattest"
	"github.com/getlantern/flashlight/statreporter"
	"github.com/getlantern/flashlight/statserver"
)

const (
	PortmapFailure = 50
)

var (
	log = golog.LoggerFor("flashlight.server")

	registerPeriod = 5 * time.Minute
)

type Server struct {
	// Addr: listen address in form of host:port
	Addr string

	// Host: FQDN that is guaranteed to hit this server
	Host string

	// ReadTimeout: (optional) timeout for read ops
	ReadTimeout time.Duration

	// WriteTimeout: (optional) timeout for write ops
	WriteTimeout time.Duration

	CertContext                *fronted.CertContext // context for certificate management
	AllowNonGlobalDestinations bool                 // if true, requests to LAN, Loopback, etc. will be allowed
	AllowedPorts               []int                // if specified, only connections to these ports will be allowed

	waddellClient  *waddell.Client
	nattywadServer *nattywad.Server
	cfg            *ServerConfig
	cfgMutex       sync.RWMutex
}

func (server *Server) Configure(newCfg *ServerConfig) {
	server.cfgMutex.Lock()
	defer server.cfgMutex.Unlock()

	oldCfg := server.cfg

	log.Debug("Server.Configure() called")
	if oldCfg != nil && reflect.DeepEqual(oldCfg, newCfg) {
		log.Debugf("Server configuration unchanged")
		return
	}

	if oldCfg == nil || newCfg.Portmap != oldCfg.Portmap {
		// Portmap changed
		if oldCfg != nil && oldCfg.Portmap > 0 {
			log.Debugf("Attempting to unmap old external port %d", oldCfg.Portmap)
			err := unmapPort(oldCfg.Portmap)
			if err != nil {
				log.Errorf("Unable to unmap old external port: %s", err)
			}
			log.Debugf("Unmapped old external port %d", oldCfg.Portmap)
		}

		if newCfg.Portmap > 0 {
			log.Debugf("Attempting to map new external port %d", newCfg.Portmap)
			err := mapPort(server.Addr, newCfg.Portmap)
			if err != nil {
				log.Errorf("Unable to map new external port: %s", err)
				os.Exit(PortmapFailure)
			}
			log.Debugf("Mapped new external port %d", newCfg.Portmap)
		}
	}

	nattywadIsEnabled := newCfg.WaddellAddr != ""
	nattywadWasEnabled := server.nattywadServer != nil
	waddellAddrChanged := oldCfg == nil && newCfg.WaddellAddr != "" || oldCfg != nil && oldCfg.WaddellAddr != newCfg.WaddellAddr

	if waddellAddrChanged {
		if nattywadWasEnabled {
			server.stopNattywad()
		}
		if nattywadIsEnabled {
			server.startNattywad(newCfg.WaddellAddr)
		}
	}

	server.cfg = newCfg
}

func (server *Server) ListenAndServe() error {
	if server.Host != "" {
		log.Debugf("Running as host %s", server.Host)
	}

	fs := &fronted.Server{
		Addr:                       server.Addr,
		Host:                       server.Host,
		ReadTimeout:                server.ReadTimeout,
		WriteTimeout:               server.WriteTimeout,
		CertContext:                server.CertContext,
		AllowNonGlobalDestinations: server.AllowNonGlobalDestinations,
		AllowedPorts:               server.AllowedPorts,
	}

	if server.cfg.Unencrypted {
		log.Debug("Running in unencrypted mode")
		fs.CertContext = nil
	}

	// Add callbacks to track bytes given
	fs.OnBytesReceived = func(ip string, destAddr string, req *http.Request, bytes int64) {
		onBytesGiven(destAddr, req, bytes)
		statserver.OnBytesReceived(ip, bytes)
	}
	fs.OnBytesSent = func(ip string, destAddr string, req *http.Request, bytes int64) {
		onBytesGiven(destAddr, req, bytes)
		statserver.OnBytesSent(ip, bytes)
	}

	l, err := fs.Listen()
	if err != nil {
		return fmt.Errorf("Unable to listen at %s: %s", server.Addr, err)
	}

	go server.register()

	return fs.Serve(l)
}

func (server *Server) register() {
	for {
		server.cfgMutex.RLock()
		baseUrl := server.cfg.RegisterAt
		server.cfgMutex.RUnlock()
		if baseUrl != "" {
			if globals.InstanceId == "" {
				log.Error("Unable to register server because no InstanceId is configured")
			} else {
				log.Debugf("Registering server at %v", baseUrl)
				registerUrl := baseUrl + "/register"
				vals := url.Values{
					"name": []string{globals.InstanceId},
					"port": []string{"443"},
				}
				resp, err := http.PostForm(registerUrl, vals)
				if err != nil {
					log.Errorf("Unable to register at %v: %v", registerUrl, err)
					return
				} else if resp.StatusCode != 200 {
					bodyString, _ := ioutil.ReadAll(resp.Body)
					log.Errorf("Unexpected response status registering at %v: %d    %v", registerUrl, resp.StatusCode, string(bodyString))
				} else {
					log.Debugf("Successfully registered server at %v", registerUrl)
				}
				resp.Body.Close()
				time.Sleep(registerPeriod)
			}
		}
	}
}

func (server *Server) startNattywad(waddellAddr string) {
	log.Debugf("Connecting to waddell at: %s", waddellAddr)
	var err error
	server.waddellClient, err = waddell.NewClient(&waddell.ClientConfig{
		Dial: func() (net.Conn, error) {
			return net.Dial("tcp", waddellAddr)
		},
		ServerCert:        globals.WaddellCert,
		ReconnectAttempts: 10,
		OnId: func(id waddell.PeerId) {
			log.Debugf("Connected to Waddell!! Id is: %s", id)
		},
	})
	if err != nil {
		log.Errorf("Unable to connect to waddell: %s", err)
		server.waddellClient = nil
		return
	}
	server.nattywadServer = &nattywad.Server{
		Client: server.waddellClient,
		OnSuccess: func(local *net.UDPAddr, remote *net.UDPAddr) bool {
			err := nattest.Serve(local)
			if err != nil {
				log.Error(err.Error())
				return false
			}
			return true
		},
	}
	server.nattywadServer.Start()
}

func (server *Server) stopNattywad() {
	log.Debug("Stopping nattywad server")
	server.nattywadServer.Stop()
	server.nattywadServer = nil
	log.Debug("Stopping waddell client")
	server.waddellClient.Close()
	server.waddellClient = nil
}

func mapPort(addr string, port int) error {
	internalIP, internalPortString, err := net.SplitHostPort(addr)
	if err != nil {
		return fmt.Errorf("Unable to split host and port for %v: %v", addr, err)
	}

	internalPort, err := strconv.Atoi(internalPortString)
	if err != nil {
		return fmt.Errorf("Unable to parse local port: ")
	}

	if internalIP == "" {
		internalIP, err = determineInternalIP()
		if err != nil {
			return fmt.Errorf("Unable to determine internal IP: %s", err)
		}
	}

	igd, err := igdman.NewIGD()
	if err != nil {
		return fmt.Errorf("Unable to get IGD: %s", err)
	}

	igd.RemovePortMapping(igdman.TCP, port)
	err = igd.AddPortMapping(igdman.TCP, internalIP, internalPort, port, 0)
	if err != nil {
		return fmt.Errorf("Unable to map port with igdman %d: %s", port, err)
	}

	return nil
}

func unmapPort(port int) error {
	igd, err := igdman.NewIGD()
	if err != nil {
		return fmt.Errorf("Unable to get IGD: %s", err)
	}

	igd.RemovePortMapping(igdman.TCP, port)
	if err != nil {
		return fmt.Errorf("Unable to unmap port with igdman %d: %s", port, err)
	}

	return nil
}

// determineInternalIP determines the internal IP to use for mapping ports. It
// does this by dialing a website on the public Internet and then finding out
// the LocalAddr for the corresponding connection. This gives us an interface
// that we know has Internet access, which makes it suitable for port mapping.
func determineInternalIP() (string, error) {
	conn, err := net.DialTimeout("tcp", "s3.amazonaws.com:443", 20*time.Second)
	if err != nil {
		return "", fmt.Errorf("Unable to determine local IP: %s", err)
	}
	defer conn.Close()
	host, _, err := net.SplitHostPort(conn.LocalAddr().String())
	return host, err
}

func onBytesGiven(destAddr string, req *http.Request, bytes int64) {
	host, port, _ := net.SplitHostPort(destAddr)
	if port == "" {
		port = "0"
	}

	given := statreporter.CountryDim().
		And("flserver", globals.InstanceId).
		And("destport", port)
	given.Increment("bytesGiven").Add(bytes)
	given.Increment("bytesGivenByFlashlight").Add(bytes)

	clientCountry := req.Header.Get("Cf-Ipcountry")
	if clientCountry != "" {
		givenTo := statreporter.Country(clientCountry)
		givenTo.Increment("bytesGivenTo").Add(bytes)
		givenTo.Increment("bytesGivenToByFlashlight").Add(bytes)
		givenTo.Member("distinctDestHosts", host)

		clientIp := req.Header.Get("X-Forwarded-For")
		if clientIp != "" {
			// clientIp may contain multiple ips, use the first
			ips := strings.Split(clientIp, ",")
			clientIp := strings.TrimSpace(ips[0])
			givenTo.Member("distinctClients", clientIp)
		}

	}
}
<commit_msg>backport of fix to https://github.com/getlantern/lantern/issues/2310<commit_after>package server

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/getlantern/fronted"
	"github.com/getlantern/go-igdman/igdman"
	"github.com/getlantern/golog"
	"github.com/getlantern/nattywad"
	"github.com/getlantern/waddell"

	"github.com/getlantern/flashlight/globals"
	"github.com/getlantern/flashlight/nattest"
	"github.com/getlantern/flashlight/statreporter"
	"github.com/getlantern/flashlight/statserver"
)

const (
	PortmapFailure = 50
)

var (
	log = golog.LoggerFor("flashlight.server")

	registerPeriod = 5 * time.Minute
)

type Server struct {
	// Addr: listen address in form of host:port
	Addr string

	// Host: FQDN that is guaranteed to hit this server
	Host string

	// ReadTimeout: (optional) timeout for read ops
	ReadTimeout time.Duration

	// WriteTimeout: (optional) timeout for write ops
	WriteTimeout time.Duration

	CertContext                *fronted.CertContext // context for certificate management
	AllowNonGlobalDestinations bool                 // if true, requests to LAN, Loopback, etc. will be allowed
	AllowedPorts               []int                // if specified, only connections to these ports will be allowed

	waddellClient  *waddell.Client
	nattywadServer *nattywad.Server
	cfg            *ServerConfig
	cfgMutex       sync.RWMutex
}

func (server *Server) Configure(newCfg *ServerConfig) {
	server.cfgMutex.Lock()
	defer server.cfgMutex.Unlock()

	oldCfg := server.cfg

	log.Debug("Server.Configure() called")
	if oldCfg != nil && reflect.DeepEqual(oldCfg, newCfg) {
		log.Debugf("Server configuration unchanged")
		return
	}

	if oldCfg == nil || newCfg.Portmap != oldCfg.Portmap {
		// Portmap changed
		if oldCfg != nil && oldCfg.Portmap > 0 {
			log.Debugf("Attempting to unmap old external port %d", oldCfg.Portmap)
			err := unmapPort(oldCfg.Portmap)
			if err != nil {
				log.Errorf("Unable to unmap old external port: %s", err)
			}
			log.Debugf("Unmapped old external port %d", oldCfg.Portmap)
		}

		if newCfg.Portmap > 0 {
			log.Debugf("Attempting to map new external port %d", newCfg.Portmap)
			err := mapPort(server.Addr, newCfg.Portmap)
			if err != nil {
				log.Errorf("Unable to map new external port: %s", err)
				os.Exit(PortmapFailure)
			}
			log.Debugf("Mapped new external port %d", newCfg.Portmap)
		}
	}

	nattywadIsEnabled := newCfg.WaddellAddr != ""
	nattywadWasEnabled := server.nattywadServer != nil
	waddellAddrChanged := oldCfg == nil && newCfg.WaddellAddr != "" || oldCfg != nil && oldCfg.WaddellAddr != newCfg.WaddellAddr

	if waddellAddrChanged {
		if nattywadWasEnabled {
			server.stopNattywad()
		}
		if nattywadIsEnabled {
			server.startNattywad(newCfg.WaddellAddr)
		}
	}

	server.cfg = newCfg
}

func (server *Server) ListenAndServe() error {
	if server.Host != "" {
		log.Debugf("Running as host %s", server.Host)
	}

	fs := &fronted.Server{
		Addr:                       server.Addr,
		Host:                       server.Host,
		ReadTimeout:                server.ReadTimeout,
		WriteTimeout:               server.WriteTimeout,
		CertContext:                server.CertContext,
		AllowNonGlobalDestinations: server.AllowNonGlobalDestinations,
		AllowedPorts:               server.AllowedPorts,
	}

	if server.cfg.Unencrypted {
		log.Debug("Running in unencrypted mode")
		fs.CertContext = nil
	}

	// Add callbacks to track bytes given
	fs.OnBytesReceived = func(ip string, destAddr string, req *http.Request, bytes int64) {
		onBytesGiven(destAddr, req, bytes)
		statserver.OnBytesReceived(ip, bytes)
	}
	fs.OnBytesSent = func(ip string, destAddr string, req *http.Request, bytes int64) {
		onBytesGiven(destAddr, req, bytes)
		statserver.OnBytesSent(ip, bytes)
	}

	l, err := fs.Listen()
	if err != nil {
		return fmt.Errorf("Unable to listen at %s: %s", server.Addr, err)
	}

	go server.register()

	return fs.Serve(l)
}

func (server *Server) register() {
	for {
		server.cfgMutex.RLock()
		baseUrl := server.cfg.RegisterAt
		server.cfgMutex.RUnlock()
		if baseUrl != "" {
			if globals.InstanceId == "" {
				log.Error("Unable to register server because no InstanceId is configured")
			} else {
				log.Debugf("Registering server at %v", baseUrl)
				registerUrl := baseUrl + "/register"
				vals := url.Values{
					"name": []string{globals.InstanceId},
					"port": []string{"443"},
				}
				resp, err := http.PostForm(registerUrl, vals)
				if err != nil {
					log.Errorf("Unable to register at %v: %v", registerUrl, err)
				} else if resp.StatusCode != 200 {
					bodyString, _ := ioutil.ReadAll(resp.Body)
					log.Errorf("Unexpected response status registering at %v: %d    %v", registerUrl, resp.StatusCode, string(bodyString))
				} else {
					log.Debugf("Successfully registered server at %v", registerUrl)
				}
				if err == nil {
					resp.Body.Close()
				}
				time.Sleep(registerPeriod)
			}
		}
	}
}

func (server *Server) startNattywad(waddellAddr string) {
	log.Debugf("Connecting to waddell at: %s", waddellAddr)
	var err error
	server.waddellClient, err = waddell.NewClient(&waddell.ClientConfig{
		Dial: func() (net.Conn, error) {
			return net.Dial("tcp", waddellAddr)
		},
		ServerCert:        globals.WaddellCert,
		ReconnectAttempts: 10,
		OnId: func(id waddell.PeerId) {
			log.Debugf("Connected to Waddell!! Id is: %s", id)
		},
	})
	if err != nil {
		log.Errorf("Unable to connect to waddell: %s", err)
		server.waddellClient = nil
		return
	}
	server.nattywadServer = &nattywad.Server{
		Client: server.waddellClient,
		OnSuccess: func(local *net.UDPAddr, remote *net.UDPAddr) bool {
			err := nattest.Serve(local)
			if err != nil {
				log.Error(err.Error())
				return false
			}
			return true
		},
	}
	server.nattywadServer.Start()
}

func (server *Server) stopNattywad() {
	log.Debug("Stopping nattywad server")
	server.nattywadServer.Stop()
	server.nattywadServer = nil
	log.Debug("Stopping waddell client")
	server.waddellClient.Close()
	server.waddellClient = nil
}

func mapPort(addr string, port int) error {
	internalIP, internalPortString, err := net.SplitHostPort(addr)
	if err != nil {
		return fmt.Errorf("Unable to split host and port for %v: %v", addr, err)
	}

	internalPort, err := strconv.Atoi(internalPortString)
	if err != nil {
		return fmt.Errorf("Unable to parse local port: ")
	}

	if internalIP == "" {
		internalIP, err = determineInternalIP()
		if err != nil {
			return fmt.Errorf("Unable to determine internal IP: %s", err)
		}
	}

	igd, err := igdman.NewIGD()
	if err != nil {
		return fmt.Errorf("Unable to get IGD: %s", err)
	}

	igd.RemovePortMapping(igdman.TCP, port)
	err = igd.AddPortMapping(igdman.TCP, internalIP, internalPort, port, 0)
	if err != nil {
		return fmt.Errorf("Unable to map port with igdman %d: %s", port, err)
	}

	return nil
}

func unmapPort(port int) error {
	igd, err := igdman.NewIGD()
	if err != nil {
		return fmt.Errorf("Unable to get IGD: %s", err)
	}

	igd.RemovePortMapping(igdman.TCP, port)
	if err != nil {
		return fmt.Errorf("Unable to unmap port with igdman %d: %s", port, err)
	}

	return nil
}

// determineInternalIP determines the internal IP to use for mapping ports. It
// does this by dialing a website on the public Internet and then finding out
// the LocalAddr for the corresponding connection. This gives us an interface
// that we know has Internet access, which makes it suitable for port mapping.
func determineInternalIP() (string, error) {
	conn, err := net.DialTimeout("tcp", "s3.amazonaws.com:443", 20*time.Second)
	if err != nil {
		return "", fmt.Errorf("Unable to determine local IP: %s", err)
	}
	defer conn.Close()
	host, _, err := net.SplitHostPort(conn.LocalAddr().String())
	return host, err
}

func onBytesGiven(destAddr string, req *http.Request, bytes int64) {
	host, port, _ := net.SplitHostPort(destAddr)
	if port == "" {
		port = "0"
	}

	given := statreporter.CountryDim().
		And("flserver", globals.InstanceId).
		And("destport", port)
	given.Increment("bytesGiven").Add(bytes)
	given.Increment("bytesGivenByFlashlight").Add(bytes)

	clientCountry := req.Header.Get("Cf-Ipcountry")
	if clientCountry != "" {
		givenTo := statreporter.Country(clientCountry)
		givenTo.Increment("bytesGivenTo").Add(bytes)
		givenTo.Increment("bytesGivenToByFlashlight").Add(bytes)
		givenTo.Member("distinctDestHosts", host)

		clientIp := req.Header.Get("X-Forwarded-For")
		if clientIp != "" {
			// clientIp may contain multiple ips, use the first
			ips := strings.Split(clientIp, ",")
			clientIp := strings.TrimSpace(ips[0])
			givenTo.Member("distinctClients", clientIp)
		}

	}
}
<|endoftext|>