예제 #1
0
func resolveBucketStore(config map[string]string) (leakybucket.Storage, error) {

	switch config["type"] {
	default:
		return nil, errors.New("must specify one of 'redis' or 'memory' storage")
	case "memory":
		return leakybucketMemory.New(), nil
	case "redis":
		return leakybucketRedis.New("tcp", fmt.Sprintf("%s:%s",
			config["host"], config["port"]))
	}
}
예제 #2
0
// Make sure limit.Add adds requests to different buckets
func TestLimitAdd(t *testing.T) {
	data, err := ioutil.ReadFile("../example.yaml")
	if err != nil {
		t.Fatal(err)
	}
	conf, err := config.LoadAndValidateYaml(data)
	if err != nil {
		t.Fatal(err)
	}

	limitconfig := config.Limit{
		Interval: 100,
		Max:      3,
		// matches name: Authorization, match: bearer (from example.yaml)
		Matches: conf.Limits["basic-simple"].Matches,
		// excludes path: /special/resoures/.*
		Excludes: conf.Limits["basic-simple"].Excludes,
		Keys:     conf.Limits["basic-simple"].Keys,
	}

	lim, err := New("test-limit", limitconfig, memory.New())
	limit := lim.(*limit)
	if err != nil {
		t.Error("Could not initialize test-limit")
	}

	request := common.Request{
		"path": "/special/resources/123",
		"headers": http.Header{
			"Authorization": []string{"Basic 12345"},
		},
	}
	for i := uint(1); i < 4; i++ {
		bucketStatus, err := limit.Add(request)
		if err != nil {
			t.Errorf("Error while adding to limit test-limit: %s", err.Error())
		}
		if bucketStatus.Remaining != limit.max-i {
			t.Errorf("Expected remaining %d, found: %d",
				limit.max-i, bucketStatus.Remaining)
		}
	}

	bucketStatus, err := limit.Add(request)
	if err == nil {
		t.Fatal("expected error")
	}
	if err != leakybucket.ErrorFull {
		t.Errorf("Expected leakybucket.ErrorFull error, got: %s", err.Error())
	}

	request2 := common.Request{
		"path": "/special/resources/123",
		"headers": http.Header{
			"Authorization": []string{"Basic ABC"},
		},
	}
	bucketStatus, err = limit.Add(request2)
	if err != nil {
		t.Errorf("Error while adding to limit test-limit: %s", err.Error())
	}
	if bucketStatus.Remaining != 2 {
		t.Errorf("Expected remaining %d, found: %d",
			2, bucketStatus.Remaining)
	}
}