<commit_before>package dependency

import (
	"fmt"
	"net/url"
	"testing"
	"time"

	"github.com/hashicorp/vault/api"
	"github.com/stretchr/testify/assert"
)

func TestNewVaultReadQuery(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name string
		i    string
		exp  *VaultReadQuery
		err  bool
	}{
		{
			"empty",
			"",
			nil,
			true,
		},
		{
			"path",
			"path",
			&VaultReadQuery{
				rawPath:     "path",
				queryValues: url.Values{},
			},
			false,
		},
		{
			"leading_slash",
			"/leading/slash",
			&VaultReadQuery{
				rawPath:     "leading/slash",
				queryValues: url.Values{},
			},
			false,
		},
		{
			"trailing_slash",
			"trailing/slash/",
			&VaultReadQuery{
				rawPath:     "trailing/slash",
				queryValues: url.Values{},
			},
			false,
		},
		{
			"query_param",
			"path?version=3",
			&VaultReadQuery{
				rawPath: "path",
				queryValues: url.Values{
					"version": []string{"3"},
				},
			},
			false,
		},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
			act, err := NewVaultReadQuery(tc.i)
			if (err != nil) != tc.err {
				t.Fatal(err)
			}

			if act != nil {
				act.stopCh = nil
			}

			assert.Equal(t, tc.exp, act)
		})
	}
}

func TestVaultReadQuery_Fetch(t *testing.T) {
	t.Parallel()

	clients, vault := testVaultServer(t)
	defer vault.Stop()

	err := vault.CreateSecret("foo/bar", map[string]interface{}{
		"ttl": "100ms", // explicitly make this a short duration for testing
		"zip": "zap",
	})
	if err != nil {
		t.Fatal(err)
	}

	cases := []struct {
		name string
		i    string
		exp  interface{}
		err  bool
	}{
		{
			"exists",
			"secret/foo/bar",
			&Secret{
				Data: map[string]interface{}{
					"ttl": "100ms",
					"zip": "zap",
				},
			},
			false,
		},
		{
			"no_exist",
			"not/a/real/path/like/ever",
			nil,
			true,
		},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
			d, err := NewVaultReadQuery(tc.i)
			if err != nil {
				t.Fatal(err)
			}

			act, _, err := d.Fetch(clients, nil)
			if (err != nil) != tc.err {
				t.Fatal(err)
			}

			if act != nil {
				act.(*Secret).RequestID = ""
				act.(*Secret).LeaseID = ""
				act.(*Secret).LeaseDuration = 0
				act.(*Secret).Renewable = false
			}

			assert.Equal(t, tc.exp, act)
		})
	}

	t.Run("stops", func(t *testing.T) {
		d, err := NewVaultReadQuery("secret/foo/bar")
		if err != nil {
			t.Fatal(err)
		}

		dataCh := make(chan interface{}, 1)
		errCh := make(chan error, 1)
		go func() {
			for {
				data, _, err := d.Fetch(clients, nil)
				if err != nil {
					errCh <- err
					return
				}
				dataCh <- data
			}
		}()

		select {
		case err := <-errCh:
			t.Fatal(err)
		case <-dataCh:
		}

		d.Stop()

		select {
		case err := <-errCh:
			if err != ErrStopped {
				t.Fatal(err)
			}
		case <-time.After(500 * time.Millisecond):
			t.Errorf("did not stop")
		}
	})

	t.Run("fires_changes", func(t *testing.T) {
		d, err := NewVaultReadQuery("secret/foo/bar")
		if err != nil {
			t.Fatal(err)
		}

		_, qm, err := d.Fetch(clients, nil)
		if err != nil {
			t.Fatal(err)
		}

		dataCh := make(chan interface{}, 1)
		errCh := make(chan error, 1)
		go func() {
			for {
				data, _, err := d.Fetch(clients, &QueryOptions{WaitIndex: qm.LastIndex})
				if err != nil {
					errCh <- err
					return
				}
				dataCh <- data
				return
			}
		}()

		select {
		case err := <-errCh:
			t.Fatal(err)
		case <-dataCh:
		}
	})
}

func TestVaultReadQuery_Fetch_KVv2(t *testing.T) {
	t.Parallel()

	clients, vault := testVaultServer(t)
	defer vault.Stop()

	// Enable v2 kv for versioned secrets
	vc := clients.Vault()
	if err := vc.Sys().TuneMount("secret", api.MountConfigInput{
		Options: map[string]string{
			"version": "2",
		},
	}); err != nil {
		t.Fatalf("Error tuning secrets engine: %s", err)
	}

	// Write an initial value to the secret path
	err := vault.CreateSecret("data/foo/bar", map[string]interface{}{
		"data": map[string]interface{}{
			"ttl": "100ms", // explicitly make this a short duration for testing
			"zip": "zap",
		},
	})
	if err != nil {
		t.Fatal(err)
	}

	// Write a new value to increment the version
	err = vault.CreateSecret("data/foo/bar", map[string]interface{}{
		"data": map[string]interface{}{
			"ttl": "100ms", // explicitly make this a short duration for testing
			"zip": "zop",
		},
	})
	if err != nil {
		t.Fatal(err)
	}

	cases := []struct {
		name string
		i    string
		exp  interface{}
		err  bool
	}{
		{
			"exists",
			"secret/foo/bar",
			&Secret{
				Data: map[string]interface{}{
					"data": map[string]interface{}{
						"ttl": "100ms", // explicitly make this a short duration for testing
						"zip": "zop",
					},
				},
			},
			false,
		},
		{
			"/data in path",
			"secret/data/foo/bar",
			&Secret{
				Data: map[string]interface{}{
					"data": map[string]interface{}{
						"ttl": "100ms", // explicitly make this a short duration for testing
						"zip": "zop",
					},
				},
			},
			false,
		},
		{
			"version=1",
			"secret/foo/bar?version=1",
			&Secret{
				Data: map[string]interface{}{
					"data": map[string]interface{}{
						"ttl": "100ms", // explicitly make this a short duration for testing
						"zip": "zap",
					},
				},
			},
			false,
		},
		{
			"no_exist",
			"not/a/real/path/like/ever",
			nil,
			true,
		},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
			d, err := NewVaultReadQuery(tc.i)
			if err != nil {
				t.Fatal(err)
			}

			act, _, err := d.Fetch(clients, nil)
			if (err != nil) != tc.err {
				t.Fatal(err)
			}

			if act != nil {
				act.(*Secret).RequestID = ""
				act.(*Secret).LeaseID = ""
				act.(*Secret).LeaseDuration = 0
				act.(*Secret).Renewable = false
				tc.exp.(*Secret).Data["metadata"] = act.(*Secret).Data["metadata"]
			}

			assert.Equal(t, tc.exp, act)
		})
	}

	t.Run("stops", func(t *testing.T) {
		d, err := NewVaultReadQuery("secret/foo/bar")
		if err != nil {
			t.Fatal(err)
		}

		dataCh := make(chan interface{}, 1)
		errCh := make(chan error, 1)
		go func() {
			for {
				data, _, err := d.Fetch(clients, nil)
				if err != nil {
					errCh <- err
					return
				}
				dataCh <- data
			}
		}()

		select {
		case err := <-errCh:
			t.Fatal(err)
		case <-dataCh:
		}

		d.Stop()

		select {
		case err := <-errCh:
			if err != ErrStopped {
				t.Fatal(err)
			}
		case <-time.After(500 * time.Millisecond):
			t.Errorf("did not stop")
		}
	})

	for _, dataPrefix := range []string{"", "/data"} {
		t.Run(fmt.Sprintf("fires_changes%s", dataPrefix), func(t *testing.T) {
			d, err := NewVaultReadQuery(fmt.Sprintf("secret%s/foo/bar", dataPrefix))
			if err != nil {
				t.Fatal(err)
			}

			_, qm, err := d.Fetch(clients, nil)
			if err != nil {
				t.Fatal(err)
			}

			dataCh := make(chan interface{}, 1)
			errCh := make(chan error, 1)
			go func() {
				for {
					data, _, err := d.Fetch(clients, &QueryOptions{WaitIndex: qm.LastIndex})
					if err != nil {
						errCh <- err
						return
					}
					dataCh <- data
					return
				}
			}()

			select {
			case err := <-errCh:
				t.Fatal(err)
			case <-dataCh:
			}
		})
	}
}

func TestVaultReadQuery_String(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name string
		i    string
		exp  string
	}{
		{
			"path",
			"path",
			"vault.read(path)",
		},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
			d, err := NewVaultReadQuery(tc.i)
			if err != nil {
				t.Fatal(err)
			}
			assert.Equal(t, tc.exp, d.String())
		})
	}
}
<commit_msg>add a test for anonymous access of pki ca cert<commit_after>package dependency

import (
	"fmt"
	"net/url"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/vault/api"
	"github.com/stretchr/testify/assert"
)

func TestNewVaultReadQuery(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name string
		i    string
		exp  *VaultReadQuery
		err  bool
	}{
		{
			"empty",
			"",
			nil,
			true,
		},
		{
			"path",
			"path",
			&VaultReadQuery{
				rawPath:     "path",
				queryValues: url.Values{},
			},
			false,
		},
		{
			"leading_slash",
			"/leading/slash",
			&VaultReadQuery{
				rawPath:     "leading/slash",
				queryValues: url.Values{},
			},
			false,
		},
		{
			"trailing_slash",
			"trailing/slash/",
			&VaultReadQuery{
				rawPath:     "trailing/slash",
				queryValues: url.Values{},
			},
			false,
		},
		{
			"query_param",
			"path?version=3",
			&VaultReadQuery{
				rawPath: "path",
				queryValues: url.Values{
					"version": []string{"3"},
				},
			},
			false,
		},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
			act, err := NewVaultReadQuery(tc.i)
			if (err != nil) != tc.err {
				t.Fatal(err)
			}

			if act != nil {
				act.stopCh = nil
			}

			assert.Equal(t, tc.exp, act)
		})
	}
}

func TestVaultReadQuery_Fetch(t *testing.T) {
	t.Parallel()

	clients, vault := testVaultServer(t)
	defer vault.Stop()

	err := vault.CreateSecret("foo/bar", map[string]interface{}{
		"ttl": "100ms", // explicitly make this a short duration for testing
		"zip": "zap",
	})
	if err != nil {
		t.Fatal(err)
	}

	cases := []struct {
		name string
		i    string
		exp  interface{}
		err  bool
	}{
		{
			"exists",
			"secret/foo/bar",
			&Secret{
				Data: map[string]interface{}{
					"ttl": "100ms",
					"zip": "zap",
				},
			},
			false,
		},
		{
			"no_exist",
			"not/a/real/path/like/ever",
			nil,
			true,
		},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
			d, err := NewVaultReadQuery(tc.i)
			if err != nil {
				t.Fatal(err)
			}

			act, _, err := d.Fetch(clients, nil)
			if (err != nil) != tc.err {
				t.Fatal(err)
			}

			if act != nil {
				act.(*Secret).RequestID = ""
				act.(*Secret).LeaseID = ""
				act.(*Secret).LeaseDuration = 0
				act.(*Secret).Renewable = false
			}

			assert.Equal(t, tc.exp, act)
		})
	}

	t.Run("stops", func(t *testing.T) {
		d, err := NewVaultReadQuery("secret/foo/bar")
		if err != nil {
			t.Fatal(err)
		}

		dataCh := make(chan interface{}, 1)
		errCh := make(chan error, 1)
		go func() {
			for {
				data, _, err := d.Fetch(clients, nil)
				if err != nil {
					errCh <- err
					return
				}
				dataCh <- data
			}
		}()

		select {
		case err := <-errCh:
			t.Fatal(err)
		case <-dataCh:
		}

		d.Stop()

		select {
		case err := <-errCh:
			if err != ErrStopped {
				t.Fatal(err)
			}
		case <-time.After(500 * time.Millisecond):
			t.Errorf("did not stop")
		}
	})

	t.Run("fires_changes", func(t *testing.T) {
		d, err := NewVaultReadQuery("secret/foo/bar")
		if err != nil {
			t.Fatal(err)
		}

		_, qm, err := d.Fetch(clients, nil)
		if err != nil {
			t.Fatal(err)
		}

		dataCh := make(chan interface{}, 1)
		errCh := make(chan error, 1)
		go func() {
			for {
				data, _, err := d.Fetch(clients, &QueryOptions{WaitIndex: qm.LastIndex})
				if err != nil {
					errCh <- err
					return
				}
				dataCh <- data
				return
			}
		}()

		select {
		case err := <-errCh:
			t.Fatal(err)
		case <-dataCh:
		}
	})
}

func TestVaultReadQuery_Fetch_KVv2(t *testing.T) {
	t.Parallel()

	clients, vault := testVaultServer(t)
	defer vault.Stop()

	// Enable v2 kv for versioned secrets
	vc := clients.Vault()
	if err := vc.Sys().TuneMount("secret", api.MountConfigInput{
		Options: map[string]string{
			"version": "2",
		},
	}); err != nil {
		t.Fatalf("Error tuning secrets engine: %s", err)
	}

	// Write an initial value to the secret path
	err := vault.CreateSecret("data/foo/bar", map[string]interface{}{
		"data": map[string]interface{}{
			"ttl": "100ms", // explicitly make this a short duration for testing
			"zip": "zap",
		},
	})
	if err != nil {
		t.Fatal(err)
	}

	// Write a new value to increment the version
	err = vault.CreateSecret("data/foo/bar", map[string]interface{}{
		"data": map[string]interface{}{
			"ttl": "100ms", // explicitly make this a short duration for testing
			"zip": "zop",
		},
	})
	if err != nil {
		t.Fatal(err)
	}

	cases := []struct {
		name string
		i    string
		exp  interface{}
		err  bool
	}{
		{
			"exists",
			"secret/foo/bar",
			&Secret{
				Data: map[string]interface{}{
					"data": map[string]interface{}{
						"ttl": "100ms", // explicitly make this a short duration for testing
						"zip": "zop",
					},
				},
			},
			false,
		},
		{
			"/data in path",
			"secret/data/foo/bar",
			&Secret{
				Data: map[string]interface{}{
					"data": map[string]interface{}{
						"ttl": "100ms", // explicitly make this a short duration for testing
						"zip": "zop",
					},
				},
			},
			false,
		},
		{
			"version=1",
			"secret/foo/bar?version=1",
			&Secret{
				Data: map[string]interface{}{
					"data": map[string]interface{}{
						"ttl": "100ms", // explicitly make this a short duration for testing
						"zip": "zap",
					},
				},
			},
			false,
		},
		{
			"no_exist",
			"not/a/real/path/like/ever",
			nil,
			true,
		},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
			d, err := NewVaultReadQuery(tc.i)
			if err != nil {
				t.Fatal(err)
			}

			act, _, err := d.Fetch(clients, nil)
			if (err != nil) != tc.err {
				t.Fatal(err)
			}

			if act != nil {
				act.(*Secret).RequestID = ""
				act.(*Secret).LeaseID = ""
				act.(*Secret).LeaseDuration = 0
				act.(*Secret).Renewable = false
				tc.exp.(*Secret).Data["metadata"] = act.(*Secret).Data["metadata"]
			}

			assert.Equal(t, tc.exp, act)
		})
	}

	t.Run("stops", func(t *testing.T) {
		d, err := NewVaultReadQuery("secret/foo/bar")
		if err != nil {
			t.Fatal(err)
		}

		dataCh := make(chan interface{}, 1)
		errCh := make(chan error, 1)
		go func() {
			for {
				data, _, err := d.Fetch(clients, nil)
				if err != nil {
					errCh <- err
					return
				}
				dataCh <- data
			}
		}()

		select {
		case err := <-errCh:
			t.Fatal(err)
		case <-dataCh:
		}

		d.Stop()

		select {
		case err := <-errCh:
			if err != ErrStopped {
				t.Fatal(err)
			}
		case <-time.After(500 * time.Millisecond):
			t.Errorf("did not stop")
		}
	})

	for _, dataPrefix := range []string{"", "/data"} {
		t.Run(fmt.Sprintf("fires_changes%s", dataPrefix), func(t *testing.T) {
			d, err := NewVaultReadQuery(fmt.Sprintf("secret%s/foo/bar", dataPrefix))
			if err != nil {
				t.Fatal(err)
			}

			_, qm, err := d.Fetch(clients, nil)
			if err != nil {
				t.Fatal(err)
			}

			dataCh := make(chan interface{}, 1)
			errCh := make(chan error, 1)
			go func() {
				for {
					data, _, err := d.Fetch(clients, &QueryOptions{WaitIndex: qm.LastIndex})
					if err != nil {
						errCh <- err
						return
					}
					dataCh <- data
					return
				}
			}()

			select {
			case err := <-errCh:
				t.Fatal(err)
			case <-dataCh:
			}
		})
	}
}

// TestVaultReadQuery_Fetch_PKI_Anonymous asserts that vault.read can fetch a
// pki ca public cert even even when running unauthenticated client.
func TestVaultReadQuery_Fetch_PKI_Anonymous(t *testing.T) {
	t.Parallel()

	clients, vault := testVaultServer(t)
	defer vault.Stop()

	err := clients.Vault().Sys().Mount("pki", &api.MountInput{
		Type: "pki",
	})
	if err != nil {
		t.Fatal(err)
	}

	vc := clients.Vault()
	_, err = vc.Logical().Write("sys/policies/acl/secrets-only", map[string]interface{}{
		"policy": `path "secret/*" { capabilities = ["create", "read"] }`,
	})
	if err != nil {
		t.Fatal(err)
	}

	_, err = vc.Logical().Write("pki/root/generate/internal", map[string]interface{}{
		"common_name": "example.com",
		"ttl":         "24h",
	})

	anonClient := NewClientSet()
	anonClient.CreateVaultClient(&CreateVaultClientInput{
		Address: vault.Address,
		Token:   "",
	})
	_, err = anonClient.vault.client.Auth().Token().LookupSelf()
	if err == nil || !strings.Contains(err.Error(), "missing client token") {
		t.Fatalf("expected a missing client token error but found: %v", err)
	}

	d, err := NewVaultReadQuery("pki/cert/ca")
	if err != nil {
		t.Fatal(err)
	}

	act, _, err := d.Fetch(anonClient, nil)
	if err != nil {
		t.Fatal(err)
	}

	sec, ok := act.(*Secret)
	if !ok {
		t.Fatalf("expected secret but found %v", reflect.TypeOf(act))
	}
	cert, ok := sec.Data["certificate"].(string)
	if !ok || !strings.Contains(cert, "BEGIN") {
		t.Fatalf("expected a cert but found: %v", cert)
	}
}

func TestVaultReadQuery_String(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name string
		i    string
		exp  string
	}{
		{
			"path",
			"path",
			"vault.read(path)",
		},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
			d, err := NewVaultReadQuery(tc.i)
			if err != nil {
				t.Fatal(err)
			}
			assert.Equal(t, tc.exp, d.String())
		})
	}
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
)

func TestResponseInitTerms(t *testing.T) {
	terms := []string{"Foo-Foo", "BAR*BAR?", "バズ"}
	InitTerms(terms)
	if terms[0] != "foo-foo" {
		t.Error("failed to initialize terms")
	}

	if terms[1] != "bar*bar?" {
		t.Error("failed to initialize terms")
	}

	if terms[2] != "バズ" {
		t.Error("failed to initialize terms")
	}
}

func TestResponseContainTerms(t *testing.T) {
	if !ContainTerms([]string{"foo-bar"}, "FOO-Bar") {
		t.Error("failed to match terms")
	}

	if !ContainTerms([]string{"foo-bar"}, "1000foo-bar2000") {
		t.Error("failed to match terms")
	}

	if !ContainTerms([]string{"foo", "oops"}, "foops") {
		t.Error("failed to match terms")
	}

	if ContainTerms([]string{"foo-bar"}, "foo--bar") {
		t.Error("failed to match terms")
	}
}
<commit_msg>Add test: TestResponseAddItem<commit_after>package main

import (
	"testing"
)

func TestResponseInitTerms(t *testing.T) {
	terms := []string{"Foo-Foo", "BAR*BAR?", "バズ"}
	InitTerms(terms)
	if terms[0] != "foo-foo" {
		t.Error("failed to initialize terms")
	}

	if terms[1] != "bar*bar?" {
		t.Error("failed to initialize terms")
	}

	if terms[2] != "バズ" {
		t.Error("failed to initialize terms")
	}
}

func TestResponseContainTerms(t *testing.T) {
	if !ContainTerms([]string{"foo-bar"}, "FOO-Bar") {
		t.Error("failed to match terms")
	}

	if !ContainTerms([]string{"foo-bar"}, "1000foo-bar2000") {
		t.Error("failed to match terms")
	}

	if !ContainTerms([]string{"foo", "oops"}, "foops") {
		t.Error("failed to match terms")
	}

	if ContainTerms([]string{"foo-bar"}, "foo--bar") {
		t.Error("failed to match terms")
	}
}

func TestResponseAddItem(t *testing.T) {
	r := NewResponse()
	r.AddItem(&ResponseItem{Title: "title-foo"})

	if r.Items[0].Title != "title-foo" {
		t.Error("failed to add item to new response")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"syscall"
	"os/exec"
	"strings"
	"time"
)

// Parse args:
var (
	benchmarkClientSshPrelude, benchmarkServerSshPrelude string
	benchmarkClientCmd, benchmarkServerCmd string
)
func init() {
	flag.StringVar(&benchmarkServerSshPrelude, "benchmark-server-ssh-prelude", "", "SSH prelude for server connection.")
	flag.StringVar(&benchmarkClientSshPrelude, "benchmark-client-ssh-prelude", "", "SSH prelude for client connection.")
	flag.StringVar(&benchmarkServerCmd, "benchmark-server-cmd", "", "Command to run on the benchmarking server.")
	flag.StringVar(&benchmarkClientCmd, "benchmark-client-cmd", "", "Command to run on the benchmarking client.")

	flag.Parse()
}
func main() {
	serverArgs := append(strings.Split(benchmarkServerSshPrelude, " "), benchmarkServerCmd)
	clientArgs := append(strings.Split(benchmarkClientSshPrelude, " "), benchmarkClientCmd)

	server := exec.Command("ssh", serverArgs...)
	client := exec.Command("ssh", clientArgs...)

	// start server, keep it running:
	err := server.Start()
	if err != nil {
		log.Fatal(err)
	}

	// wait for server to initialize:
	time.Sleep(2 * time.Second)

	// run client to completion, printing collected stdout:
	out, err := client.Output()
	fmt.Println(string(out))
	if err != nil {
		log.Fatal(err)
	}

	// signal the server to shut down:
	err = server.Process.Signal(syscall.SIGINT)
	if err != nil {
		log.Fatal(err)
	}

	// wait for server to shut down:
	_, err = server.Process.Wait()
	if err != nil {
		log.Fatal("wait error ", err)
	}
}
<commit_msg>remove unneeded script<commit_after><|endoftext|>


<commit_before>package http

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

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"

	cmds "github.com/jbenet/go-ipfs/commands"
	u "github.com/jbenet/go-ipfs/util"
)

var log = u.Logger("commands/http")

type Handler struct {
	ctx    cmds.Context
	root   *cmds.Command
	origin string
}

var ErrNotFound = errors.New("404 page not found")

const (
	streamHeader           = "X-Stream-Output"
	channelHeader          = "X-Chunked-Output"
	contentTypeHeader      = "Content-Type"
	contentLengthHeader    = "Content-Length"
	transferEncodingHeader = "Transfer-Encoding"
	applicationJson        = "application/json"
)

var mimeTypes = map[string]string{
	cmds.JSON: "application/json",
	cmds.XML:  "application/xml",
	cmds.Text: "text/plain",
}

func NewHandler(ctx cmds.Context, root *cmds.Command, origin string) *Handler {
	// allow whitelisted origins (so we can make API requests from the browser)
	if len(origin) > 0 {
		log.Info("Allowing API requests from origin: " + origin)
	}

	return &Handler{ctx, root, origin}
}

func (i Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// create a context.Context to pass into the commands.
	ctx, cancel := context.WithCancel(context.TODO())
	defer cancel()
	i.ctx.Context = ctx

	log.Debug("Incoming API request: ", r.URL)

	if len(i.origin) > 0 {
		w.Header().Set("Access-Control-Allow-Origin", i.origin)
	}
	w.Header().Set("Access-Control-Allow-Headers", "Content-Type")

	req, err := Parse(r, i.root)
	if err != nil {
		if err == ErrNotFound {
			w.WriteHeader(http.StatusNotFound)
		} else {
			w.WriteHeader(http.StatusBadRequest)
		}
		w.Write([]byte(err.Error()))
		return
	}
	req.SetContext(i.ctx)

	// call the command
	res := i.root.Call(req)

	// set the Content-Type based on res output
	if _, ok := res.Output().(io.Reader); ok {
		// we don't set the Content-Type for streams, so that browsers can MIME-sniff the type themselves
		// we set this header so clients have a way to know this is an output stream
		// (not marshalled command output)
		// TODO: set a specific Content-Type if the command response needs it to be a certain type
		w.Header().Set(streamHeader, "1")

	} else {
		enc, found, err := req.Option(cmds.EncShort).String()
		if err != nil || !found {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		mime := mimeTypes[enc]
		w.Header().Set(contentTypeHeader, mime)
	}

	// if response contains an error, write an HTTP error status code
	if e := res.Error(); e != nil {
		if e.Code == cmds.ErrClient {
			w.WriteHeader(http.StatusBadRequest)
		} else {
			w.WriteHeader(http.StatusInternalServerError)
		}
	}

	out, err := res.Reader()
	if err != nil {
		w.Header().Set(contentTypeHeader, "text/plain")
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(err.Error()))
		return
	}

	// if output is a channel and user requested streaming channels,
	// use chunk copier for the output
	_, isChan := res.Output().(chan interface{})
	streamChans, _, _ := req.Option("stream-channels").Bool()
	if isChan && streamChans {
		err = copyChunks(applicationJson, w, out)
		if err != nil {
			log.Error(err)
		}
		return
	}

	flushCopy(w, out)
}

// flushCopy Copies from an io.Reader to a http.ResponseWriter.
// Flushes chunks over HTTP stream as they are read (if supported by transport).
func flushCopy(w http.ResponseWriter, out io.Reader) error {
	if _, ok := w.(http.Flusher); !ok {
		return copyChunks("", w, out)
	}

	io.Copy(&flushResponse{w}, out)
	return nil
}

// Copies from an io.Reader to a http.ResponseWriter.
// Flushes chunks over HTTP stream as they are read (if supported by transport).
func copyChunks(contentType string, w http.ResponseWriter, out io.Reader) error {
	hijacker, ok := w.(http.Hijacker)
	if !ok {
		return errors.New("Could not create hijacker")
	}
	conn, writer, err := hijacker.Hijack()
	if err != nil {
		return err
	}
	defer conn.Close()

	writer.WriteString("HTTP/1.1 200 OK\r\n")
	if contentType != "" {
		writer.WriteString(contentTypeHeader + ": " + contentType + "\r\n")
	}
	writer.WriteString(transferEncodingHeader + ": chunked\r\n")
	writer.WriteString(channelHeader + ": 1\r\n\r\n")

	buf := make([]byte, 32*1024)

	for {
		n, err := out.Read(buf)

		if n > 0 {
			length := fmt.Sprintf("%x\r\n", n)
			writer.WriteString(length)

			_, err := writer.Write(buf[0:n])
			if err != nil {
				return err
			}

			writer.WriteString("\r\n")
			writer.Flush()
		}

		if err != nil && err != io.EOF {
			return err
		}
		if err == io.EOF {
			break
		}
	}

	writer.WriteString("0\r\n\r\n")
	writer.Flush()

	return nil
}

type flushResponse struct {
	W http.ResponseWriter
}

func (fr *flushResponse) Write(buf []byte) (int, error) {
	n, err := fr.W.Write(buf)
	if err != nil {
		return n, err
	}

	if flusher, ok := fr.W.(http.Flusher); ok {
		flusher.Flush()
	}
	return n, err
}
<commit_msg>cmds: use flushCopy instrad of copychunks<commit_after>package http

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

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"

	cmds "github.com/jbenet/go-ipfs/commands"
	u "github.com/jbenet/go-ipfs/util"
)

var log = u.Logger("commands/http")

type Handler struct {
	ctx    cmds.Context
	root   *cmds.Command
	origin string
}

var ErrNotFound = errors.New("404 page not found")

const (
	streamHeader           = "X-Stream-Output"
	channelHeader          = "X-Chunked-Output"
	contentTypeHeader      = "Content-Type"
	contentLengthHeader    = "Content-Length"
	transferEncodingHeader = "Transfer-Encoding"
	applicationJson        = "application/json"
)

var mimeTypes = map[string]string{
	cmds.JSON: "application/json",
	cmds.XML:  "application/xml",
	cmds.Text: "text/plain",
}

func NewHandler(ctx cmds.Context, root *cmds.Command, origin string) *Handler {
	// allow whitelisted origins (so we can make API requests from the browser)
	if len(origin) > 0 {
		log.Info("Allowing API requests from origin: " + origin)
	}

	return &Handler{ctx, root, origin}
}

func (i Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// create a context.Context to pass into the commands.
	ctx, cancel := context.WithCancel(context.TODO())
	defer cancel()
	i.ctx.Context = ctx

	log.Debug("Incoming API request: ", r.URL)

	if len(i.origin) > 0 {
		w.Header().Set("Access-Control-Allow-Origin", i.origin)
	}
	w.Header().Set("Access-Control-Allow-Headers", "Content-Type")

	req, err := Parse(r, i.root)
	if err != nil {
		if err == ErrNotFound {
			w.WriteHeader(http.StatusNotFound)
		} else {
			w.WriteHeader(http.StatusBadRequest)
		}
		w.Write([]byte(err.Error()))
		return
	}
	req.SetContext(i.ctx)

	// call the command
	res := i.root.Call(req)

	// set the Content-Type based on res output
	if _, ok := res.Output().(io.Reader); ok {
		// we don't set the Content-Type for streams, so that browsers can MIME-sniff the type themselves
		// we set this header so clients have a way to know this is an output stream
		// (not marshalled command output)
		// TODO: set a specific Content-Type if the command response needs it to be a certain type
		w.Header().Set(streamHeader, "1")

	} else {
		enc, found, err := req.Option(cmds.EncShort).String()
		if err != nil || !found {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		mime := mimeTypes[enc]
		w.Header().Set(contentTypeHeader, mime)
	}

	// if response contains an error, write an HTTP error status code
	if e := res.Error(); e != nil {
		if e.Code == cmds.ErrClient {
			w.WriteHeader(http.StatusBadRequest)
		} else {
			w.WriteHeader(http.StatusInternalServerError)
		}
	}

	out, err := res.Reader()
	if err != nil {
		w.Header().Set(contentTypeHeader, "text/plain")
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(err.Error()))
		return
	}

	// if output is a channel and user requested streaming channels,
	// use chunk copier for the output
	_, isChan := res.Output().(chan interface{})
	streamChans, _, _ := req.Option("stream-channels").Bool()
	if isChan && streamChans {
		// w.WriteString(transferEncodingHeader + ": chunked\r\n")
		// w.Header().Set(channelHeader, "1")
		// w.WriteHeader(200)
		err = copyChunks(applicationJson, w, out)
		if err != nil {
			log.Error(err)
		}
		return
	}

	flushCopy(w, out)
}

// flushCopy Copies from an io.Reader to a http.ResponseWriter.
// Flushes chunks over HTTP stream as they are read (if supported by transport).
func flushCopy(w http.ResponseWriter, out io.Reader) error {
	if _, ok := w.(http.Flusher); !ok {
		return copyChunks("", w, out)
	}

	io.Copy(&flushResponse{w}, out)
	return nil
}

// Copies from an io.Reader to a http.ResponseWriter.
// Flushes chunks over HTTP stream as they are read (if supported by transport).
func copyChunks(contentType string, w http.ResponseWriter, out io.Reader) error {
	hijacker, ok := w.(http.Hijacker)
	if !ok {
		return errors.New("Could not create hijacker")
	}
	conn, writer, err := hijacker.Hijack()
	if err != nil {
		return err
	}
	defer conn.Close()

	writer.WriteString("HTTP/1.1 200 OK\r\n")
	if contentType != "" {
		writer.WriteString(contentTypeHeader + ": " + contentType + "\r\n")
	}
	writer.WriteString(transferEncodingHeader + ": chunked\r\n")
	writer.WriteString(channelHeader + ": 1\r\n\r\n")

	buf := make([]byte, 32*1024)

	for {
		n, err := out.Read(buf)

		if n > 0 {
			length := fmt.Sprintf("%x\r\n", n)
			writer.WriteString(length)

			_, err := writer.Write(buf[0:n])
			if err != nil {
				return err
			}

			writer.WriteString("\r\n")
			writer.Flush()
		}

		if err != nil && err != io.EOF {
			return err
		}
		if err == io.EOF {
			break
		}
	}

	writer.WriteString("0\r\n\r\n")
	writer.Flush()

	return nil
}

type flushResponse struct {
	W http.ResponseWriter
}

func (fr *flushResponse) Write(buf []byte) (int, error) {
	n, err := fr.W.Write(buf)
	if err != nil {
		return n, err
	}

	if flusher, ok := fr.W.(http.Flusher); ok {
		flusher.Flush()
	}
	return n, err
}
<|endoftext|>


<commit_before>package goboots

import (
	"compress/gzip"
	"io"
	"net/http"
	"strings"
)

type Filter func(in *In) bool

// gzip compression filter

type gzipRespWriter struct {
	io.Writer
	http.ResponseWriter
}

func (w gzipRespWriter) Write(b []byte) (int, error) {
	if w.Header().Get("Content-Type") == "" {
		// If no content type, apply sniffing algorithm to un-gzipped body.
		w.Header().Set("Content-Type", http.DetectContentType(b))
	}
	return w.Writer.Write(b)
}

// based on https://gist.github.com/the42/1956518
func CompressFilter(in *In) bool {
	if in.hijacked {
		return true
	}
	if !strings.Contains(in.R.Header.Get("Accept-Encoding"), "gzip") || !in.App.Config.GZipDynamic {
		return true
	}
	in.W.Header().Set("Content-Encoding", "gzip")
	gz := gzip.NewWriter(in.W)
	gzr := &gzipRespWriter{gz, in.W}
	in.W = gzr
	in.closers = append(in.closers, gz)
	return true
}

func ServedByProxyFilter(in *In) bool {
	if v := in.R.Header.Get("X-Forwarded-For"); v != "" {
		in.R.RemoteAddr = v
	}
	return true
}

func NoCacheFilter(in *In) bool {
	hh := in.W.Header()
	hh.Set("Cache-Control", "no-cache, no-store, must-revalidate")
	hh.Set("Pragma", "no-cache")
	hh.Set("Expires", "0")
	return true
}
<commit_msg>added X-Real-IP header check<commit_after>package goboots

import (
	"compress/gzip"
	"io"
	"net/http"
	"strings"
)

type Filter func(in *In) bool

// gzip compression filter

type gzipRespWriter struct {
	io.Writer
	http.ResponseWriter
}

func (w gzipRespWriter) Write(b []byte) (int, error) {
	if w.Header().Get("Content-Type") == "" {
		// If no content type, apply sniffing algorithm to un-gzipped body.
		w.Header().Set("Content-Type", http.DetectContentType(b))
	}
	return w.Writer.Write(b)
}

// based on https://gist.github.com/the42/1956518
func CompressFilter(in *In) bool {
	if in.hijacked {
		return true
	}
	if !strings.Contains(in.R.Header.Get("Accept-Encoding"), "gzip") || !in.App.Config.GZipDynamic {
		return true
	}
	in.W.Header().Set("Content-Encoding", "gzip")
	gz := gzip.NewWriter(in.W)
	gzr := &gzipRespWriter{gz, in.W}
	in.W = gzr
	in.closers = append(in.closers, gz)
	return true
}

func ServedByProxyFilter(in *In) bool {
	if v := in.R.Header.Get("X-Real-IP"); v != "" {
		in.R.RemoteAddr = v
	} else if v := in.R.Header.Get("X-Forwarded-For"); v != "" {
		in.R.RemoteAddr = v
	}
	return true
}

func NoCacheFilter(in *In) bool {
	hh := in.W.Header()
	hh.Set("Cache-Control", "no-cache, no-store, must-revalidate")
	hh.Set("Pragma", "no-cache")
	hh.Set("Expires", "0")
	return true
}
<|endoftext|>


<commit_before>package gnuplot

import (
	"github.com/yassu/gnuplot.go/conf"
	"testing"
	// "fmt"
	// "math"
)

// Plotter
func TestNewPlotter(t *testing.T) {
	NewPlotter()
}

func TestConfigure2(t *testing.T) {
	p := NewPlotter()
	p.Configure(conf.WithConf())
	if len(p.configures) != 1 {
		t.Errorf("fails in TestConfigure2")
	}
}

// // Function2d
// func TestFunction2dSplitNum(t *testing.T) {
// 	fun := new(Function2d)
// 	fun.Init()
// 	if fun.splitNum != 1000 {
// 		t.Errorf("fails in TestFunction2dSplitNum")
// 	}
// }
//
// func TestFunction2dUpdatePlotter(t *testing.T) {
// 	fun := new(Function2d)
// 	fun.Init()
//
// 	p := new(Plotter)
// 	p.Init()
// 	p.Configure("xMin", "12345")
// 	fun.UpdatePlotter(p)
// 	if fun.plotter.GetC("xMin") != "12345" {
// 		t.Errorf("fails in TestFunction2dUpdatePlotter")
// 	}
// }
//
// // func TestGetGnuData(t *testing.T){
// //     // without panic
// //     fun := new(Function2d)
// //     fun.Init()
// //
// //     f := func(x float64) float64 {
// //         return x * x
// //     }
// //     fun.SetF(f)
// //     fmt.Println(fun.getGnuData())
// // }
//
// // Curve2d
// func TestCurve2dSplitNum(t *testing.T) {
// 	c := new(Curve2d)
// 	c.Init()
// 	if c.splitNum != 100 {
// 		t.Errorf("fails in TestFunction2dSplitNum")
// 	}
// }
//
// func TestCurve2dUpdatePlotter(t *testing.T) {
// 	c := new(Curve2d)
// 	c.Init()
//
// 	p := new(Plotter)
// 	p.Init()
// 	p.Configure("tMin", "12345")
// 	c.UpdatePlotter(p)
// 	if c.plotter.GetC("tMin") != "12345" {
// 		t.Errorf("fails in TestCurve2dUpdatePlotter")
// 	}
// }
<commit_msg>test for function2d<commit_after>package gnuplot

import (
	"github.com/yassu/gnuplot.go/conf"
	"testing"
	// "fmt"
	// "math"
)

// Plotter
func TestNewPlotter(t *testing.T) {
	NewPlotter()
}

func TestConfigure2(t *testing.T) {
	p := NewPlotter()
	p.Configure(conf.WithConf())
	if len(p.configures) != 1 {
		t.Errorf("fails in TestConfigure2")
	}
}

// Function2d
func TestFunction2dSplitNum(t *testing.T) {
	fun := NewFunction2d()
	if fun.splitNum != 1000 {
		t.Errorf("fails in TestFunction2dSplitNum")
	}
}

func TestGetGnuData(t *testing.T) {
	fun := NewFunction2d()

	f := func(x float64) float64 {
		return x * x
	}
	fun.SetF(f)
	fun.getGnuData()
}

// // Curve2d
// func TestCurve2dSplitNum(t *testing.T) {
// 	c := new(Curve2d)
// 	c.Init()
// 	if c.splitNum != 100 {
// 		t.Errorf("fails in TestFunction2dSplitNum")
// 	}
// }
//
// func TestCurve2dUpdatePlotter(t *testing.T) {
// 	c := new(Curve2d)
// 	c.Init()
//
// 	p := new(Plotter)
// 	p.Init()
// 	p.Configure("tMin", "12345")
// 	c.UpdatePlotter(p)
// 	if c.plotter.GetC("tMin") != "12345" {
// 		t.Errorf("fails in TestCurve2dUpdatePlotter")
// 	}
// }
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2014 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package cpu6502

import (
	"io/ioutil"
	"testing"
)

type Ram []byte

func (r Ram) Read(v uint16) byte     { return r[v] }
func (r Ram) Write(v uint16, b byte) { r[v] = b }

// Download from https://github.com/Klaus2m5/6502_65C02_functional_tests/blob/master/bin_files/6502_functional_test.bin
// GPL, so not included here.
func TestFunctional(t *testing.T) {
	b, err := ioutil.ReadFile("6502_functional_test.bin")
	if err != nil {
		t.Fatal(err)
	}
	r := make(Ram, 0xffff+1)
	copy(r[:], b)
	c := New(r)
	c.L = make([]Log, 20)
	c.PC = 0x0400
	i := 0
	for !c.Halt {
		if c.PC == 0x3399 {
			break
		}
		pc := c.PC
		c.Step()
		if c.PC == pc {
			t.Log(c.StringLog())
			t.Log(c.String())
			t.Log(i, "instructions ran")
			t.Fatalf("repeated PC: 0x%04X", pc)
		} else if c.PC <= 0x1ff {
			t.Log(c.StringLog())
			t.Log(c.String())
			t.Fatalf("low PC: 0x%04X", c.PC)
		}
		i++
	}
}
<commit_msg>Fix test<commit_after>/*
 * Copyright (c) 2014 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package cpu6502

import (
	"io/ioutil"
	"testing"
)

type Ram []byte

func (r Ram) Read(v uint16) byte     { return r[v] }
func (r Ram) Write(v uint16, b byte) { r[v] = b }

// Download from https://github.com/Klaus2m5/6502_65C02_functional_tests/blob/master/bin_files/6502_functional_test.bin
// GPL, so not included here.
func TestFunctional(t *testing.T) {
	b, err := ioutil.ReadFile("6502_functional_test.bin")
	if err != nil {
		t.Fatal(err)
	}
	r := make(Ram, 0xffff+1)
	copy(r[:], b)
	c := New(r)
	c.L = make([]Log, 20)
	c.PC = 0x0400
	i := 0
	for {
		if c.PC == 0x3399 {
			break
		}
		pc := c.PC
		c.Step()
		if c.PC == pc {
			t.Log(c.StringLog())
			t.Log(c.String())
			t.Log(i, "instructions ran")
			t.Fatalf("repeated PC: 0x%04X", pc)
		} else if c.PC <= 0x1ff {
			t.Log(c.StringLog())
			t.Log(c.String())
			t.Fatalf("low PC: 0x%04X", c.PC)
		}
		i++
	}
}
<|endoftext|>


<commit_before>package socks

import (
	"net"

	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy/socks/protocol"
)

const (
	bufferSize = 2 * 1024
)

var udpAddress v2net.Address

func (server *SocksServer) ListenUDP(port uint16) error {
	addr := &net.UDPAddr{
		IP:   net.IP{0, 0, 0, 0},
		Port: int(port),
		Zone: "",
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Error("Socks failed to listen UDP on port %d: %v", port, err)
		return err
	}
	udpAddress = v2net.IPAddress([]byte{127, 0, 0, 1}, port)

	go server.AcceptPackets(conn)
	return nil
}

func (server *SocksServer) getUDPAddr() v2net.Address {
	return udpAddress
}

func (server *SocksServer) AcceptPackets(conn *net.UDPConn) error {
	for {
		buffer := make([]byte, bufferSize)
		nBytes, addr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			log.Error("Socks failed to read UDP packets: %v", err)
			return err
		}
		request, err := protocol.ReadUDPRequest(buffer[:nBytes])
		if err != nil {
			log.Error("Socks failed to parse UDP request: %v", err)
			return err
		}
		if request.Fragment != 0 {
			// TODO handle fragments
			continue
		}

		udpPacket := v2net.NewPacket(request.Destination(), request.Data, false)
		go server.handlePacket(conn, udpPacket, addr, request)
	}
}

func (server *SocksServer) handlePacket(conn *net.UDPConn, packet v2net.Packet, clientAddr *net.UDPAddr, request protocol.Socks5UDPRequest) {
	ray := server.vPoint.DispatchToOutbound(packet)
	close(ray.InboundInput())

	if data, ok := <-ray.InboundOutput(); ok {
		request.Data = data
		udpMessage := request.Bytes(nil)
		nBytes, err := conn.WriteToUDP(udpMessage, clientAddr)
    if err != nil {
      log.Error("Socks failed to write UDP message (%d bytes) to %s: %v", nBytes, clientAddr.String(), err)
    }
	}
}
<commit_msg>format code<commit_after>package socks

import (
	"net"

	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy/socks/protocol"
)

const (
	bufferSize = 2 * 1024
)

var udpAddress v2net.Address

func (server *SocksServer) ListenUDP(port uint16) error {
	addr := &net.UDPAddr{
		IP:   net.IP{0, 0, 0, 0},
		Port: int(port),
		Zone: "",
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Error("Socks failed to listen UDP on port %d: %v", port, err)
		return err
	}
	udpAddress = v2net.IPAddress([]byte{127, 0, 0, 1}, port)

	go server.AcceptPackets(conn)
	return nil
}

func (server *SocksServer) getUDPAddr() v2net.Address {
	return udpAddress
}

func (server *SocksServer) AcceptPackets(conn *net.UDPConn) error {
	for {
		buffer := make([]byte, bufferSize)
		nBytes, addr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			log.Error("Socks failed to read UDP packets: %v", err)
			return err
		}
		request, err := protocol.ReadUDPRequest(buffer[:nBytes])
		if err != nil {
			log.Error("Socks failed to parse UDP request: %v", err)
			return err
		}
		if request.Fragment != 0 {
			// TODO handle fragments
			continue
		}

		udpPacket := v2net.NewPacket(request.Destination(), request.Data, false)
		go server.handlePacket(conn, udpPacket, addr, request)
	}
}

func (server *SocksServer) handlePacket(conn *net.UDPConn, packet v2net.Packet, clientAddr *net.UDPAddr, request protocol.Socks5UDPRequest) {
	ray := server.vPoint.DispatchToOutbound(packet)
	close(ray.InboundInput())

	if data, ok := <-ray.InboundOutput(); ok {
		request.Data = data
		udpMessage := request.Bytes(nil)
		nBytes, err := conn.WriteToUDP(udpMessage, clientAddr)
		if err != nil {
			log.Error("Socks failed to write UDP message (%d bytes) to %s: %v", nBytes, clientAddr.String(), err)
		}
	}
}
<|endoftext|>


<commit_before>package unit

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"
	"syscall"

	systemdDbus "github.com/coreos/go-systemd/dbus"
	"github.com/guelfey/go.dbus"

	"github.com/coreos/coreinit/job"
	"github.com/coreos/coreinit/machine"
)

const (
	defaultSystemdRuntimePath = "/run/systemd/system/"
	defaultSystemdDbusPath    = "/org/freedesktop/systemd1/unit/"
)

type SystemdManager struct {
	Systemd  *systemdDbus.Conn
	Target   *SystemdTarget
	Machine  *machine.Machine
	unitPath string
	dbusPath string
}

func NewSystemdManager(machine *machine.Machine) *SystemdManager {
	systemd := systemdDbus.New()

	name := "coreinit-" + machine.BootId + ".target"
	target := NewSystemdTarget(name)

	mgr := &SystemdManager{systemd, target, machine, defaultSystemdRuntimePath, defaultSystemdDbusPath}

	mgr.writeUnit(target.Name(), "")

	return mgr
}

func (m *SystemdManager) getUnitByName(name string) (*SystemdUnit, error) {
	var unit SystemdUnit
	if strings.HasSuffix(name, ".service") {
		unit = NewSystemdService(m, name)
	} else if strings.HasSuffix(name, ".socket") {
		unit = NewSystemdSocket(m, name)
	} else {
		panic("WAT")
	}

	return &unit, nil
}

func (m *SystemdManager) getUnitsByTarget(target *SystemdTarget) []SystemdUnit {
	object := m.getDbusPath(target.Name())
	info, err := m.Systemd.GetUnitInfo(object)

	if err != nil {
		panic(err)
	}

	names := info["Wants"].Value().([]string)

	var units []SystemdUnit
	for _, name := range names {
		unit, err := m.getUnitByName(name)
		if err == nil {
			units = append(units, *unit)
		} else {
			log.Printf("Unit %s seems to exist, yet unable to get corresponding SystemdUnit object", name)
		}
	}

	return units
}

func (m *SystemdManager) GetJobs() map[string]job.Job {
	units := m.getUnitsByTarget(m.Target)
	jobs := make(map[string]job.Job, len(units))
	for _, u := range units {
		state := m.getJobStateFromUnit(&u)
		j, _ := job.NewJob(u.Name(), state, nil)
		jobs[j.Name] = *j
	}

	return jobs
}

func (m *SystemdManager) getJobStateFromUnit(u *SystemdUnit) *job.JobState {
	loadState, activeState, subState, sockets, err := (*u).State()
	if err != nil {
		log.Printf("Failed to get state for unit %s", (*u).Name())
		return nil
	} else {
		return job.NewJobState(loadState, activeState, subState, sockets, m.Machine)
	}
}

func (m *SystemdManager) GetJobState(j *job.Job) *job.JobState {
	unit, err := m.getUnitByName(j.Name)
	if err != nil {
		log.Printf("No local unit corresponding to job %s", j.Name)
		return nil
	}

	return m.getJobStateFromUnit(unit)
}

func (m *SystemdManager) StartJob(job *job.Job) {
	//This is probably not the right place to force the service to be
	// WantedBy our systemd target
	job.Payload.Value += "\r\n\r\n[Install]\r\nWantedBy=" + m.Target.Name()

	m.writeUnit(job.Name, job.Payload.Value)
	m.startUnit(job.Name)
}

func (m *SystemdManager) StopJob(job *job.Job) {
	m.stopUnit(job.Name)
	m.removeUnit(job.Name)
}

func (m *SystemdManager) getUnitStates(name string) (string, string, string, error) {
	info, err := m.Systemd.GetUnitInfo(m.getDbusPath(name))

	if err != nil {
		return "", "", "", err
	} else {
		loadState := info["LoadState"].Value().(string)
		activeState := info["ActiveState"].Value().(string)
		subState := info["SubState"].Value().(string)
		return loadState, activeState, subState, nil
	}
}

func (m *SystemdManager) startUnit(name string) {
	log.Println("Starting systemd unit", name)

	files := []string{name}
	m.Systemd.EnableUnitFiles(files, true, false)

	m.Systemd.StartUnit(name, "replace")
}

func (m *SystemdManager) stopUnit(name string) {
	log.Println("Stopping systemd unit", name)

	m.Systemd.StopUnit(name, "replace")

	// go-systemd does not yet have this implemented
	//files := []string{name}
	//Systemd.DisableUnitFiles(files, true, false)
}

func (m *SystemdManager) removeUnit(name string) {
	log.Printf("Unlinking systemd unit %s from target %s", name, m.Target.Name())
	link := m.getLocalPath(path.Join(m.Target.Name()+".wants", name))
	syscall.Unlink(link)
}

func (m *SystemdManager) readUnit(name string) (string, error) {
	path := m.getLocalPath(name)
	contents, err := ioutil.ReadFile(path)
	if err == nil {
		return string(contents), nil
	} else {
		return "", errors.New(fmt.Sprintf("No unit file at local path %s", path))
	}
}

func (m *SystemdManager) writeUnit(name string, contents string) error {
	log.Println("Writing systemd unit file", name)

	path := path.Join(m.unitPath, name)
	file, err := os.Create(path)
	defer file.Close()

	if err != nil {
		return err
	}

	file.Write([]byte(contents))
	return nil
}

func (m *SystemdManager) getDbusPath(name string) dbus.ObjectPath {
	path := path.Join(m.dbusPath, name)
	path = serializeDbusPath(path)
	return dbus.ObjectPath(path)
}

func (m *SystemdManager) getLocalPath(name string) string {
	return path.Join(m.unitPath, name)
}
<commit_msg>fix(agent): Unlink unit files when jobs are stopped<commit_after>package unit

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"
	"syscall"

	systemdDbus "github.com/coreos/go-systemd/dbus"
	"github.com/guelfey/go.dbus"

	"github.com/coreos/coreinit/job"
	"github.com/coreos/coreinit/machine"
)

const (
	defaultSystemdRuntimePath = "/run/systemd/system/"
	defaultSystemdDbusPath    = "/org/freedesktop/systemd1/unit/"
)

type SystemdManager struct {
	Systemd  *systemdDbus.Conn
	Target   *SystemdTarget
	Machine  *machine.Machine
	unitPath string
	dbusPath string
}

func NewSystemdManager(machine *machine.Machine) *SystemdManager {
	systemd := systemdDbus.New()

	name := "coreinit-" + machine.BootId + ".target"
	target := NewSystemdTarget(name)

	mgr := &SystemdManager{systemd, target, machine, defaultSystemdRuntimePath, defaultSystemdDbusPath}

	mgr.writeUnit(target.Name(), "")

	return mgr
}

func (m *SystemdManager) getUnitByName(name string) (*SystemdUnit, error) {
	var unit SystemdUnit
	if strings.HasSuffix(name, ".service") {
		unit = NewSystemdService(m, name)
	} else if strings.HasSuffix(name, ".socket") {
		unit = NewSystemdSocket(m, name)
	} else {
		panic("WAT")
	}

	return &unit, nil
}

func (m *SystemdManager) getUnitsByTarget(target *SystemdTarget) []SystemdUnit {
	object := m.getDbusPath(target.Name())
	info, err := m.Systemd.GetUnitInfo(object)

	if err != nil {
		panic(err)
	}

	names := info["Wants"].Value().([]string)

	var units []SystemdUnit
	for _, name := range names {
		unit, err := m.getUnitByName(name)
		if err == nil {
			units = append(units, *unit)
		} else {
			log.Printf("Unit %s seems to exist, yet unable to get corresponding SystemdUnit object", name)
		}
	}

	return units
}

func (m *SystemdManager) GetJobs() map[string]job.Job {
	units := m.getUnitsByTarget(m.Target)
	jobs := make(map[string]job.Job, len(units))
	for _, u := range units {
		state := m.getJobStateFromUnit(&u)
		j, _ := job.NewJob(u.Name(), state, nil)
		jobs[j.Name] = *j
	}

	return jobs
}

func (m *SystemdManager) getJobStateFromUnit(u *SystemdUnit) *job.JobState {
	loadState, activeState, subState, sockets, err := (*u).State()
	if err != nil {
		log.Printf("Failed to get state for unit %s", (*u).Name())
		return nil
	} else {
		return job.NewJobState(loadState, activeState, subState, sockets, m.Machine)
	}
}

func (m *SystemdManager) GetJobState(j *job.Job) *job.JobState {
	unit, err := m.getUnitByName(j.Name)
	if err != nil {
		log.Printf("No local unit corresponding to job %s", j.Name)
		return nil
	}

	return m.getJobStateFromUnit(unit)
}

func (m *SystemdManager) StartJob(job *job.Job) {
	//This is probably not the right place to force the service to be
	// WantedBy our systemd target
	job.Payload.Value += "\r\n\r\n[Install]\r\nWantedBy=" + m.Target.Name()

	m.writeUnit(job.Name, job.Payload.Value)
	m.startUnit(job.Name)
}

func (m *SystemdManager) StopJob(job *job.Job) {
	m.stopUnit(job.Name)
	m.removeUnit(job.Name)
}

func (m *SystemdManager) getUnitStates(name string) (string, string, string, error) {
	info, err := m.Systemd.GetUnitInfo(m.getDbusPath(name))

	if err != nil {
		return "", "", "", err
	} else {
		loadState := info["LoadState"].Value().(string)
		activeState := info["ActiveState"].Value().(string)
		subState := info["SubState"].Value().(string)
		return loadState, activeState, subState, nil
	}
}

func (m *SystemdManager) startUnit(name string) {
	log.Println("Starting systemd unit", name)

	files := []string{name}
	m.Systemd.EnableUnitFiles(files, true, false)

	m.Systemd.StartUnit(name, "replace")
}

func (m *SystemdManager) stopUnit(name string) {
	log.Println("Stopping systemd unit", name)

	m.Systemd.StopUnit(name, "replace")

	// go-systemd does not yet have this implemented
	//files := []string{name}
	//Systemd.DisableUnitFiles(files, true, false)
}

func (m *SystemdManager) removeUnit(name string) {
	log.Printf("Unlinking systemd unit %s from target %s", name, m.Target.Name())
	link := m.getLocalPath(path.Join(m.Target.Name()+".wants", name))
	syscall.Unlink(link)

	file := m.getLocalPath(name)
	log.Printf("Removing systemd unit file %s", file)
	syscall.Unlink(file)
}

func (m *SystemdManager) readUnit(name string) (string, error) {
	path := m.getLocalPath(name)
	contents, err := ioutil.ReadFile(path)
	if err == nil {
		return string(contents), nil
	} else {
		return "", errors.New(fmt.Sprintf("No unit file at local path %s", path))
	}
}

func (m *SystemdManager) writeUnit(name string, contents string) error {
	log.Println("Writing systemd unit file", name)

	path := path.Join(m.unitPath, name)
	file, err := os.Create(path)
	defer file.Close()

	if err != nil {
		return err
	}

	file.Write([]byte(contents))
	return nil
}

func (m *SystemdManager) getDbusPath(name string) dbus.ObjectPath {
	path := path.Join(m.dbusPath, name)
	path = serializeDbusPath(path)
	return dbus.ObjectPath(path)
}

func (m *SystemdManager) getLocalPath(name string) string {
	return path.Join(m.unitPath, name)
}
<|endoftext|>


<commit_before>package handler

import (
	"github.com/charles-l/gitamite"

	"github.com/labstack/echo"

	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
)

func CreateRepo(c echo.Context) error {
	blob, err := ioutil.ReadAll(c.Request().Body)
	if err != nil {
		return err
	}

	var a gitamite.AuthRequest
	err = json.Unmarshal(blob, &a)
	if err != nil {
		return err
	}

	if len(a.Signature) == 0 || a.Data == nil {
		return fmt.Errorf("Need data and signature")
	}

	if err := a.VerifyRequest(); err != nil {
		return fmt.Errorf("Invalid signature")
	}

	log.Printf("creating new repo: %s", a.Data.(map[string]interface{})["Name"])
	return nil
}
<commit_msg>create repo page actually creates repos now<commit_after>package handler

import (
	"github.com/charles-l/gitamite"
	"github.com/charles-l/gitamite/server/context"
	"github.com/libgit2/git2go"

	"github.com/labstack/echo"

	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"path"
)

func CreateRepo(c echo.Context) error {
	blob, err := ioutil.ReadAll(c.Request().Body)
	if err != nil {
		return err
	}

	var a gitamite.AuthRequest
	err = json.Unmarshal(blob, &a)
	if err != nil {
		return err
	}

	if len(a.Signature) == 0 || a.Data == nil {
		return fmt.Errorf("Need data and signature")
	}

	if err := a.VerifyRequest(); err != nil {
		return fmt.Errorf("Invalid signature")
	}

	m := a.Data.(map[string]interface{})
	var name string
	for k, v := range m {
		switch vv := v.(type) {
		case string:
			if k == "Name" {
				name = vv
			}
		}
	}
	if name == "" {
		return fmt.Errorf("need a repository name")
	}

	newRepoPath := path.Join(gitamite.GlobalConfig.RepoDir, name)
	log.Printf("creating new repo: %s", newRepoPath)

	repo, err := git.InitRepository(newRepoPath, true)
	if err != nil {
		return err
	}
	c.(*server.Context).Repos[name] = &gitamite.Repo{name, newRepoPath, "", repo}
	return nil
}
<|endoftext|>


<commit_before>package rabbus

import (
	"context"
	"errors"
	"sync"
	"time"

	amqpwrap "github.com/rafaeljesus/rabbus/internal/amqp"
	"github.com/rafaeljesus/retry-go"
	"github.com/sony/gobreaker"
	"github.com/streadway/amqp"
)

const (
	// Transient means higher throughput but messages will not be restored on broker restart.
	Transient uint8 = 1
	// Persistent messages will be restored to durable queues and lost on non-durable queues during server restart.
	Persistent uint8 = 2
	// ContentTypeJSON define json content type
	ContentTypeJSON = "application/json"
	// ContentTypePlain define plain text content type
	ContentTypePlain = "plain/text"

	contentEncoding = "UTF-8"
)

type (
	// Option represents an option you can pass to New.
	// See the documentation for the individual options.
	Option func(*Rabbus) error
	// OnStateChangeFunc is the callback function when circuit breaker state changes.
	OnStateChangeFunc func(name, from, to string)

	// Message carries fields for sending messages.
	Message struct {
		// Exchange the exchange name.
		Exchange string
		// Kind the exchange type.
		Kind string
		// Key the routing key name.
		Key string
		// Payload the message payload.
		Payload []byte
		// DeliveryMode indicates if the is Persistent or Transient.
		DeliveryMode uint8
		// ContentType the message content-type.
		ContentType string
		// Headers the message application headers
		Headers map[string]interface{}
	}

	// ListenConfig carries fields for listening messages.
	ListenConfig struct {
		// Exchange the exchange name.
		Exchange string
		// Kind the exchange type.
		Kind string
		// Key the routing key name.
		Key string
		// PassiveExchange determines a passive exchange connection it uses
		// amqp's ExchangeDeclarePassive instead the default ExchangeDeclare
		PassiveExchange bool
		// Queue the queue name
		Queue string
	}

	// Delivery wraps amqp.Delivery struct
	Delivery struct {
		amqp.Delivery
	}

	// Rabbus interpret (implement) Rabbus interface definition
	Rabbus struct {
		Amqp
		mu         sync.RWMutex
		breaker    *gobreaker.CircuitBreaker
		emit       chan Message
		emitErr    chan error
		emitOk     chan struct{}
		reconn     chan struct{}
		exDeclared map[string]struct{}
		config
	}

	// Amqp expose a interface for interacting with amqp broker
	Amqp interface {
		// Publish wraps amqp.Publish method
		Publish(exchange, key string, opts amqp.Publishing) error
		// CreateConsumer creates a amqp consumer
		CreateConsumer(exchange, key, kind, queue string, durable bool) (<-chan amqp.Delivery, error)
		// WithExchange creates a amqp exchange
		WithExchange(exchange, kind string, durable bool) error
		// WithQos wrapper over amqp.Qos method
		WithQos(count, size int, global bool) error
		// NotifyClose wrapper over notifyClose method
		NotifyClose(c chan *amqp.Error) chan *amqp.Error
		// Close closes the running amqp connection and channel
		Close() error
	}

	config struct {
		dsn                string
		durable, passiveex bool
		retrycfg
		breaker
		qos
	}

	retrycfg struct {
		attempts              int
		sleep, reconnectSleep time.Duration
	}

	breaker struct {
		interval, timeout time.Duration
		threshold         uint32
		onStateChange     OnStateChangeFunc
	}

	qos struct {
		prefetchCount, prefetchSize int
		global                      bool
	}
)

func (lc ListenConfig) validate() error {
	if lc.Exchange == "" {
		return ErrMissingExchange
	}

	if lc.Kind == "" {
		return ErrMissingKind
	}

	if lc.Queue == "" {
		return ErrMissingQueue
	}

	return nil
}

// New returns a new Rabbus configured with the
// variables from the config parameter, or returning an non-nil err
// if an error occurred while creating connection and channel.
func New(dsn string, options ...Option) (*Rabbus, error) {
	r := &Rabbus{
		emit:       make(chan Message),
		emitErr:    make(chan error),
		emitOk:     make(chan struct{}),
		reconn:     make(chan struct{}, 10),
		exDeclared: make(map[string]struct{}),
	}

	for _, o := range options {
		if err := o(r); err != nil {
			return nil, err
		}
	}

	if r.Amqp == nil {
		amqpWrapper, err := amqpwrap.New(dsn, r.config.passiveex)
		if err != nil {
			return nil, err
		}
		r.Amqp = amqpWrapper
	}

	if err := r.WithQos(
		r.config.qos.prefetchCount,
		r.config.qos.prefetchSize,
		r.config.qos.global,
	); err != nil {
		return nil, err
	}

	r.config.dsn = dsn
	r.breaker = gobreaker.NewCircuitBreaker(newBreakerSettings(r.config))

	return r, nil
}

// Run starts rabbus channels for emiting and listeting for amqp connection close
// returns ctx error in case of any.
func (r *Rabbus) Run(ctx context.Context) error {
	for {
		select {
		case m, ok := <-r.emit:
			if ok {
				r.produce(m)
			}
		case err, ok := <-r.NotifyClose(make(chan *amqp.Error)):
			if ok {
				r.handleAmqpClose(err)
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// EmitAsync emits a message to RabbitMQ, but does not wait for the response from broker.
func (r *Rabbus) EmitAsync() chan<- Message { return r.emit }

// EmitErr returns an error if encoding payload fails, or if after circuit breaker is open or retries attempts exceed.
func (r *Rabbus) EmitErr() <-chan error { return r.emitErr }

// EmitOk returns true when the message was sent.
func (r *Rabbus) EmitOk() <-chan struct{} { return r.emitOk }

// Listen to a message from RabbitMQ, returns
// an error if exchange, queue name and function handler not passed or if an error occurred while creating
// amqp consumer.
func (r *Rabbus) Listen(c ListenConfig) (chan ConsumerMessage, error) {
	if err := c.validate(); err != nil {
		return nil, err
	}

	msgs, err := r.CreateConsumer(c.Exchange, c.Key, c.Kind, c.Queue, r.config.durable)
	if err != nil {
		return nil, err
	}

	r.mu.Lock()
	r.exDeclared[c.Exchange] = struct{}{}
	r.mu.Unlock()

	messages := make(chan ConsumerMessage, 256)
	go r.wrapMessage(c, msgs, messages)
	go r.listenReconn(c, messages)

	return messages, nil
}

// Close channels and attempt to close channel and connection.
func (r *Rabbus) Close() error {
	close(r.emit)
	close(r.emitOk)
	close(r.emitErr)
	close(r.reconn)

	return r.Amqp.Close()
}

func (r *Rabbus) produce(m Message) {
	if _, ok := r.exDeclared[m.Exchange]; !ok {
		if err := r.WithExchange(m.Exchange, m.Kind, r.config.durable); err != nil {
			r.emitErr <- err
			return
		}
		r.exDeclared[m.Exchange] = struct{}{}
	}

	if m.ContentType == "" {
		m.ContentType = ContentTypeJSON
	}

	if m.DeliveryMode == 0 {
		m.DeliveryMode = Persistent
	}

	opts := amqp.Publishing{
		Headers:         amqp.Table(m.Headers),
		ContentType:     m.ContentType,
		ContentEncoding: contentEncoding,
		DeliveryMode:    m.DeliveryMode,
		Timestamp:       time.Now(),
		Body:            m.Payload,
	}

	if _, err := r.breaker.Execute(func() (interface{}, error) {
		return nil, retry.Do(func() error {
			return r.Publish(m.Exchange, m.Key, opts)
		}, r.config.retrycfg.attempts, r.config.retrycfg.sleep)
	}); err != nil {
		r.emitErr <- err
		return
	}

	r.emitOk <- struct{}{}
}

// Durable indicates of the queue will survive broker restarts. Default to true.
func Durable(durable bool) Option {
	return func(r *Rabbus) error {
		r.config.durable = durable
		return nil
	}
}

// PassiveExchange forces passive connection with all exchanges using
// amqp's ExchangeDeclarePassive instead the default ExchangeDeclare
func PassiveExchange(passiveex bool) Option {
	return func(r *Rabbus) error {
		r.config.passiveex = passiveex
		return nil
	}
}

// PrefetchCount limit the number of unacknowledged messages.
func PrefetchCount(count int) Option {
	return func(r *Rabbus) error {
		r.config.qos.prefetchCount = count
		return nil
	}
}

// PrefetchSize when greater than zero, the server will try to keep at least
// that many bytes of deliveries flushed to the network before receiving
// acknowledgments from the consumers.
func PrefetchSize(size int) Option {
	return func(r *Rabbus) error {
		r.config.qos.prefetchSize = size
		return nil
	}
}

// QosGlobal when global is true, these Qos settings apply to all existing and future
// consumers on all channels on the same connection. When false, the Channel.Qos
// settings will apply to all existing and future consumers on this channel.
// RabbitMQ does not implement the global flag.
func QosGlobal(global bool) Option {
	return func(r *Rabbus) error {
		r.config.qos.global = global
		return nil
	}
}

// Attempts is the max number of retries on broker outages.
func Attempts(attempts int) Option {
	return func(r *Rabbus) error {
		r.config.retrycfg.attempts = attempts
		return nil
	}
}

// Sleep is the sleep time of the retry mechanism.
func Sleep(sleep time.Duration) Option {
	return func(r *Rabbus) error {
		if sleep == 0 {
			r.config.retrycfg.reconnectSleep = time.Second * 10
		}
		r.config.retrycfg.sleep = sleep
		return nil
	}
}

// BreakerInterval is the cyclic period of the closed state for CircuitBreaker to clear the internal counts,
// If Interval is 0, CircuitBreaker doesn't clear the internal counts during the closed state.
func BreakerInterval(interval time.Duration) Option {
	return func(r *Rabbus) error {
		r.config.breaker.interval = interval
		return nil
	}
}

// BreakerTimeout is the period of the open state, after which the state of CircuitBreaker becomes half-open.
// If Timeout is 0, the timeout value of CircuitBreaker is set to 60 seconds.
func BreakerTimeout(timeout time.Duration) Option {
	return func(r *Rabbus) error {
		r.config.breaker.timeout = timeout
		return nil
	}
}

// Threshold when a threshold of failures has been reached, future calls to the broker will not run.
// During this state, the circuit breaker will periodically allow the calls to run and, if it is successful,
// will start running the function again. Default value is 5.
func Threshold(threshold uint32) Option {
	return func(r *Rabbus) error {
		if threshold == 0 {
			threshold = 5
		}
		r.config.breaker.threshold = threshold
		return nil
	}
}

// OnStateChange is called whenever the state of CircuitBreaker changes.
func OnStateChange(fn OnStateChangeFunc) Option {
	return func(r *Rabbus) error {
		r.config.breaker.onStateChange = fn
		return nil
	}
}

// AmqpProvider expose a interface for interacting with amqp broker
func AmqpProvider(provider Amqp) Option {
	return func(r *Rabbus) error {
		if provider != nil {
			r.Amqp = provider
			return nil
		}
		return errors.New("unexpected amqp provider")
	}
}

func (r *Rabbus) wrapMessage(c ListenConfig, sourceChan <-chan amqp.Delivery, targetChan chan ConsumerMessage) {
	for m := range sourceChan {
		targetChan <- newConsumerMessage(m)
	}
}

func (r *Rabbus) handleAmqpClose(err error) {
	for {
		time.Sleep(time.Second)
		aw, err := amqpwrap.New(r.config.dsn, r.config.passiveex)
		if err != nil {
			continue
		}

		r.mu.Lock()
		r.Amqp = aw
		r.mu.Unlock()
		r.reconn <- struct{}{}
		break
	}
}

func (r *Rabbus) listenReconn(c ListenConfig, messages chan ConsumerMessage) {
	for range r.reconn {
		msgs, err := r.CreateConsumer(c.Exchange, c.Key, c.Kind, c.Queue, r.config.durable)
		if err != nil {
			r.Close()
			continue
		}

		go r.wrapMessage(c, msgs, messages)
		go r.listenReconn(c, messages)
		break
	}
}

func newBreakerSettings(c config) gobreaker.Settings {
	s := gobreaker.Settings{}
	s.Name = "rabbus-circuit-breaker"
	s.Interval = c.breaker.interval
	s.Timeout = c.breaker.timeout
	s.ReadyToTrip = func(counts gobreaker.Counts) bool {
		return counts.ConsecutiveFailures > c.breaker.threshold
	}
	if c.breaker.onStateChange != nil {
		s.OnStateChange = func(name string, from gobreaker.State, to gobreaker.State) {
			c.breaker.onStateChange(name, from.String(), to.String())
		}
	}
	return s
}
<commit_msg>Making docs more professional<commit_after>package rabbus

import (
	"context"
	"errors"
	"sync"
	"time"

	amqpwrap "github.com/rafaeljesus/rabbus/internal/amqp"
	"github.com/rafaeljesus/retry-go"
	"github.com/sony/gobreaker"
	"github.com/streadway/amqp"
)

const (
	// Transient means higher throughput but messages will not be restored on broker restart.
	Transient uint8 = 1
	// Persistent messages will be restored to durable queues and lost on non-durable queues during server restart.
	Persistent uint8 = 2
	// ContentTypeJSON define json content type
	ContentTypeJSON = "application/json"
	// ContentTypePlain define plain text content type
	ContentTypePlain = "plain/text"

	contentEncoding = "UTF-8"
)

type (
	// Option represents an option you can pass to New.
	// See the documentation for the individual options.
	Option func(*Rabbus) error
	// OnStateChangeFunc is the callback function when circuit breaker state changes.
	OnStateChangeFunc func(name, from, to string)

	// Message carries fields for sending messages.
	Message struct {
		// Exchange the exchange name.
		Exchange string
		// Kind the exchange type.
		Kind string
		// Key the routing key name.
		Key string
		// Payload the message payload.
		Payload []byte
		// DeliveryMode indicates if the is Persistent or Transient.
		DeliveryMode uint8
		// ContentType the message content-type.
		ContentType string
		// Headers the message application headers
		Headers map[string]interface{}
	}

	// ListenConfig carries fields for listening messages.
	ListenConfig struct {
		// Exchange the exchange name.
		Exchange string
		// Kind the exchange type.
		Kind string
		// Key the routing key name.
		Key string
		// PassiveExchange determines a passive exchange connection it uses
		// amqp's ExchangeDeclarePassive instead the default ExchangeDeclare
		PassiveExchange bool
		// Queue the queue name
		Queue string
	}

	// Delivery wraps amqp.Delivery struct
	Delivery struct {
		amqp.Delivery
	}

	// Rabbus interpret (implement) Rabbus interface definition
	Rabbus struct {
		Amqp
		mu         sync.RWMutex
		breaker    *gobreaker.CircuitBreaker
		emit       chan Message
		emitErr    chan error
		emitOk     chan struct{}
		reconn     chan struct{}
		exDeclared map[string]struct{}
		config
	}

	// Amqp expose a interface for interacting with amqp broker
	Amqp interface {
		// Publish wraps amqp.Publish method
		Publish(exchange, key string, opts amqp.Publishing) error
		// CreateConsumer creates a amqp consumer
		CreateConsumer(exchange, key, kind, queue string, durable bool) (<-chan amqp.Delivery, error)
		// WithExchange creates a amqp exchange
		WithExchange(exchange, kind string, durable bool) error
		// WithQos wrapper over amqp.Qos method
		WithQos(count, size int, global bool) error
		// NotifyClose wrapper over notifyClose method
		NotifyClose(c chan *amqp.Error) chan *amqp.Error
		// Close closes the running amqp connection and channel
		Close() error
	}

	config struct {
		dsn                string
		durable, passiveex bool
		retrycfg
		breaker
		qos
	}

	retrycfg struct {
		attempts              int
		sleep, reconnectSleep time.Duration
	}

	breaker struct {
		interval, timeout time.Duration
		threshold         uint32
		onStateChange     OnStateChangeFunc
	}

	qos struct {
		prefetchCount, prefetchSize int
		global                      bool
	}
)

func (lc ListenConfig) validate() error {
	if lc.Exchange == "" {
		return ErrMissingExchange
	}

	if lc.Kind == "" {
		return ErrMissingKind
	}

	if lc.Queue == "" {
		return ErrMissingQueue
	}

	return nil
}

// New returns a new Rabbus configured with the
// variables from the config parameter, or returning an non-nil err
// if an error occurred while creating connection and channel.
func New(dsn string, options ...Option) (*Rabbus, error) {
	r := &Rabbus{
		emit:       make(chan Message),
		emitErr:    make(chan error),
		emitOk:     make(chan struct{}),
		reconn:     make(chan struct{}, 10),
		exDeclared: make(map[string]struct{}),
	}

	for _, o := range options {
		if err := o(r); err != nil {
			return nil, err
		}
	}

	if r.Amqp == nil {
		amqpWrapper, err := amqpwrap.New(dsn, r.config.passiveex)
		if err != nil {
			return nil, err
		}
		r.Amqp = amqpWrapper
	}

	if err := r.WithQos(
		r.config.qos.prefetchCount,
		r.config.qos.prefetchSize,
		r.config.qos.global,
	); err != nil {
		return nil, err
	}

	r.config.dsn = dsn
	r.breaker = gobreaker.NewCircuitBreaker(newBreakerSettings(r.config))

	return r, nil
}

// Run starts rabbus channels for emitting and listening for amqp connection close
// returns ctx error in case of any.
func (r *Rabbus) Run(ctx context.Context) error {
	for {
		select {
		case m, ok := <-r.emit:
			if ok {
				r.produce(m)
			}
		case err, ok := <-r.NotifyClose(make(chan *amqp.Error)):
			if ok {
				r.handleAmqpClose(err)
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// EmitAsync emits a message to RabbitMQ, but does not wait for the response from broker.
func (r *Rabbus) EmitAsync() chan<- Message { return r.emit }

// EmitErr returns an error if encoding payload fails, or if after circuit breaker is open or retries attempts exceed.
func (r *Rabbus) EmitErr() <-chan error { return r.emitErr }

// EmitOk returns true when the message was sent.
func (r *Rabbus) EmitOk() <-chan struct{} { return r.emitOk }

// Listen to a message from RabbitMQ, returns
// an error if exchange, queue name and function handler not passed or if an error occurred while creating
// amqp consumer.
func (r *Rabbus) Listen(c ListenConfig) (chan ConsumerMessage, error) {
	if err := c.validate(); err != nil {
		return nil, err
	}

	msgs, err := r.CreateConsumer(c.Exchange, c.Key, c.Kind, c.Queue, r.config.durable)
	if err != nil {
		return nil, err
	}

	r.mu.Lock()
	r.exDeclared[c.Exchange] = struct{}{}
	r.mu.Unlock()

	messages := make(chan ConsumerMessage, 256)
	go r.wrapMessage(c, msgs, messages)
	go r.listenReconn(c, messages)

	return messages, nil
}

// Close channels and attempt to close channel and connection.
func (r *Rabbus) Close() error {
	close(r.emit)
	close(r.emitOk)
	close(r.emitErr)
	close(r.reconn)

	return r.Amqp.Close()
}

func (r *Rabbus) produce(m Message) {
	if _, ok := r.exDeclared[m.Exchange]; !ok {
		if err := r.WithExchange(m.Exchange, m.Kind, r.config.durable); err != nil {
			r.emitErr <- err
			return
		}
		r.exDeclared[m.Exchange] = struct{}{}
	}

	if m.ContentType == "" {
		m.ContentType = ContentTypeJSON
	}

	if m.DeliveryMode == 0 {
		m.DeliveryMode = Persistent
	}

	opts := amqp.Publishing{
		Headers:         amqp.Table(m.Headers),
		ContentType:     m.ContentType,
		ContentEncoding: contentEncoding,
		DeliveryMode:    m.DeliveryMode,
		Timestamp:       time.Now(),
		Body:            m.Payload,
	}

	if _, err := r.breaker.Execute(func() (interface{}, error) {
		return nil, retry.Do(func() error {
			return r.Publish(m.Exchange, m.Key, opts)
		}, r.config.retrycfg.attempts, r.config.retrycfg.sleep)
	}); err != nil {
		r.emitErr <- err
		return
	}

	r.emitOk <- struct{}{}
}

// Durable indicates of the queue will survive broker restarts. Default to true.
func Durable(durable bool) Option {
	return func(r *Rabbus) error {
		r.config.durable = durable
		return nil
	}
}

// PassiveExchange forces passive connection with all exchanges using
// amqp's ExchangeDeclarePassive instead the default ExchangeDeclare
func PassiveExchange(passiveex bool) Option {
	return func(r *Rabbus) error {
		r.config.passiveex = passiveex
		return nil
	}
}

// PrefetchCount limit the number of unacknowledged messages.
func PrefetchCount(count int) Option {
	return func(r *Rabbus) error {
		r.config.qos.prefetchCount = count
		return nil
	}
}

// PrefetchSize when greater than zero, the server will try to keep at least
// that many bytes of deliveries flushed to the network before receiving
// acknowledgments from the consumers.
func PrefetchSize(size int) Option {
	return func(r *Rabbus) error {
		r.config.qos.prefetchSize = size
		return nil
	}
}

// QosGlobal when global is true, these Qos settings apply to all existing and future
// consumers on all channels on the same connection. When false, the Channel.Qos
// settings will apply to all existing and future consumers on this channel.
// RabbitMQ does not implement the global flag.
func QosGlobal(global bool) Option {
	return func(r *Rabbus) error {
		r.config.qos.global = global
		return nil
	}
}

// Attempts is the max number of retries on broker outages.
func Attempts(attempts int) Option {
	return func(r *Rabbus) error {
		r.config.retrycfg.attempts = attempts
		return nil
	}
}

// Sleep is the sleep time of the retry mechanism.
func Sleep(sleep time.Duration) Option {
	return func(r *Rabbus) error {
		if sleep == 0 {
			r.config.retrycfg.reconnectSleep = time.Second * 10
		}
		r.config.retrycfg.sleep = sleep
		return nil
	}
}

// BreakerInterval is the cyclic period of the closed state for CircuitBreaker to clear the internal counts,
// If Interval is 0, CircuitBreaker doesn't clear the internal counts during the closed state.
func BreakerInterval(interval time.Duration) Option {
	return func(r *Rabbus) error {
		r.config.breaker.interval = interval
		return nil
	}
}

// BreakerTimeout is the period of the open state, after which the state of CircuitBreaker becomes half-open.
// If Timeout is 0, the timeout value of CircuitBreaker is set to 60 seconds.
func BreakerTimeout(timeout time.Duration) Option {
	return func(r *Rabbus) error {
		r.config.breaker.timeout = timeout
		return nil
	}
}

// Threshold when a threshold of failures has been reached, future calls to the broker will not run.
// During this state, the circuit breaker will periodically allow the calls to run and, if it is successful,
// will start running the function again. Default value is 5.
func Threshold(threshold uint32) Option {
	return func(r *Rabbus) error {
		if threshold == 0 {
			threshold = 5
		}
		r.config.breaker.threshold = threshold
		return nil
	}
}

// OnStateChange is called whenever the state of CircuitBreaker changes.
func OnStateChange(fn OnStateChangeFunc) Option {
	return func(r *Rabbus) error {
		r.config.breaker.onStateChange = fn
		return nil
	}
}

// AmqpProvider expose a interface for interacting with amqp broker
func AmqpProvider(provider Amqp) Option {
	return func(r *Rabbus) error {
		if provider != nil {
			r.Amqp = provider
			return nil
		}
		return errors.New("unexpected amqp provider")
	}
}

func (r *Rabbus) wrapMessage(c ListenConfig, sourceChan <-chan amqp.Delivery, targetChan chan ConsumerMessage) {
	for m := range sourceChan {
		targetChan <- newConsumerMessage(m)
	}
}

func (r *Rabbus) handleAmqpClose(err error) {
	for {
		time.Sleep(time.Second)
		aw, err := amqpwrap.New(r.config.dsn, r.config.passiveex)
		if err != nil {
			continue
		}

		r.mu.Lock()
		r.Amqp = aw
		r.mu.Unlock()
		r.reconn <- struct{}{}
		break
	}
}

func (r *Rabbus) listenReconn(c ListenConfig, messages chan ConsumerMessage) {
	for range r.reconn {
		msgs, err := r.CreateConsumer(c.Exchange, c.Key, c.Kind, c.Queue, r.config.durable)
		if err != nil {
			r.Close()
			continue
		}

		go r.wrapMessage(c, msgs, messages)
		go r.listenReconn(c, messages)
		break
	}
}

func newBreakerSettings(c config) gobreaker.Settings {
	s := gobreaker.Settings{}
	s.Name = "rabbus-circuit-breaker"
	s.Interval = c.breaker.interval
	s.Timeout = c.breaker.timeout
	s.ReadyToTrip = func(counts gobreaker.Counts) bool {
		return counts.ConsecutiveFailures > c.breaker.threshold
	}
	if c.breaker.onStateChange != nil {
		s.OnStateChange = func(name string, from gobreaker.State, to gobreaker.State) {
			c.breaker.onStateChange(name, from.String(), to.String())
		}
	}
	return s
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      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 quota

import (
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/apigee/istio-mixer-adapter/adapter/auth"
	"github.com/apigee/istio-mixer-adapter/adapter/product"
	"istio.io/istio/mixer/pkg/adapter"
)

const (
	quotaPath             = "/quotas/organization/%s/environment/%s"
	defaultSyncRate       = time.Second
	defaultNumSyncWorkers = 10
	defaultRefreshAfter   = 1 * time.Minute
	defaultDeleteAfter    = 10 * time.Minute
	syncQueueSize         = 100
	resultCacheBufferSize = 30
)

// A Manager tracks multiple Apigee quotas
type Manager struct {
	baseURL        *url.URL
	close          chan bool
	closed         chan bool
	client         *http.Client
	now            func() time.Time
	log            adapter.Logger
	syncRate       time.Duration
	bucketsLock    sync.RWMutex
	buckets        map[string]*bucket // Map from ID -> bucket
	syncQueue      chan *bucket
	numSyncWorkers int
	dupCache       ResultCache
	syncingBuckets map[*bucket]struct{}
}

// NewManager constructs and starts a new Manager. Call Close when done.
func NewManager(env adapter.Env, options Options) (*Manager, error) {
	if err := options.validate(); err != nil {
		return nil, err
	}
	m := newManager(options.BaseURL, options.Client)
	m.Start(env)
	return m, nil
}

// newManager constructs a new Manager
func newManager(baseURL *url.URL, client *http.Client) *Manager {
	return &Manager{
		close:          make(chan bool),
		closed:         make(chan bool),
		client:         client,
		now:            time.Now,
		syncRate:       defaultSyncRate,
		buckets:        map[string]*bucket{},
		syncQueue:      make(chan *bucket, syncQueueSize),
		baseURL:        baseURL,
		numSyncWorkers: defaultNumSyncWorkers,
		dupCache:       ResultCache{size: resultCacheBufferSize},
		syncingBuckets: map[*bucket]struct{}{},
	}
}

// Start starts the manager.
func (m *Manager) Start(env adapter.Env) {
	m.log = env.Logger()
	m.log.Infof("starting quota manager")

	env.ScheduleDaemon(func() {
		m.syncLoop()
	})

	for i := 0; i < m.numSyncWorkers; i++ {
		env.ScheduleDaemon(func() {
			m.syncBucketWorker()
		})
	}
	m.log.Infof("started quota manager with %d workers", m.numSyncWorkers)
}

// Close shuts down the manager.
func (m *Manager) Close() {
	if m == nil {
		return
	}
	m.log.Infof("closing quota manager")
	m.close <- true
	close(m.syncQueue)
	for i := 0; i <= m.numSyncWorkers; i++ {
		<-m.closed
	}
	m.log.Infof("closed quota manager")
}

// Apply a quota request to the local quota bucket and schedule for sync
func (m *Manager) Apply(auth *auth.Context, p *product.APIProduct, args adapter.QuotaArgs) (*Result, error) {

	if result := m.dupCache.Get(args.DeduplicationID); result != nil {
		return result, nil
	}

	quotaID := fmt.Sprintf("%s-%s", auth.Application, p.Name)

	req := &Request{
		Identifier: quotaID,
		Weight:     args.QuotaAmount,
		Interval:   p.QuotaIntervalInt,
		Allow:      p.QuotaLimitInt,
		TimeUnit:   p.QuotaTimeUnit,
	}

	// a new bucket is created if missing or if product is no longer compatible
	var result *Result
	var err error
	forceSync := false
	m.bucketsLock.RLock()
	b, ok := m.buckets[quotaID]
	m.bucketsLock.RUnlock()
	if !ok || !b.compatible(req) {
		m.bucketsLock.Lock()
		b, ok = m.buckets[quotaID]
		if !ok || !b.compatible(req) {
			forceSync = true
			b = newBucket(*req, m, auth)
			m.syncingBuckets[b] = struct{}{}
			defer delete(m.syncingBuckets, b)
			m.buckets[quotaID] = b
			m.log.Debugf("new quota bucket: %s", quotaID)
		}
		m.bucketsLock.Unlock()
	}

	if forceSync {
		err = b.sync() // force sync for new bucket
		result = b.result
	} else {
		result, err = b.apply(req)
	}

	if result != nil && err == nil && args.DeduplicationID != "" {
		m.dupCache.Add(args.DeduplicationID, result)
	}

	return result, err
}

// loop to sync active buckets and deletes old buckets
func (m *Manager) syncLoop() {
	t := time.NewTicker(m.syncRate)
	for {
		select {
		case <-t.C:
			var deleteIDs []string
			m.bucketsLock.RLock()
			for id, b := range m.buckets {
				if b.needToDelete() {
					deleteIDs = append(deleteIDs, id)
				} else if b.needToSync() {
					bucket := b
					m.syncQueue <- bucket
				}
			}
			m.bucketsLock.RUnlock()
			if deleteIDs != nil {
				m.log.Debugf("deleting quota buckets: %v", deleteIDs)
				m.bucketsLock.Lock()
				for _, id := range deleteIDs {
					delete(m.buckets, id)
				}
				m.bucketsLock.Unlock()
			}
		case <-m.close:
			m.log.Debugf("closing quota sync loop")
			t.Stop()
			m.closed <- true
			return
		}
	}
}

// worker routine for syncing a bucket with the server
func (m *Manager) syncBucketWorker() {
	for {
		bucket, ok := <-m.syncQueue
		if ok {
			if _, ok := m.syncingBuckets[bucket]; !ok {
				m.syncingBuckets[bucket] = struct{}{}
				bucket.sync()
				delete(m.syncingBuckets, bucket)
			}
		} else {
			m.log.Debugf("closing quota sync worker")
			m.closed <- true
			return
		}
	}
}

// Options allows us to specify options for how this auth manager will run
type Options struct {
	// Client is a configured HTTPClient
	Client *http.Client
	// BaseURL of the Apigee internal proxy
	BaseURL *url.URL
}

func (o *Options) validate() error {
	if o.Client == nil ||
		o.BaseURL == nil {
		return fmt.Errorf("all quota options are required")
	}
	return nil
}
<commit_msg>guard syncingBuckets with mutex<commit_after>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      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 quota

import (
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/apigee/istio-mixer-adapter/adapter/auth"
	"github.com/apigee/istio-mixer-adapter/adapter/product"
	"istio.io/istio/mixer/pkg/adapter"
)

const (
	quotaPath             = "/quotas/organization/%s/environment/%s"
	defaultSyncRate       = time.Second
	defaultNumSyncWorkers = 10
	defaultRefreshAfter   = 1 * time.Minute
	defaultDeleteAfter    = 10 * time.Minute
	syncQueueSize         = 100
	resultCacheBufferSize = 30
)

// A Manager tracks multiple Apigee quotas
type Manager struct {
	baseURL            *url.URL
	close              chan bool
	closed             chan bool
	client             *http.Client
	now                func() time.Time
	log                adapter.Logger
	syncRate           time.Duration
	bucketsLock        sync.RWMutex
	buckets            map[string]*bucket // Map from ID -> bucket
	syncQueue          chan *bucket
	numSyncWorkers     int
	dupCache           ResultCache
	syncingBuckets     map[*bucket]struct{}
	syncingBucketsLock sync.Mutex
}

// NewManager constructs and starts a new Manager. Call Close when done.
func NewManager(env adapter.Env, options Options) (*Manager, error) {
	if err := options.validate(); err != nil {
		return nil, err
	}
	m := newManager(options.BaseURL, options.Client)
	m.Start(env)
	return m, nil
}

// newManager constructs a new Manager
func newManager(baseURL *url.URL, client *http.Client) *Manager {
	return &Manager{
		close:          make(chan bool),
		closed:         make(chan bool),
		client:         client,
		now:            time.Now,
		syncRate:       defaultSyncRate,
		buckets:        map[string]*bucket{},
		syncQueue:      make(chan *bucket, syncQueueSize),
		baseURL:        baseURL,
		numSyncWorkers: defaultNumSyncWorkers,
		dupCache:       ResultCache{size: resultCacheBufferSize},
		syncingBuckets: map[*bucket]struct{}{},
	}
}

// Start starts the manager.
func (m *Manager) Start(env adapter.Env) {
	m.log = env.Logger()
	m.log.Infof("starting quota manager")

	env.ScheduleDaemon(func() {
		m.syncLoop()
	})

	for i := 0; i < m.numSyncWorkers; i++ {
		env.ScheduleDaemon(func() {
			m.syncBucketWorker()
		})
	}
	m.log.Infof("started quota manager with %d workers", m.numSyncWorkers)
}

// Close shuts down the manager.
func (m *Manager) Close() {
	if m == nil {
		return
	}
	m.log.Infof("closing quota manager")
	m.close <- true
	close(m.syncQueue)
	for i := 0; i <= m.numSyncWorkers; i++ {
		<-m.closed
	}
	m.log.Infof("closed quota manager")
}

// Apply a quota request to the local quota bucket and schedule for sync
func (m *Manager) Apply(auth *auth.Context, p *product.APIProduct, args adapter.QuotaArgs) (*Result, error) {

	if result := m.dupCache.Get(args.DeduplicationID); result != nil {
		return result, nil
	}

	quotaID := fmt.Sprintf("%s-%s", auth.Application, p.Name)

	req := &Request{
		Identifier: quotaID,
		Weight:     args.QuotaAmount,
		Interval:   p.QuotaIntervalInt,
		Allow:      p.QuotaLimitInt,
		TimeUnit:   p.QuotaTimeUnit,
	}

	// a new bucket is created if missing or if product is no longer compatible
	var result *Result
	var err error
	forceSync := false
	m.bucketsLock.RLock()
	b, ok := m.buckets[quotaID]
	m.bucketsLock.RUnlock()
	if !ok || !b.compatible(req) {
		m.bucketsLock.Lock()
		b, ok = m.buckets[quotaID]
		if !ok || !b.compatible(req) {
			forceSync = true
			b = newBucket(*req, m, auth)
			m.syncingBucketsLock.Lock()
			m.syncingBuckets[b] = struct{}{}
			m.syncingBucketsLock.Unlock()
			defer func() {
				m.syncingBucketsLock.Lock()
				delete(m.syncingBuckets, b)
				m.syncingBucketsLock.Unlock()
			}()
			m.buckets[quotaID] = b
			m.log.Debugf("new quota bucket: %s", quotaID)
		}
		m.bucketsLock.Unlock()
	}

	if forceSync {
		err = b.sync() // force sync for new bucket
		result = b.result
	} else {
		result, err = b.apply(req)
	}

	if result != nil && err == nil && args.DeduplicationID != "" {
		m.dupCache.Add(args.DeduplicationID, result)
	}

	return result, err
}

// loop to sync active buckets and deletes old buckets
func (m *Manager) syncLoop() {
	t := time.NewTicker(m.syncRate)
	for {
		select {
		case <-t.C:
			var deleteIDs []string
			m.bucketsLock.RLock()
			for id, b := range m.buckets {
				if b.needToDelete() {
					deleteIDs = append(deleteIDs, id)
				} else if b.needToSync() {
					bucket := b
					m.syncQueue <- bucket
				}
			}
			m.bucketsLock.RUnlock()
			if deleteIDs != nil {
				m.log.Debugf("deleting quota buckets: %v", deleteIDs)
				m.bucketsLock.Lock()
				for _, id := range deleteIDs {
					delete(m.buckets, id)
				}
				m.bucketsLock.Unlock()
			}
		case <-m.close:
			m.log.Debugf("closing quota sync loop")
			t.Stop()
			m.closed <- true
			return
		}
	}
}

// worker routine for syncing a bucket with the server
func (m *Manager) syncBucketWorker() {
	for {
		bucket, ok := <-m.syncQueue
		if ok {
			m.syncingBucketsLock.Lock()
			if _, ok := m.syncingBuckets[bucket]; !ok {
				m.syncingBuckets[bucket] = struct{}{}
				m.syncingBucketsLock.Unlock()
				bucket.sync()
				m.syncingBucketsLock.Lock()
				delete(m.syncingBuckets, bucket)
			}
			m.syncingBucketsLock.Unlock()
		} else {
			m.log.Debugf("closing quota sync worker")
			m.closed <- true
			return
		}
	}
}

// Options allows us to specify options for how this auth manager will run
type Options struct {
	// Client is a configured HTTPClient
	Client *http.Client
	// BaseURL of the Apigee internal proxy
	BaseURL *url.URL
}

func (o *Options) validate() error {
	if o.Client == nil ||
		o.BaseURL == nil {
		return fmt.Errorf("all quota options are required")
	}
	return nil
}
<|endoftext|>