<commit_before>package cli

import (
	"flag"
	"testing"

	rpcpb "github.com/google/shipshape/shipshape/proto/shipshape_rpc_proto"
	"github.com/google/shipshape/shipshape/util/docker"
)

var (
	// There are two ways to specify test flags when using Bazel:
	// 1) In the BUILD file with an args stanza in the _test rule.
	// 2) On the command line using --test_arg (i.e. bazel test --test_arg=-shipshape_test_docker_tag=TAG ...).
	//
	// As of 9 Oct 2015, there are multiple Bazel targets that use --shipshape_test_docker_tag (:shipshape_test_prod
	// and :shipshape_test_local) but there are no targets that set local Kythe.
	dockerTag  = flag.String("shipshape_test_docker_tag", "", "the docker tag for the images to use for testing")
	localKythe = flag.Bool("shipshape_test_local_kythe", false, "if true, don't pull the Kythe docker image")
)

func countFailures(resp rpcpb.ShipshapeResponse) int {
	failures := 0
	for _, analyzeResp := range resp.AnalyzeResponse {
		failures += len(analyzeResp.Failure)
	}
	return failures
}

func countNotes(resp rpcpb.ShipshapeResponse) int {
	notes := 0
	for _, analyzeResp := range resp.AnalyzeResponse {
		notes += len(analyzeResp.Note)
	}
	return notes
}

func countCategoryNotes(resp rpcpb.ShipshapeResponse, category string) int {
	notes := 0
	for _, analyzeResp := range resp.AnalyzeResponse {
		for _, note := range analyzeResp.Note {
			if *note.Category == category {
				notes += 1
			}
		}
	}
	return notes
}

func TestExternalAnalyzers(t *testing.T) {
	// Replaces part of the e2e test
	// Create a fake maven project with android failures

	// Run CLI using a .shipshape file
}

func TestBuiltInAnalyzersPreBuild(t *testing.T) {
	options := Options{
		File:                "shipshape/cli/testdata/workspace1",
		ThirdPartyAnalyzers: []string{},
		Build:               "",
		TriggerCats:         []string{"PostMessage", "JSHint", "go vet", "PyLint"},
		Dind:                false,
		Event:               DefaultEvent,
		Repo:                DefaultRepo,
		StayUp:              true,
		Tag:                 *dockerTag,
		LocalKythe:          *localKythe,
	}
	var allResponses rpcpb.ShipshapeResponse
	options.HandleResponse = func(shipshapeResp *rpcpb.ShipshapeResponse, _ string) error {
		allResponses.AnalyzeResponse = append(allResponses.AnalyzeResponse, shipshapeResp.AnalyzeResponse...)
		return nil
	}
	returnedNotesCount, err := New(options).Run()
	if err != nil {
		t.Fatal(err)
	}
	testName := "TestBuiltInAnalyzerPreBuild"

	if got, want := countFailures(allResponses), 0; got != want {
		t.Errorf("%v: Wrong number of failures; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
	if countedNotes := countNotes(allResponses); returnedNotesCount != countedNotes {
		t.Errorf("%v: Inconsistent note count: returned %v, counted %v (proto data: %v", testName, returnedNotesCount, countedNotes, allResponses)
	}
	if got, want := returnedNotesCount, 39; got != want {
		t.Errorf("%v: Wrong number of notes; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
	if got, want := countCategoryNotes(allResponses, "PostMessage"), 2; got != want {
		t.Errorf("%v: Wrong number of PostMessage notes; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
	if got, want := countCategoryNotes(allResponses, "JSHint"), 3; got != want {
		t.Errorf("%v: Wrong number of JSHint notes; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
	if got, want := countCategoryNotes(allResponses, "go vet"), 1; got != want {
		t.Errorf("%v: Wrong number of go vet notes; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
	if got, want := countCategoryNotes(allResponses, "PyLint"), 33; got != want {
		t.Errorf("%v: Wrong number of PyLint notes; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
}

func TestBuiltInAnalyzersPostBuild(t *testing.T) {
	// Replaces part of the e2e test
	// Test with a kythe maven build
	// PostMessage and ErrorProne
}

func TestStreamsMode(t *testing.T) {
	// Test whether it works in streams mode
	// Before creating this, ensure that streams mode
	// is actually still something we need to support.
}

func TestChangingDirs(t *testing.T) {
	tests := []struct {
		name           string
		file           string
		expectedJSHint int
		expectedGovet  int
		expectedPyLint int
	}{
		{
			name:           "ChildDir",
			file:           "shipshape/cli/testdata/workspace2/subworkspace1",
			expectedJSHint: 0,
			expectedGovet:  1,
			expectedPyLint: 0,
		},
		{
			name:           "SiblingDir",
			file:           "shipshape/cli/testdata/workspace2/subworkspace2",
			expectedJSHint: 0,
			expectedGovet:  0,
			expectedPyLint: 22,
		},
		{
			name:           "ParentDir",
			file:           "shipshape/cli/testdata/workspace2",
			expectedJSHint: 3,
			expectedGovet:  1,
			expectedPyLint: 22,
		},
		{
			name:           "File",
			file:           "shipshape/cli/testdata/workspace2/test.js",
			expectedJSHint: 3,
			expectedGovet:  0,
			expectedPyLint: 0,
		},
	}

	// Clean up the docker state
	if result := docker.Stop("shipping_container", 0, true); result.Err != nil {
		t.Fatalf("Problem cleaning up the docker state; err: %v", result.Err)
	}

	for _, test := range tests {
		options := Options{
			File:                test.file,
			ThirdPartyAnalyzers: []string{},
			Build:               "",
			TriggerCats:         []string{"PostMessage", "JSHint", "go vet", "PyLint"},
			Dind:                false,
			Event:               DefaultEvent,
			Repo:                DefaultRepo,
			StayUp:              true,
			Tag:                 *dockerTag,
			LocalKythe:          *localKythe,
		}
		var allResponses rpcpb.ShipshapeResponse
		options.HandleResponse = func(shipshapeResp *rpcpb.ShipshapeResponse, _ string) error {
			allResponses.AnalyzeResponse =
				append(allResponses.AnalyzeResponse, shipshapeResp.AnalyzeResponse...)
			return nil
		}
		testName := test.name
		if _, err := New(options).Run(); err != nil {
			t.Fatalf("%v: Failure on service call; err: %v", testName, err)
		}
		if got, want := countFailures(allResponses), 0; got != want {
			t.Errorf("%v: Wrong number of failures; got %v, want %v (proto data: %v)",
				testName, got, want, allResponses)
		}
		if got, want := countCategoryNotes(allResponses, "JSHint"), test.expectedJSHint; got != want {
			t.Errorf("%v: Wrong number of JSHint notes; got %v, want %v (proto data: %v)",
				testName, got, want, allResponses)
		}
		if got, want := countCategoryNotes(allResponses, "go vet"), test.expectedGovet; got != want {
			t.Errorf("%v: Wrong number of go vet notes; got %v, want %v (proto data: %v)",
				testName, got, want, allResponses)
		}
		if got, want := countCategoryNotes(allResponses, "PyLint"), test.expectedPyLint; got != want {
			t.Errorf("%v: Wrong number of PyLint notes; got %v, want %v (proto data: %v)",
				testName, got, want, allResponses)
		}
	}
}

func dumpLogs() {

}

func checkOutput(category string, numResults int) {

}
<commit_msg>Adds container exists check to CLI test<commit_after>package cli

import (
	"flag"
	"testing"

	rpcpb "github.com/google/shipshape/shipshape/proto/shipshape_rpc_proto"
	"github.com/google/shipshape/shipshape/util/docker"
)

var (
	// There are two ways to specify test flags when using Bazel:
	// 1) In the BUILD file with an args stanza in the _test rule.
	// 2) On the command line using --test_arg (i.e. bazel test --test_arg=-shipshape_test_docker_tag=TAG ...).
	//
	// As of 9 Oct 2015, there are multiple Bazel targets that use --shipshape_test_docker_tag (:shipshape_test_prod
	// and :shipshape_test_local) but there are no targets that set local Kythe.
	dockerTag  = flag.String("shipshape_test_docker_tag", "", "the docker tag for the images to use for testing")
	localKythe = flag.Bool("shipshape_test_local_kythe", false, "if true, don't pull the Kythe docker image")
)

func countFailures(resp rpcpb.ShipshapeResponse) int {
	failures := 0
	for _, analyzeResp := range resp.AnalyzeResponse {
		failures += len(analyzeResp.Failure)
	}
	return failures
}

func countNotes(resp rpcpb.ShipshapeResponse) int {
	notes := 0
	for _, analyzeResp := range resp.AnalyzeResponse {
		notes += len(analyzeResp.Note)
	}
	return notes
}

func countCategoryNotes(resp rpcpb.ShipshapeResponse, category string) int {
	notes := 0
	for _, analyzeResp := range resp.AnalyzeResponse {
		for _, note := range analyzeResp.Note {
			if *note.Category == category {
				notes += 1
			}
		}
	}
	return notes
}

func TestExternalAnalyzers(t *testing.T) {
	// Replaces part of the e2e test
	// Create a fake maven project with android failures

	// Run CLI using a .shipshape file
}

func TestBuiltInAnalyzersPreBuild(t *testing.T) {
	options := Options{
		File:                "shipshape/cli/testdata/workspace1",
		ThirdPartyAnalyzers: []string{},
		Build:               "",
		TriggerCats:         []string{"PostMessage", "JSHint", "go vet", "PyLint"},
		Dind:                false,
		Event:               DefaultEvent,
		Repo:                DefaultRepo,
		StayUp:              true,
		Tag:                 *dockerTag,
		LocalKythe:          *localKythe,
	}
	var allResponses rpcpb.ShipshapeResponse
	options.HandleResponse = func(shipshapeResp *rpcpb.ShipshapeResponse, _ string) error {
		allResponses.AnalyzeResponse = append(allResponses.AnalyzeResponse, shipshapeResp.AnalyzeResponse...)
		return nil
	}
	returnedNotesCount, err := New(options).Run()
	if err != nil {
		t.Fatal(err)
	}
	testName := "TestBuiltInAnalyzerPreBuild"

	if got, want := countFailures(allResponses), 0; got != want {
		t.Errorf("%v: Wrong number of failures; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
	if countedNotes := countNotes(allResponses); returnedNotesCount != countedNotes {
		t.Errorf("%v: Inconsistent note count: returned %v, counted %v (proto data: %v", testName, returnedNotesCount, countedNotes, allResponses)
	}
	if got, want := returnedNotesCount, 39; got != want {
		t.Errorf("%v: Wrong number of notes; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
	if got, want := countCategoryNotes(allResponses, "PostMessage"), 2; got != want {
		t.Errorf("%v: Wrong number of PostMessage notes; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
	if got, want := countCategoryNotes(allResponses, "JSHint"), 3; got != want {
		t.Errorf("%v: Wrong number of JSHint notes; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
	if got, want := countCategoryNotes(allResponses, "go vet"), 1; got != want {
		t.Errorf("%v: Wrong number of go vet notes; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
	if got, want := countCategoryNotes(allResponses, "PyLint"), 33; got != want {
		t.Errorf("%v: Wrong number of PyLint notes; got %v, want %v (proto data: %v)", testName, got, want, allResponses)
	}
}

func TestBuiltInAnalyzersPostBuild(t *testing.T) {
	// Replaces part of the e2e test
	// Test with a kythe maven build
	// PostMessage and ErrorProne
}

func TestStreamsMode(t *testing.T) {
	// Test whether it works in streams mode
	// Before creating this, ensure that streams mode
	// is actually still something we need to support.
}

func TestChangingDirs(t *testing.T) {
	tests := []struct {
		name           string
		file           string
		expectedJSHint int
		expectedGovet  int
		expectedPyLint int
	}{
		{
			name:           "ChildDir",
			file:           "shipshape/cli/testdata/workspace2/subworkspace1",
			expectedJSHint: 0,
			expectedGovet:  1,
			expectedPyLint: 0,
		},
		{
			name:           "SiblingDir",
			file:           "shipshape/cli/testdata/workspace2/subworkspace2",
			expectedJSHint: 0,
			expectedGovet:  0,
			expectedPyLint: 22,
		},
		{
			name:           "ParentDir",
			file:           "shipshape/cli/testdata/workspace2",
			expectedJSHint: 3,
			expectedGovet:  1,
			expectedPyLint: 22,
		},
		{
			name:           "File",
			file:           "shipshape/cli/testdata/workspace2/test.js",
			expectedJSHint: 3,
			expectedGovet:  0,
			expectedPyLint: 0,
		},
	}

	// Clean up the docker state
	container := "shipping_container"
	exists, err := docker.ContainerExists(container)
	if err != nil {
		t.Fatalf("Problem checking docker state; err: %v", err)
	}
	if exists {
		if result := docker.Stop(container, 0, true); result.Err != nil {
			t.Fatalf("Problem cleaning up the docker state; err: %v", result.Err)
		}
	}

	for _, test := range tests {
		options := Options{
			File:                test.file,
			ThirdPartyAnalyzers: []string{},
			Build:               "",
			TriggerCats:         []string{"PostMessage", "JSHint", "go vet", "PyLint"},
			Dind:                false,
			Event:               DefaultEvent,
			Repo:                DefaultRepo,
			StayUp:              true,
			Tag:                 *dockerTag,
			LocalKythe:          *localKythe,
		}
		var allResponses rpcpb.ShipshapeResponse
		options.HandleResponse = func(shipshapeResp *rpcpb.ShipshapeResponse, _ string) error {
			allResponses.AnalyzeResponse =
				append(allResponses.AnalyzeResponse, shipshapeResp.AnalyzeResponse...)
			return nil
		}
		testName := test.name
		if _, err := New(options).Run(); err != nil {
			t.Fatalf("%v: Failure on service call; err: %v", testName, err)
		}
		if got, want := countFailures(allResponses), 0; got != want {
			t.Errorf("%v: Wrong number of failures; got %v, want %v (proto data: %v)",
				testName, got, want, allResponses)
		}
		if got, want := countCategoryNotes(allResponses, "JSHint"), test.expectedJSHint; got != want {
			t.Errorf("%v: Wrong number of JSHint notes; got %v, want %v (proto data: %v)",
				testName, got, want, allResponses)
		}
		if got, want := countCategoryNotes(allResponses, "go vet"), test.expectedGovet; got != want {
			t.Errorf("%v: Wrong number of go vet notes; got %v, want %v (proto data: %v)",
				testName, got, want, allResponses)
		}
		if got, want := countCategoryNotes(allResponses, "PyLint"), test.expectedPyLint; got != want {
			t.Errorf("%v: Wrong number of PyLint notes; got %v, want %v (proto data: %v)",
				testName, got, want, allResponses)
		}
	}
}

func dumpLogs() {

}

func checkOutput(category string, numResults int) {

}
<|endoftext|>


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

package gensupport

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"
)

// ResumableUpload is used by the generated APIs to provide resumable uploads.
// It is not used by developers directly.
type ResumableUpload struct {
	Client *http.Client
	// URI is the resumable resource destination provided by the server after specifying "&uploadType=resumable".
	URI       string
	UserAgent string // User-Agent for header of the request
	// Media is the object being uploaded.
	Media *MediaBuffer
	// MediaType defines the media type, e.g. "image/jpeg".
	MediaType string

	mu       sync.Mutex // guards progress
	progress int64      // number of bytes uploaded so far

	// Callback is an optional function that will be periodically called with the cumulative number of bytes uploaded.
	Callback func(int64)

	// Retry optionally configures retries for requests made against the upload.
	Retry *RetryConfig

	// ChunkRetryDeadline configures the per-chunk deadline after which no further
	// retries should happen.
	ChunkRetryDeadline time.Duration
}

// Progress returns the number of bytes uploaded at this point.
func (rx *ResumableUpload) Progress() int64 {
	rx.mu.Lock()
	defer rx.mu.Unlock()
	return rx.progress
}

// doUploadRequest performs a single HTTP request to upload data.
// off specifies the offset in rx.Media from which data is drawn.
// size is the number of bytes in data.
// final specifies whether data is the final chunk to be uploaded.
func (rx *ResumableUpload) doUploadRequest(ctx context.Context, data io.Reader, off, size int64, final bool) (*http.Response, error) {
	req, err := http.NewRequest("POST", rx.URI, data)
	if err != nil {
		return nil, err
	}

	req.ContentLength = size
	var contentRange string
	if final {
		if size == 0 {
			contentRange = fmt.Sprintf("bytes */%v", off)
		} else {
			contentRange = fmt.Sprintf("bytes %v-%v/%v", off, off+size-1, off+size)
		}
	} else {
		contentRange = fmt.Sprintf("bytes %v-%v/*", off, off+size-1)
	}
	req.Header.Set("Content-Range", contentRange)
	req.Header.Set("Content-Type", rx.MediaType)
	req.Header.Set("User-Agent", rx.UserAgent)

	// Google's upload endpoint uses status code 308 for a
	// different purpose than the "308 Permanent Redirect"
	// since-standardized in RFC 7238. Because of the conflict in
	// semantics, Google added this new request header which
	// causes it to not use "308" and instead reply with 200 OK
	// and sets the upload-specific "X-HTTP-Status-Code-Override:
	// 308" response header.
	req.Header.Set("X-GUploader-No-308", "yes")

	return SendRequest(ctx, rx.Client, req)
}

func statusResumeIncomplete(resp *http.Response) bool {
	// This is how the server signals "status resume incomplete"
	// when X-GUploader-No-308 is set to "yes":
	return resp != nil && resp.Header.Get("X-Http-Status-Code-Override") == "308"
}

// reportProgress calls a user-supplied callback to report upload progress.
// If old==updated, the callback is not called.
func (rx *ResumableUpload) reportProgress(old, updated int64) {
	if updated-old == 0 {
		return
	}
	rx.mu.Lock()
	rx.progress = updated
	rx.mu.Unlock()
	if rx.Callback != nil {
		rx.Callback(updated)
	}
}

// transferChunk performs a single HTTP request to upload a single chunk from rx.Media.
func (rx *ResumableUpload) transferChunk(ctx context.Context) (*http.Response, error) {
	chunk, off, size, err := rx.Media.Chunk()

	done := err == io.EOF
	if !done && err != nil {
		return nil, err
	}

	res, err := rx.doUploadRequest(ctx, chunk, off, int64(size), done)
	if err != nil {
		return res, err
	}

	// We sent "X-GUploader-No-308: yes" (see comment elsewhere in
	// this file), so we don't expect to get a 308.
	if res.StatusCode == 308 {
		return nil, errors.New("unexpected 308 response status code")
	}

	if res.StatusCode == http.StatusOK {
		rx.reportProgress(off, off+int64(size))
	}

	if statusResumeIncomplete(res) {
		rx.Media.Next()
	}
	return res, nil
}

// Upload starts the process of a resumable upload with a cancellable context.
// It retries using the provided back off strategy until cancelled or the
// strategy indicates to stop retrying.
// It is called from the auto-generated API code and is not visible to the user.
// Before sending an HTTP request, Upload calls any registered hook functions,
// and calls the returned functions after the request returns (see send.go).
// rx is private to the auto-generated API code.
// Exactly one of resp or err will be nil.  If resp is non-nil, the caller must call resp.Body.Close.
func (rx *ResumableUpload) Upload(ctx context.Context) (resp *http.Response, err error) {

	// There are a couple of cases where it's possible for err and resp to both
	// be non-nil. However, we expose a simpler contract to our callers: exactly
	// one of resp and err will be non-nil. This means that any response body
	// must be closed here before returning a non-nil error.
	var prepareReturn = func(resp *http.Response, err error) (*http.Response, error) {
		if err != nil {
			if resp != nil && resp.Body != nil {
				resp.Body.Close()
			}
			return nil, err
		}
		return resp, nil
	}
	// Configure retryable error criteria.
	errorFunc := rx.Retry.errorFunc()

	// Configure per-chunk retry deadline.
	var retryDeadline time.Duration
	if rx.ChunkRetryDeadline != 0 {
		retryDeadline = rx.ChunkRetryDeadline
	} else {
		retryDeadline = defaultRetryDeadline
	}

	// Send all chunks.
	for {
		var pause time.Duration

		// Each chunk gets its own initialized-at-zero backoff.
		bo := rx.Retry.backoff()
		quitAfter := time.After(retryDeadline)

		// Retry loop for a single chunk.
		for {
			select {
			case <-ctx.Done():
				if err == nil {
					err = ctx.Err()
				}
				return prepareReturn(resp, err)
			case <-time.After(pause):
			case <-quitAfter:
				return prepareReturn(resp, err)
			}

			// Check for context cancellation or timeout once more. If more than one
			// case in the select statement above was satisfied at the same time, Go
			// will choose one arbitrarily.
			// That can cause an operation to go through even if the context was
			// canceled before or the timeout was reached.
			select {
			case <-ctx.Done():
				if err == nil {
					err = ctx.Err()
				}
				return prepareReturn(resp, err)
			case <-quitAfter:
				return prepareReturn(resp, err)
			default:
			}

			resp, err = rx.transferChunk(ctx)

			var status int
			if resp != nil {
				status = resp.StatusCode
			}

			// Check if we should retry the request.
			if !errorFunc(status, err) {
				break
			}

			pause = bo.Pause()
			if resp != nil && resp.Body != nil {
				resp.Body.Close()
			}
		}

		// If the chunk was uploaded successfully, but there's still
		// more to go, upload the next chunk without any delay.
		if statusResumeIncomplete(resp) {
			resp.Body.Close()
			continue
		}

		return prepareReturn(resp, err)
	}
}
<commit_msg>fix(gensupport): cover ChunkRetryDeadline edge case (#1430)<commit_after>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gensupport

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"
)

// ResumableUpload is used by the generated APIs to provide resumable uploads.
// It is not used by developers directly.
type ResumableUpload struct {
	Client *http.Client
	// URI is the resumable resource destination provided by the server after specifying "&uploadType=resumable".
	URI       string
	UserAgent string // User-Agent for header of the request
	// Media is the object being uploaded.
	Media *MediaBuffer
	// MediaType defines the media type, e.g. "image/jpeg".
	MediaType string

	mu       sync.Mutex // guards progress
	progress int64      // number of bytes uploaded so far

	// Callback is an optional function that will be periodically called with the cumulative number of bytes uploaded.
	Callback func(int64)

	// Retry optionally configures retries for requests made against the upload.
	Retry *RetryConfig

	// ChunkRetryDeadline configures the per-chunk deadline after which no further
	// retries should happen.
	ChunkRetryDeadline time.Duration
}

// Progress returns the number of bytes uploaded at this point.
func (rx *ResumableUpload) Progress() int64 {
	rx.mu.Lock()
	defer rx.mu.Unlock()
	return rx.progress
}

// doUploadRequest performs a single HTTP request to upload data.
// off specifies the offset in rx.Media from which data is drawn.
// size is the number of bytes in data.
// final specifies whether data is the final chunk to be uploaded.
func (rx *ResumableUpload) doUploadRequest(ctx context.Context, data io.Reader, off, size int64, final bool) (*http.Response, error) {
	req, err := http.NewRequest("POST", rx.URI, data)
	if err != nil {
		return nil, err
	}

	req.ContentLength = size
	var contentRange string
	if final {
		if size == 0 {
			contentRange = fmt.Sprintf("bytes */%v", off)
		} else {
			contentRange = fmt.Sprintf("bytes %v-%v/%v", off, off+size-1, off+size)
		}
	} else {
		contentRange = fmt.Sprintf("bytes %v-%v/*", off, off+size-1)
	}
	req.Header.Set("Content-Range", contentRange)
	req.Header.Set("Content-Type", rx.MediaType)
	req.Header.Set("User-Agent", rx.UserAgent)

	// Google's upload endpoint uses status code 308 for a
	// different purpose than the "308 Permanent Redirect"
	// since-standardized in RFC 7238. Because of the conflict in
	// semantics, Google added this new request header which
	// causes it to not use "308" and instead reply with 200 OK
	// and sets the upload-specific "X-HTTP-Status-Code-Override:
	// 308" response header.
	req.Header.Set("X-GUploader-No-308", "yes")

	return SendRequest(ctx, rx.Client, req)
}

func statusResumeIncomplete(resp *http.Response) bool {
	// This is how the server signals "status resume incomplete"
	// when X-GUploader-No-308 is set to "yes":
	return resp != nil && resp.Header.Get("X-Http-Status-Code-Override") == "308"
}

// reportProgress calls a user-supplied callback to report upload progress.
// If old==updated, the callback is not called.
func (rx *ResumableUpload) reportProgress(old, updated int64) {
	if updated-old == 0 {
		return
	}
	rx.mu.Lock()
	rx.progress = updated
	rx.mu.Unlock()
	if rx.Callback != nil {
		rx.Callback(updated)
	}
}

// transferChunk performs a single HTTP request to upload a single chunk from rx.Media.
func (rx *ResumableUpload) transferChunk(ctx context.Context) (*http.Response, error) {
	chunk, off, size, err := rx.Media.Chunk()

	done := err == io.EOF
	if !done && err != nil {
		return nil, err
	}

	res, err := rx.doUploadRequest(ctx, chunk, off, int64(size), done)
	if err != nil {
		return res, err
	}

	// We sent "X-GUploader-No-308: yes" (see comment elsewhere in
	// this file), so we don't expect to get a 308.
	if res.StatusCode == 308 {
		return nil, errors.New("unexpected 308 response status code")
	}

	if res.StatusCode == http.StatusOK {
		rx.reportProgress(off, off+int64(size))
	}

	if statusResumeIncomplete(res) {
		rx.Media.Next()
	}
	return res, nil
}

// Upload starts the process of a resumable upload with a cancellable context.
// It retries using the provided back off strategy until cancelled or the
// strategy indicates to stop retrying.
// It is called from the auto-generated API code and is not visible to the user.
// Before sending an HTTP request, Upload calls any registered hook functions,
// and calls the returned functions after the request returns (see send.go).
// rx is private to the auto-generated API code.
// Exactly one of resp or err will be nil.  If resp is non-nil, the caller must call resp.Body.Close.
func (rx *ResumableUpload) Upload(ctx context.Context) (resp *http.Response, err error) {

	// There are a couple of cases where it's possible for err and resp to both
	// be non-nil. However, we expose a simpler contract to our callers: exactly
	// one of resp and err will be non-nil. This means that any response body
	// must be closed here before returning a non-nil error.
	var prepareReturn = func(resp *http.Response, err error) (*http.Response, error) {
		if err != nil {
			if resp != nil && resp.Body != nil {
				resp.Body.Close()
			}
			return nil, err
		}
		// This case is very unlikely but possible only if rx.ChunkRetryDeadline is
		// set to a very small value, in which case no requests will be sent before
		// the deadline. Return an error to avoid causing a panic.
		if resp == nil {
			return nil, fmt.Errorf("upload request to %v not sent, choose larger value for ChunkRetryDealine", rx.URI)
		}
		return resp, nil
	}
	// Configure retryable error criteria.
	errorFunc := rx.Retry.errorFunc()

	// Configure per-chunk retry deadline.
	var retryDeadline time.Duration
	if rx.ChunkRetryDeadline != 0 {
		retryDeadline = rx.ChunkRetryDeadline
	} else {
		retryDeadline = defaultRetryDeadline
	}

	// Send all chunks.
	for {
		var pause time.Duration

		// Each chunk gets its own initialized-at-zero backoff.
		bo := rx.Retry.backoff()
		quitAfter := time.After(retryDeadline)

		// Retry loop for a single chunk.
		for {
			select {
			case <-ctx.Done():
				if err == nil {
					err = ctx.Err()
				}
				return prepareReturn(resp, err)
			case <-time.After(pause):
			case <-quitAfter:
				return prepareReturn(resp, err)
			}

			// Check for context cancellation or timeout once more. If more than one
			// case in the select statement above was satisfied at the same time, Go
			// will choose one arbitrarily.
			// That can cause an operation to go through even if the context was
			// canceled before or the timeout was reached.
			select {
			case <-ctx.Done():
				if err == nil {
					err = ctx.Err()
				}
				return prepareReturn(resp, err)
			case <-quitAfter:
				return prepareReturn(resp, err)
			default:
			}

			resp, err = rx.transferChunk(ctx)

			var status int
			if resp != nil {
				status = resp.StatusCode
			}

			// Check if we should retry the request.
			if !errorFunc(status, err) {
				break
			}

			pause = bo.Pause()
			if resp != nil && resp.Body != nil {
				resp.Body.Close()
			}
		}

		// If the chunk was uploaded successfully, but there's still
		// more to go, upload the next chunk without any delay.
		if statusResumeIncomplete(resp) {
			resp.Body.Close()
			continue
		}

		return prepareReturn(resp, err)
	}
}
<|endoftext|>


<commit_before>/*
Package station grabs fallback data from the radio station.
*/
package station

import (
	"encoding/json"
	"flag"
	"io/ioutil"
	"net/http"
	"time"
)

var (
	latestInfo Wrapper

	bugTime = flag.Int("station-poke-delay", 15, "how stale the info can get")
)

type Wrapper struct {
	Age  time.Time
	Info Info
}

type Info struct {
	Icestats struct {
		Admin              string `json:"admin"`
		Host               string `json:"host"`
		Location           string `json:"location"`
		ServerID           string `json:"server_id"`
		ServerStart        string `json:"server_start"`
		ServerStartIso8601 string `json:"server_start_iso8601"`
		Source             []struct {
			Artist             string      `json:"artist"`
			AudioBitrate       int         `json:"audio_bitrate"`
			AudioChannels      int         `json:"audio_channels"`
			AudioInfo          string      `json:"audio_info"`
			AudioSamplerate    int         `json:"audio_samplerate"`
			Channels           int         `json:"channels"`
			Genre              string      `json:"genre"`
			IceBitrate         int         `json:"ice-bitrate"`
			ListenerPeak       int         `json:"listener_peak"`
			Listeners          int         `json:"listeners"`
			Listenurl          string      `json:"listenurl"`
			Quality            string      `json:"quality"`
			Samplerate         int         `json:"samplerate"`
			ServerDescription  string      `json:"server_description"`
			ServerName         string      `json:"server_name"`
			ServerType         string      `json:"server_type"`
			ServerURL          string      `json:"server_url"`
			StreamStart        string      `json:"stream_start"`
			StreamStartIso8601 string      `json:"stream_start_iso8601"`
			Subtype            string      `json:"subtype"`
			Title              string      `json:"title"`
			Dummy              interface{} `json:"dummy"`
		} `json:"source"`
	} `json:"icestats"`
}

func GetStats() (Info, error) {
	now := time.Now()
	if now.Before(latestInfo.Age.Add(time.Second * time.Duration(*bugTime))) {
		return latestInfo.Info, nil
	}

	i := Info{}

	resp, err := http.Get("http://dj.bronyradio.com:7090/status-json.xsl")
	if err != nil {
		return Info{}, err
	}

	defer resp.Body.Close()

	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return Info{}, err
	}

	err = json.Unmarshal(content, &i)
	if err != nil {
		return Info{}, err
	}

	latestInfo.Info = i
	latestInfo.Age = now

	return latestInfo.Info, nil
}
<commit_msg>internal/pvfm: update station metadata fetcher<commit_after>/*
Package station grabs fallback data from the radio station.
*/
package station

import (
	"encoding/json"
	"flag"
	"io/ioutil"
	"net/http"
	"time"
)

var (
	latestInfo Wrapper

	bugTime = flag.Int("station-poke-delay", 15, "how stale the info can get")
)

type Wrapper struct {
	Age  time.Time
	Info Info
}

type Info struct {
	Icestats struct {
		Admin              string `json:"admin"`
		Host               string `json:"host"`
		Location           string `json:"location"`
		ServerID           string `json:"server_id"`
		ServerStart        string `json:"server_start"`
		ServerStartIso8601 string `json:"server_start_iso8601"`
		Source             []struct {
			Artist             string      `json:"artist"`
			AudioBitrate       int         `json:"audio_bitrate"`
			AudioChannels      int         `json:"audio_channels"`
			AudioInfo          string      `json:"audio_info"`
			AudioSamplerate    int         `json:"audio_samplerate"`
			Channels           int         `json:"channels"`
			Genre              string      `json:"genre"`
			IceBitrate         int         `json:"ice-bitrate"`
			ListenerPeak       int         `json:"listener_peak"`
			Listeners          int         `json:"listeners"`
			Listenurl          string      `json:"listenurl"`
			Quality            string      `json:"quality"`
			Samplerate         int         `json:"samplerate"`
			ServerDescription  string      `json:"server_description"`
			ServerName         string      `json:"server_name"`
			ServerType         string      `json:"server_type"`
			ServerURL          string      `json:"server_url"`
			StreamStart        string      `json:"stream_start"`
			StreamStartIso8601 string      `json:"stream_start_iso8601"`
			Subtype            string      `json:"subtype"`
			Title              string      `json:"title"`
			Dummy              interface{} `json:"dummy"`
		} `json:"source"`
	} `json:"icestats"`
}

func GetStats() (Info, error) {
	now := time.Now()
	if now.Before(latestInfo.Age.Add(time.Second * time.Duration(*bugTime))) {
		return latestInfo.Info, nil
	}

	i := Info{}

	resp, err := http.Get("http://dj.bronyradio.com:8000/status-json.xsl")
	if err != nil {
		return Info{}, err
	}

	defer resp.Body.Close()

	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return Info{}, err
	}

	err = json.Unmarshal(content, &i)
	if err != nil {
		return Info{}, err
	}

	latestInfo.Info = i
	latestInfo.Age = now

	return latestInfo.Info, nil
}
<|endoftext|>


<commit_before>package Marshal

import (
	. "github.com/xitongsys/parquet-go/Common"
	. "github.com/xitongsys/parquet-go/SchemaHandler"
	"reflect"
)

type KeyValue struct {
	Key   reflect.Value
	Value reflect.Value
}

type MapRecord struct {
	KeyValues []KeyValue
	Index     int
}

//Convert the table map to objects slice. desInterface is a slice of pointers of objects
func Unmarshal(tableMap *map[string]*Table, bgn int, end int, dstInterface interface{}, schemaHandler *SchemaHandler) {
	ot := reflect.TypeOf(dstInterface).Elem().Elem()
	tableIndex := make(map[string]int)
	tableBgn, tableEnd := make(map[string]int), make(map[string]int)

	tmpRes := reflect.ValueOf(dstInterface).Elem()

	for name, table := range *tableMap {
		ln := len(table.Values)
		num := -1
		tableBgn[name], tableEnd[name] = -1, -1
		for i := 0; i < ln; i++ {
			if table.RepetitionLevels[i] == 0 {
				num++
				if num == bgn {
					tableBgn[name] = i
					tableIndex[name] = i
				}
				if num == end {
					tableEnd[name] = i
					break
				}
			}
		}
		if tableEnd[name] < 0 {
			tableEnd[name] = ln
		}
		if tableBgn[name] < 0 {
			return
		}
	}

	flag := true
	for flag {
		flag = false
		obj := reflect.New(ot).Elem()
		mapRecord := make(map[reflect.Value]*MapRecord)
		sliceRecord := make(map[reflect.Value]int)

		for name, table := range *tableMap {
			path := table.Path
			end := tableEnd[name]

			if tableIndex[name] >= end {
				continue
			}

			for key, _ := range mapRecord {
				mapRecord[key].Index = -1
			}
			for key, _ := range sliceRecord {
				sliceRecord[key] = -1
			}

			for {
				rl, dl := 0, 0
				po := obj
				pathIndex := 0
				for pathIndex < len(path) {
					if po.Type().Kind() == reflect.Struct {
						if (table.DefinitionLevels[tableIndex[name]] < table.MaxDefinitionLevel &&
							table.DefinitionLevels[tableIndex[name]] > int32(dl)) ||
							table.DefinitionLevels[tableIndex[name]] == table.MaxDefinitionLevel {
							pathIndex++
							po = po.FieldByName(HeadToUpper(path[pathIndex])) //HeadToUpper is for some filed is lowercase

						} else {
							break
						}

					} else if po.Type().Kind() == reflect.Slice {
						if po.IsNil() {
							po.Set(reflect.MakeSlice(po.Type(), 0, 0))
						}
						if _, ok := sliceRecord[po]; !ok {
							sliceRecord[po] = -1
						}

						if table.DefinitionLevels[tableIndex[name]] > int32(dl) {
							pathIndex += 2
							dl += 1
							rl += 1
							if table.RepetitionLevels[tableIndex[name]] <= int32(rl) {
								sliceRecord[po]++
								if sliceRecord[po] >= po.Len() {
									potmp := reflect.Append(po, reflect.New(po.Type().Elem()).Elem())
									po.Set(potmp)
								}
								po = po.Index(sliceRecord[po])
							} else {
								po = po.Index(sliceRecord[po])
							}

						} else {
							break
						}

					} else if po.Type().Kind() == reflect.Map {
						if po.IsNil() {
							po.Set(reflect.MakeMap(po.Type()))
						}

						if _, ok := mapRecord[po]; !ok {
							mapRecord[po] = &MapRecord{KeyValues: make([]KeyValue, 0), Index: -1}
						}

						if table.DefinitionLevels[tableIndex[name]] > int32(dl) {
							if path[pathIndex+2] == "value" {
								pathIndex += 2
								dl += 1
								rl += 1

								if table.RepetitionLevels[tableIndex[name]] <= int32(rl) {
									mapRecord[po].Index++
									if mapRecord[po].Index >= len(mapRecord[po].KeyValues) {
										mapRecord[po].KeyValues = append(mapRecord[po].KeyValues,
											KeyValue{Key: reflect.ValueOf(nil), Value: reflect.ValueOf(nil)})
										value := reflect.New(po.Type().Elem()).Elem()
										mapRecord[po].KeyValues[mapRecord[po].Index].Value = value
									}
									if !mapRecord[po].KeyValues[mapRecord[po].Index].Value.IsValid() {
										mapRecord[po].KeyValues[mapRecord[po].Index].Value = reflect.New(po.Type().Elem()).Elem()
									}
									po = mapRecord[po].KeyValues[mapRecord[po].Index].Value

								} else {
									po = mapRecord[po].KeyValues[mapRecord[po].Index].Value
								}

							} else if path[pathIndex+2] == "key" {
								mapRecord[po].Index++
								if mapRecord[po].Index >= len(mapRecord[po].KeyValues) {
									mapRecord[po].KeyValues = append(mapRecord[po].KeyValues,
										KeyValue{Key: reflect.ValueOf(nil), Value: reflect.ValueOf(nil)})
								}
								mapRecord[po].KeyValues[mapRecord[po].Index].Key = reflect.ValueOf(table.Values[tableIndex[name]])
								break
							}
						} else {
							break
						}

					} else if po.Type().Kind() == reflect.Ptr {
						dl += 1
						if int32(dl) > table.DefinitionLevels[tableIndex[name]] {
							break
						}
						if po.IsNil() {
							po.Set(reflect.New(po.Type().Elem()))
						}
						po = po.Elem()

					} else {
						po.Set(reflect.ValueOf(table.Values[tableIndex[name]]))
						break
					}
				} //for pathIndex < len(path) {

				tableIndex[name]++
				if (tableIndex[name] < end && table.RepetitionLevels[tableIndex[name]] == 0) ||
					(tableIndex[name] >= end) {
					break
				}
			}
			if tableIndex[name] < end {
				flag = true
			}

		} //for name, table := range tableMap

		for m, record := range mapRecord {
			for _, kv := range record.KeyValues {
				m.SetMapIndex(kv.Key, kv.Value)
			}
		}

		tmpRes = reflect.Append(tmpRes, obj)

	}
	reflect.ValueOf(dstInterface).Elem().Set(tmpRes)

}
<commit_msg>del unmarshal dot import<commit_after>package Marshal

import (
	"github.com/xitongsys/parquet-go/Common"
	"github.com/xitongsys/parquet-go/SchemaHandler"
	"reflect"
)

type KeyValue struct {
	Key   reflect.Value
	Value reflect.Value
}

type MapRecord struct {
	KeyValues []KeyValue
	Index     int
}

//Convert the table map to objects slice. desInterface is a slice of pointers of objects
func Unmarshal(tableMap *map[string]*Common.Table, bgn int, end int, dstInterface interface{}, schemaHandler *SchemaHandler.SchemaHandler) {
	ot := reflect.TypeOf(dstInterface).Elem().Elem()
	tableIndex := make(map[string]int)
	tableBgn, tableEnd := make(map[string]int), make(map[string]int)

	tmpRes := reflect.ValueOf(dstInterface).Elem()

	for name, table := range *tableMap {
		ln := len(table.Values)
		num := -1
		tableBgn[name], tableEnd[name] = -1, -1
		for i := 0; i < ln; i++ {
			if table.RepetitionLevels[i] == 0 {
				num++
				if num == bgn {
					tableBgn[name] = i
					tableIndex[name] = i
				}
				if num == end {
					tableEnd[name] = i
					break
				}
			}
		}
		if tableEnd[name] < 0 {
			tableEnd[name] = ln
		}
		if tableBgn[name] < 0 {
			return
		}
	}

	flag := true
	for flag {
		flag = false
		obj := reflect.New(ot).Elem()
		mapRecord := make(map[reflect.Value]*MapRecord)
		sliceRecord := make(map[reflect.Value]int)

		for name, table := range *tableMap {
			path := table.Path
			end := tableEnd[name]

			if tableIndex[name] >= end {
				continue
			}

			for key, _ := range mapRecord {
				mapRecord[key].Index = -1
			}
			for key, _ := range sliceRecord {
				sliceRecord[key] = -1
			}

			for {
				rl, dl := 0, 0
				po := obj
				pathIndex := 0
				for pathIndex < len(path) {
					if po.Type().Kind() == reflect.Struct {
						if (table.DefinitionLevels[tableIndex[name]] < table.MaxDefinitionLevel &&
							table.DefinitionLevels[tableIndex[name]] > int32(dl)) ||
							table.DefinitionLevels[tableIndex[name]] == table.MaxDefinitionLevel {
							pathIndex++
							po = po.FieldByName(Common.HeadToUpper(path[pathIndex])) //HeadToUpper is for some filed is lowercase

						} else {
							break
						}

					} else if po.Type().Kind() == reflect.Slice {
						if po.IsNil() {
							po.Set(reflect.MakeSlice(po.Type(), 0, 0))
						}
						if _, ok := sliceRecord[po]; !ok {
							sliceRecord[po] = -1
						}

						if table.DefinitionLevels[tableIndex[name]] > int32(dl) {
							pathIndex += 2
							dl += 1
							rl += 1
							if table.RepetitionLevels[tableIndex[name]] <= int32(rl) {
								sliceRecord[po]++
								if sliceRecord[po] >= po.Len() {
									potmp := reflect.Append(po, reflect.New(po.Type().Elem()).Elem())
									po.Set(potmp)
								}
								po = po.Index(sliceRecord[po])
							} else {
								po = po.Index(sliceRecord[po])
							}

						} else {
							break
						}

					} else if po.Type().Kind() == reflect.Map {
						if po.IsNil() {
							po.Set(reflect.MakeMap(po.Type()))
						}

						if _, ok := mapRecord[po]; !ok {
							mapRecord[po] = &MapRecord{KeyValues: make([]KeyValue, 0), Index: -1}
						}

						if table.DefinitionLevels[tableIndex[name]] > int32(dl) {
							if path[pathIndex+2] == "value" {
								pathIndex += 2
								dl += 1
								rl += 1

								if table.RepetitionLevels[tableIndex[name]] <= int32(rl) {
									mapRecord[po].Index++
									if mapRecord[po].Index >= len(mapRecord[po].KeyValues) {
										mapRecord[po].KeyValues = append(mapRecord[po].KeyValues,
											KeyValue{Key: reflect.ValueOf(nil), Value: reflect.ValueOf(nil)})
										value := reflect.New(po.Type().Elem()).Elem()
										mapRecord[po].KeyValues[mapRecord[po].Index].Value = value
									}
									if !mapRecord[po].KeyValues[mapRecord[po].Index].Value.IsValid() {
										mapRecord[po].KeyValues[mapRecord[po].Index].Value = reflect.New(po.Type().Elem()).Elem()
									}
									po = mapRecord[po].KeyValues[mapRecord[po].Index].Value

								} else {
									po = mapRecord[po].KeyValues[mapRecord[po].Index].Value
								}

							} else if path[pathIndex+2] == "key" {
								mapRecord[po].Index++
								if mapRecord[po].Index >= len(mapRecord[po].KeyValues) {
									mapRecord[po].KeyValues = append(mapRecord[po].KeyValues,
										KeyValue{Key: reflect.ValueOf(nil), Value: reflect.ValueOf(nil)})
								}
								mapRecord[po].KeyValues[mapRecord[po].Index].Key = reflect.ValueOf(table.Values[tableIndex[name]])
								break
							}
						} else {
							break
						}

					} else if po.Type().Kind() == reflect.Ptr {
						dl += 1
						if int32(dl) > table.DefinitionLevels[tableIndex[name]] {
							break
						}
						if po.IsNil() {
							po.Set(reflect.New(po.Type().Elem()))
						}
						po = po.Elem()

					} else {
						po.Set(reflect.ValueOf(table.Values[tableIndex[name]]))
						break
					}
				} //for pathIndex < len(path) {

				tableIndex[name]++
				if (tableIndex[name] < end && table.RepetitionLevels[tableIndex[name]] == 0) ||
					(tableIndex[name] >= end) {
					break
				}
			}
			if tableIndex[name] < end {
				flag = true
			}

		} //for name, table := range tableMap

		for m, record := range mapRecord {
			for _, kv := range record.KeyValues {
				m.SetMapIndex(kv.Key, kv.Value)
			}
		}

		tmpRes = reflect.Append(tmpRes, obj)

	}
	reflect.ValueOf(dstInterface).Elem().Set(tmpRes)

}
<|endoftext|>


<commit_before>package core

import (
	"fmt"
	"log"
	"os"
	"os/user"
	"sync"

	"github.com/tinzenite/channel"
	"github.com/tinzenite/model"
	"github.com/tinzenite/shared"
)

/*
Tinzenite is the struct on which all important operations should be called.
*/
type Tinzenite struct {
	Path           string
	auth           *Authentication
	selfpeer       *shared.Peer
	channel        *channel.Channel
	cInterface     *chaninterface
	allPeers       []*shared.Peer
	model          *model.Model
	sendChannel    chan shared.UpdateMessage
	muteFlag       bool
	stop           chan bool
	wg             sync.WaitGroup
	peerValidation PeerValidation
}

/*
Sync the entire file system model of the directory, first locally and then
remotely if other peers are connected. NOTE: All Sync{|Local|Remote} methods can
block for a potentially long time, especially when first run!
*/
func (t *Tinzenite) Sync() error {
	// mute updates because we'll sync models later
	t.muteFlag = true
	// first ensure that local model is up to date
	err := t.SyncLocal()
	// and reset
	t.muteFlag = false
	if err != nil {
		return err
	}
	return t.SyncRemote()
}

/*
SyncLocal changes. Will send updates to connected peers but not synchronize with
other peers.
*/
func (t *Tinzenite) SyncLocal() error {
	// update peers
	err := t.applyPeers()
	if err != nil {
		return err
	}
	// update model
	return t.model.Update()
}

/*
SyncRemote changes. Will request the models of connected peers and merge them
to the local peer.

TODO: fetches model from other peers and syncs (this is for manual sync)
*/
func (t *Tinzenite) SyncRemote() error {
	online, err := t.channel.OnlineAddresses()
	if err != nil {
		return err
	}
	for _, address := range online {
		t.cInterface.SyncModel(address)
	}
	return nil
}

/*
Address of this Tinzenite peer that can be used to connect to.
*/
func (t *Tinzenite) Address() (string, error) {
	return t.channel.ConnectionAddress()
}

/*
Close cleanly stores everything and shuts Tinzenite down.
*/
func (t *Tinzenite) Close() {
	// send stop signal
	t.stop <- false
	// wait for it to close
	t.wg.Wait()
	// store all information
	t.Store()
	// FINALLY close (afterwards because I still need info from channel for store!)
	t.channel.Close()
}

/*
Store the tinzenite directory structure to disk. Will resolve all important
objects and store them so that it can later be reloaded. NOTE: Will not update
the full model, so be sure to have called Update() to guarantee an up to date
save.
*/
func (t *Tinzenite) Store() error {
	err := shared.MakeDotTinzenite(t.Path)
	if err != nil {
		return err
	}
	// write all peers to files
	for _, peer := range t.allPeers {
		err := peer.Store(t.Path)
		if err != nil {
			return err
		}
	}
	// store local peer info with toxdata
	toxData, err := t.channel.ToxData()
	if err != nil {
		return err
	}
	toxPeerDump := &shared.ToxPeerDump{
		SelfPeer: t.selfpeer,
		ToxData:  toxData}
	err = toxPeerDump.Store(t.Path)
	if err != nil {
		return err
	}
	// store auth file
	err = t.auth.Store(t.Path)
	if err != nil {
		return err
	}
	// update model for tinzenite dir to catch above stores
	err = t.model.PartialUpdate(t.Path + "/" + shared.TINZENITEDIR)
	if err != nil {
		return err
	}
	// finally store model (last because peers and org stuff is included!)
	return t.model.Store()
}

/*
PrintStatus returns a formatted string of the peer status.
*/
func (t *Tinzenite) PrintStatus() string {
	var out string
	out += "Online:\n"
	addresses, err := t.channel.FriendAddresses()
	if err != nil {
		out += "channel.FriendAddresses failed!"
	} else {
		var count int
		for _, address := range addresses {
			online, err := t.channel.IsOnline(address)
			var insert string
			if err != nil {
				insert = "ERROR"
			} else {
				insert = fmt.Sprintf("%v", online)
			}
			out += address[:16] + " :: " + insert + "\n"
			count++
		}
		out += "Total friends: " + fmt.Sprintf("%d", count)
	}
	return out
}

/*
DisconnectPeer does exactly that. NOTE: this is a passive action and not due
anything except remove the peer from the network. The peer is not further
notified.

TODO: maybe not use name but Identification?
*/
func (t *Tinzenite) DisconnectPeer(peerName string) {
	var newPeers []*shared.Peer
	for _, peer := range t.allPeers {
		if t.selfpeer.Identification == peer.Identification {
			continue
		}
		if peer.Name == peerName {
			log.Println("Removing", peer.Name, "at", peer.Address[:8])
			// delete peer file
			path := shared.CreatePath(t.Path, shared.TINZENITEDIR+"/"+shared.ORGDIR+"/"+shared.PEERSDIR+"/"+peer.Identification+shared.ENDING)
			err := t.model.ApplyRemove(path, nil)
			if err != nil {
				log.Println("DisconnectPeer:", err)
			}
			// remove from channel
			err = t.channel.RemoveConnection(peer.Address)
			if err != nil {
				log.Println("DisconnectPeer:", err)
			}
			// continue does not readd to tinzenite, removing the reference to it
			continue
		}
		newPeers = append(newPeers, peer)
	}
	t.allPeers = newPeers
}

/*
applyPeers loads all peers and readies the communication channel accordingly.
*/
func (t *Tinzenite) applyPeers() error {
	peers, err := shared.LoadPeers(t.Path)
	if err != nil {
		return err
	}
	// make sure they are all tox ready
	for _, peer := range peers {
		// tox will return an error if the address has already been added, so we just ignore it
		_ = t.channel.AcceptConnection(peer.Address)
	}
	// finally apply
	t.allPeers = peers
	return nil
}

/*
Send the given message to the peer with the address if online.
*/
func (t *Tinzenite) send(address, msg string) error {
	return t.channel.Send(address, msg)
}

/*
Merge an update message to the local model.

TODO: with move implemented one of the merged file copies can be kept and simply
renamed, also solving the ID problem. Look into this!
*/
func (t *Tinzenite) merge(msg *shared.UpdateMessage) error {
	relPath := shared.CreatePath(t.Path, msg.Object.Path)
	// first: apply local changes to model (this is why writing PartialUpdate was no waste of time, isn't this cool?! :D)
	err := t.model.PartialUpdate(relPath.FullPath())
	if err != nil {
		return err
	}
	// check if content is same, no need for merge then (except for version)
	stin, err := t.model.GetInfo(relPath)
	if err != nil {
		log.Println("Core:", "Can not check if content is same!")
	} else {
		if stin.Content == msg.Object.Content {
			log.Println("Core:", "Merge not required as updates are in sync!")
			// so all we need to do is apply the version update
			t.model.ApplyModify(relPath, &msg.Object)
			return nil
		}
	}
	// second: move to new name
	err = os.Rename(relPath.FullPath(), relPath.FullPath()+LOCAL)
	if err != nil {
		log.Println("Original can not be found!")
		return err
	}
	err = t.model.PartialUpdate(relPath.FullPath() + LOCAL)
	if err != nil {
		return err
	}
	// third: remove original
	err = t.model.ApplyRemove(relPath, nil)
	if err != nil {
		return err
	}
	// fourth: change path and apply remote as create
	msg.Operation = shared.OpCreate
	msg.Object.Path = relPath.SubPath() + REMOTE
	msg.Object.Name = relPath.LastElement() + REMOTE
	oldID := msg.Object.Identification
	msg.Object.Identification, err = shared.NewIdentifier()
	if err != nil {
		return err
	}
	// new id --> rename temp file
	tempPath := t.Path + "/" + shared.TINZENITEDIR + "/" + shared.TEMPDIR
	err = os.Rename(tempPath+"/"+oldID, tempPath+"/"+msg.Object.Identification)
	if err != nil {
		log.Println("Updating remote object file failed!")
		return err
	}
	// fifth: create remote file
	err = t.model.ApplyCreate(relPath.Apply(relPath.FullPath()+REMOTE), &msg.Object)
	if err != nil {
		return err
	}
	return nil
}

/*
storeGlobalConfig stores the path value into the user's home directory so that clients
can locate it.
*/
func (t *Tinzenite) storeGlobalConfig() error {
	// ready outside data
	user, err := user.Current()
	if err != nil {
		return err
	}
	path := user.HomeDir + "/.config/tinzenite"
	err = shared.MakeDirectory(path)
	if err != nil {
		return err
	}
	path += "/" + shared.DIRECTORYLIST
	file, err := os.OpenFile(path, shared.FILEFLAGCREATEAPPEND, shared.FILEPERMISSIONMODE)
	if err != nil {
		return err
	}
	defer file.Close()
	// write path to file
	_, err = file.WriteString(t.Path + "\n")
	if err != nil {
		return err
	}
	// ensure that the file is valid
	return shared.PrettifyDirectoryList()
}

/*
initialize the background process.
*/
func (t *Tinzenite) initialize() {
	// prepare send channel that will distribute updates
	t.wg.Add(1)
	t.stop = make(chan bool, 1)
	t.sendChannel = make(chan shared.UpdateMessage, 1)
	go t.background()
	t.model.Register(t.sendChannel)
}

/*
background function that handles all async stuff that needs to be done.
*/
func (t *Tinzenite) background() {
	for {
		select {
		case <-t.stop:
			t.wg.Done()
			return
		case msg := <-t.sendChannel:
			// if muted don't send updates
			if t.muteFlag {
				continue
			}
			online, err := t.channel.OnlineAddresses()
			if err != nil {
				log.Println("Background:", err)
			}
			for _, address := range online {
				t.channel.Send(address, msg.String())
			}
		} // select
	} // for
}
<commit_msg>use defer so that better behaviour<commit_after>package core

import (
	"fmt"
	"log"
	"os"
	"os/user"
	"sync"

	"github.com/tinzenite/channel"
	"github.com/tinzenite/model"
	"github.com/tinzenite/shared"
)

/*
Tinzenite is the struct on which all important operations should be called.
*/
type Tinzenite struct {
	Path           string
	auth           *Authentication
	selfpeer       *shared.Peer
	channel        *channel.Channel
	cInterface     *chaninterface
	allPeers       []*shared.Peer
	model          *model.Model
	sendChannel    chan shared.UpdateMessage
	muteFlag       bool
	stop           chan bool
	wg             sync.WaitGroup
	peerValidation PeerValidation
}

/*
Sync the entire file system model of the directory, first locally and then
remotely if other peers are connected. NOTE: All Sync{|Local|Remote} methods can
block for a potentially long time, especially when first run!
*/
func (t *Tinzenite) Sync() error {
	// mute updates because we'll sync models later
	t.muteFlag = true
	// defer setting it back guaranteed
	defer func() { t.muteFlag = false }()
	// first ensure that local model is up to date
	err := t.SyncLocal()
	if err != nil {
		return err
	}
	return t.SyncRemote()
}

/*
SyncLocal changes. Will send updates to connected peers but not synchronize with
other peers.
*/
func (t *Tinzenite) SyncLocal() error {
	// update peers
	err := t.applyPeers()
	if err != nil {
		return err
	}
	// update model
	return t.model.Update()
}

/*
SyncRemote changes. Will request the models of connected peers and merge them
to the local peer.

TODO: fetches model from other peers and syncs (this is for manual sync)
*/
func (t *Tinzenite) SyncRemote() error {
	online, err := t.channel.OnlineAddresses()
	if err != nil {
		return err
	}
	for _, address := range online {
		t.cInterface.SyncModel(address)
	}
	return nil
}

/*
Address of this Tinzenite peer that can be used to connect to.
*/
func (t *Tinzenite) Address() (string, error) {
	return t.channel.ConnectionAddress()
}

/*
Close cleanly stores everything and shuts Tinzenite down.
*/
func (t *Tinzenite) Close() {
	// send stop signal
	t.stop <- false
	// wait for it to close
	t.wg.Wait()
	// store all information
	t.Store()
	// FINALLY close (afterwards because I still need info from channel for store!)
	t.channel.Close()
}

/*
Store the tinzenite directory structure to disk. Will resolve all important
objects and store them so that it can later be reloaded. NOTE: Will not update
the full model, so be sure to have called Update() to guarantee an up to date
save.
*/
func (t *Tinzenite) Store() error {
	err := shared.MakeDotTinzenite(t.Path)
	if err != nil {
		return err
	}
	// write all peers to files
	for _, peer := range t.allPeers {
		err := peer.Store(t.Path)
		if err != nil {
			return err
		}
	}
	// store local peer info with toxdata
	toxData, err := t.channel.ToxData()
	if err != nil {
		return err
	}
	toxPeerDump := &shared.ToxPeerDump{
		SelfPeer: t.selfpeer,
		ToxData:  toxData}
	err = toxPeerDump.Store(t.Path)
	if err != nil {
		return err
	}
	// store auth file
	err = t.auth.Store(t.Path)
	if err != nil {
		return err
	}
	// update model for tinzenite dir to catch above stores
	err = t.model.PartialUpdate(t.Path + "/" + shared.TINZENITEDIR)
	if err != nil {
		return err
	}
	// finally store model (last because peers and org stuff is included!)
	return t.model.Store()
}

/*
PrintStatus returns a formatted string of the peer status.
*/
func (t *Tinzenite) PrintStatus() string {
	var out string
	out += "Online:\n"
	addresses, err := t.channel.FriendAddresses()
	if err != nil {
		out += "channel.FriendAddresses failed!"
	} else {
		var count int
		for _, address := range addresses {
			online, err := t.channel.IsOnline(address)
			var insert string
			if err != nil {
				insert = "ERROR"
			} else {
				insert = fmt.Sprintf("%v", online)
			}
			out += address[:16] + " :: " + insert + "\n"
			count++
		}
		out += "Total friends: " + fmt.Sprintf("%d", count)
	}
	return out
}

/*
DisconnectPeer does exactly that. NOTE: this is a passive action and not due
anything except remove the peer from the network. The peer is not further
notified.

TODO: maybe not use name but Identification?
*/
func (t *Tinzenite) DisconnectPeer(peerName string) {
	var newPeers []*shared.Peer
	for _, peer := range t.allPeers {
		if t.selfpeer.Identification == peer.Identification {
			continue
		}
		if peer.Name == peerName {
			log.Println("Removing", peer.Name, "at", peer.Address[:8])
			// delete peer file
			path := shared.CreatePath(t.Path, shared.TINZENITEDIR+"/"+shared.ORGDIR+"/"+shared.PEERSDIR+"/"+peer.Identification+shared.ENDING)
			err := t.model.ApplyRemove(path, nil)
			if err != nil {
				log.Println("DisconnectPeer:", err)
			}
			// remove from channel
			err = t.channel.RemoveConnection(peer.Address)
			if err != nil {
				log.Println("DisconnectPeer:", err)
			}
			// continue does not readd to tinzenite, removing the reference to it
			continue
		}
		newPeers = append(newPeers, peer)
	}
	t.allPeers = newPeers
}

/*
applyPeers loads all peers and readies the communication channel accordingly.
*/
func (t *Tinzenite) applyPeers() error {
	peers, err := shared.LoadPeers(t.Path)
	if err != nil {
		return err
	}
	// make sure they are all tox ready
	for _, peer := range peers {
		// tox will return an error if the address has already been added, so we just ignore it
		_ = t.channel.AcceptConnection(peer.Address)
	}
	// finally apply
	t.allPeers = peers
	return nil
}

/*
Send the given message to the peer with the address if online.
*/
func (t *Tinzenite) send(address, msg string) error {
	return t.channel.Send(address, msg)
}

/*
Merge an update message to the local model.

TODO: with move implemented one of the merged file copies can be kept and simply
renamed, also solving the ID problem. Look into this!
*/
func (t *Tinzenite) merge(msg *shared.UpdateMessage) error {
	relPath := shared.CreatePath(t.Path, msg.Object.Path)
	// first: apply local changes to model (this is why writing PartialUpdate was no waste of time, isn't this cool?! :D)
	err := t.model.PartialUpdate(relPath.FullPath())
	if err != nil {
		return err
	}
	// check if content is same, no need for merge then (except for version)
	stin, err := t.model.GetInfo(relPath)
	if err != nil {
		log.Println("Core:", "Can not check if content is same!")
	} else {
		if stin.Content == msg.Object.Content {
			log.Println("Core:", "Merge not required as updates are in sync!")
			// so all we need to do is apply the version update
			t.model.ApplyModify(relPath, &msg.Object)
			return nil
		}
	}
	// second: move to new name
	err = os.Rename(relPath.FullPath(), relPath.FullPath()+LOCAL)
	if err != nil {
		log.Println("Original can not be found!")
		return err
	}
	err = t.model.PartialUpdate(relPath.FullPath() + LOCAL)
	if err != nil {
		return err
	}
	// third: remove original
	err = t.model.ApplyRemove(relPath, nil)
	if err != nil {
		return err
	}
	// fourth: change path and apply remote as create
	msg.Operation = shared.OpCreate
	msg.Object.Path = relPath.SubPath() + REMOTE
	msg.Object.Name = relPath.LastElement() + REMOTE
	oldID := msg.Object.Identification
	msg.Object.Identification, err = shared.NewIdentifier()
	if err != nil {
		return err
	}
	// new id --> rename temp file
	tempPath := t.Path + "/" + shared.TINZENITEDIR + "/" + shared.TEMPDIR
	err = os.Rename(tempPath+"/"+oldID, tempPath+"/"+msg.Object.Identification)
	if err != nil {
		log.Println("Updating remote object file failed!")
		return err
	}
	// fifth: create remote file
	err = t.model.ApplyCreate(relPath.Apply(relPath.FullPath()+REMOTE), &msg.Object)
	if err != nil {
		return err
	}
	return nil
}

/*
storeGlobalConfig stores the path value into the user's home directory so that clients
can locate it.
*/
func (t *Tinzenite) storeGlobalConfig() error {
	// ready outside data
	user, err := user.Current()
	if err != nil {
		return err
	}
	path := user.HomeDir + "/.config/tinzenite"
	err = shared.MakeDirectory(path)
	if err != nil {
		return err
	}
	path += "/" + shared.DIRECTORYLIST
	file, err := os.OpenFile(path, shared.FILEFLAGCREATEAPPEND, shared.FILEPERMISSIONMODE)
	if err != nil {
		return err
	}
	defer file.Close()
	// write path to file
	_, err = file.WriteString(t.Path + "\n")
	if err != nil {
		return err
	}
	// ensure that the file is valid
	return shared.PrettifyDirectoryList()
}

/*
initialize the background process.
*/
func (t *Tinzenite) initialize() {
	// prepare send channel that will distribute updates
	t.wg.Add(1)
	t.stop = make(chan bool, 1)
	t.sendChannel = make(chan shared.UpdateMessage, 1)
	go t.background()
	t.model.Register(t.sendChannel)
}

/*
background function that handles all async stuff that needs to be done.
*/
func (t *Tinzenite) background() {
	for {
		select {
		case <-t.stop:
			t.wg.Done()
			return
		case msg := <-t.sendChannel:
			// if muted don't send updates
			if t.muteFlag {
				continue
			}
			online, err := t.channel.OnlineAddresses()
			if err != nil {
				log.Println("Background:", err)
			}
			for _, address := range online {
				t.channel.Send(address, msg.String())
			}
		} // select
	} // for
}
<|endoftext|>


<commit_before>// Package tollbooth provides rate-limiting logic to HTTP request handler.
package tollbooth

import (
	"net/http"
	"strings"

	"fmt"
	"math"

	"github.com/didip/tollbooth/errors"
	"github.com/didip/tollbooth/libstring"
	"github.com/didip/tollbooth/limiter"
)

// setResponseHeaders configures X-Rate-Limit-Limit and X-Rate-Limit-Duration
func setResponseHeaders(lmt *limiter.Limiter, w http.ResponseWriter, r *http.Request) {
	w.Header().Add("X-Rate-Limit-Limit", fmt.Sprintf("%.2f", lmt.GetMax()))
	w.Header().Add("X-Rate-Limit-Duration", "1")
	w.Header().Add("X-Rate-Limit-Request-Forwarded-For", r.Header.Get("X-Forwarded-For"))
	w.Header().Add("X-Rate-Limit-Request-Remote-Addr", r.RemoteAddr)
}

// NewLimiter is a convenience function to limiter.New.
func NewLimiter(max float64, tbOptions *limiter.ExpirableOptions) *limiter.Limiter {
	return limiter.New(tbOptions).SetMax(max).SetBurst(int(math.Max(1, max)))
}

// LimitByKeys keeps track number of request made by keys separated by pipe.
// It returns HTTPError when limit is exceeded.
func LimitByKeys(lmt *limiter.Limiter, keys []string) *errors.HTTPError {
	if lmt.LimitReached(strings.Join(keys, "|")) {
		return &errors.HTTPError{Message: lmt.GetMessage(), StatusCode: lmt.GetStatusCode()}
	}

	return nil
}

// BuildKeys generates a slice of keys to rate-limit by given limiter and request structs.
func BuildKeys(lmt *limiter.Limiter, r *http.Request) [][]string {
	remoteIP := libstring.RemoteIP(lmt.GetIPLookups(), lmt.GetForwardedForIndexFromBehind(), r)
	path := r.URL.Path
	sliceKeys := make([][]string, 0)

	// Don't BuildKeys if remoteIP is blank.
	if remoteIP == "" {
		return sliceKeys
	}

	lmtMethods := lmt.GetMethods()
	lmtHeaders := lmt.GetHeaders()
	lmtBasicAuthUsers := lmt.GetBasicAuthUsers()

	lmtHeadersIsSet := len(lmtHeaders) > 0
	lmtBasicAuthUsersIsSet := len(lmtBasicAuthUsers) > 0

	method := ""
	if lmtMethods != nil && libstring.StringInSlice(lmtMethods, r.Method) {
		method = r.Method
	}

	usernameToLimit := ""
	if lmtBasicAuthUsersIsSet {
		username, _, ok := r.BasicAuth()
		if ok && libstring.StringInSlice(lmtBasicAuthUsers, username) {
			usernameToLimit = username
		}
	}

	headerValuesToLimit := [][]string{}
	if lmtHeadersIsSet {
		for headerKey, headerValues := range lmtHeaders {
			reqHeaderValue := r.Header.Get(headerKey)
			if reqHeaderValue == "" {
				continue
			}

			if headerValues == nil || len(headerValues) <= 0 {
				// If header values are empty, rate-limit all request containing headerKey.
				headerValuesToLimit = append(headerValuesToLimit, []string{headerKey, reqHeaderValue})

			} else if len(headerValues) > 0 {
				// If header values are not empty, rate-limit all request with headerKey and headerValues.
				for _, headerValue := range headerValues {
					if r.Header.Get(headerKey) == headerValue {
						headerValuesToLimit = append(headerValuesToLimit, []string{headerKey, headerValue})
						break
					}
				}
			}
		}
	}
	if len(headerValuesToLimit) <= 0 {
		headerValuesToLimit = append(headerValuesToLimit, []string{"", ""})
	}

	for _, header := range headerValuesToLimit {
		sliceKeys = append(sliceKeys, []string{remoteIP, path, method, header[0], header[1], usernameToLimit})
	}

	return sliceKeys
}

// LimitByRequest builds keys based on http.Request struct,
// loops through all the keys, and check if any one of them returns HTTPError.
func LimitByRequest(lmt *limiter.Limiter, w http.ResponseWriter, r *http.Request) *errors.HTTPError {
	setResponseHeaders(lmt, w, r)

	sliceKeys := BuildKeys(lmt, r)

	// Loop sliceKeys and check if one of them has error.
	for _, keys := range sliceKeys {
		httpError := LimitByKeys(lmt, keys)
		if httpError != nil {
			return httpError
		}
	}

	return nil
}

// LimitHandler is a middleware that performs rate-limiting given http.Handler struct.
func LimitHandler(lmt *limiter.Limiter, next http.Handler) http.Handler {
	middle := func(w http.ResponseWriter, r *http.Request) {
		httpError := LimitByRequest(lmt, w, r)
		if httpError != nil {
			lmt.ExecOnLimitReached(w, r)
			w.Header().Add("Content-Type", lmt.GetMessageContentType())
			w.WriteHeader(httpError.StatusCode)
			w.Write([]byte(httpError.Message))
			return
		}

		// There's no rate-limit error, serve the next handler.
		next.ServeHTTP(w, r)
	}

	return http.HandlerFunc(middle)
}

// LimitFuncHandler is a middleware that performs rate-limiting given request handler function.
func LimitFuncHandler(lmt *limiter.Limiter, nextFunc func(http.ResponseWriter, *http.Request)) http.Handler {
	return LimitHandler(lmt, http.HandlerFunc(nextFunc))
}
<commit_msg>Fix test of issue 66<commit_after>// Package tollbooth provides rate-limiting logic to HTTP request handler.
package tollbooth

import (
	"net/http"
	"strings"

	"fmt"
	"math"

	"github.com/didip/tollbooth/errors"
	"github.com/didip/tollbooth/libstring"
	"github.com/didip/tollbooth/limiter"
)

// setResponseHeaders configures X-Rate-Limit-Limit and X-Rate-Limit-Duration
func setResponseHeaders(lmt *limiter.Limiter, w http.ResponseWriter, r *http.Request) {
	w.Header().Add("X-Rate-Limit-Limit", fmt.Sprintf("%.2f", lmt.GetMax()))
	w.Header().Add("X-Rate-Limit-Duration", "1")
	w.Header().Add("X-Rate-Limit-Request-Forwarded-For", r.Header.Get("X-Forwarded-For"))
	w.Header().Add("X-Rate-Limit-Request-Remote-Addr", r.RemoteAddr)
}

// NewLimiter is a convenience function to limiter.New.
func NewLimiter(max float64, tbOptions *limiter.ExpirableOptions) *limiter.Limiter {
	return limiter.New(tbOptions).SetMax(max).SetBurst(int(math.Max(1, max)))
}

// LimitByKeys keeps track number of request made by keys separated by pipe.
// It returns HTTPError when limit is exceeded.
func LimitByKeys(lmt *limiter.Limiter, keys []string) *errors.HTTPError {
	if lmt.LimitReached(strings.Join(keys, "|")) {
		return &errors.HTTPError{Message: lmt.GetMessage(), StatusCode: lmt.GetStatusCode()}
	}

	return nil
}

// BuildKeys generates a slice of keys to rate-limit by given limiter and request structs.
func BuildKeys(lmt *limiter.Limiter, r *http.Request) [][]string {
	remoteIP := libstring.RemoteIP(lmt.GetIPLookups(), lmt.GetForwardedForIndexFromBehind(), r)
	path := r.URL.Path
	sliceKeys := make([][]string, 0)

	// Don't BuildKeys if remoteIP is blank.
	if remoteIP == "" {
		return sliceKeys
	}

	lmtMethods := lmt.GetMethods()
	lmtHeaders := lmt.GetHeaders()
	lmtBasicAuthUsers := lmt.GetBasicAuthUsers()

	lmtHeadersIsSet := len(lmtHeaders) > 0
	lmtBasicAuthUsersIsSet := len(lmtBasicAuthUsers) > 0

	method := ""
	if lmtMethods != nil && libstring.StringInSlice(lmtMethods, r.Method) {
		method = r.Method
	}

	usernameToLimit := ""
	if lmtBasicAuthUsersIsSet {
		username, _, ok := r.BasicAuth()
		if ok && libstring.StringInSlice(lmtBasicAuthUsers, username) {
			usernameToLimit = username
		}
	}

	headerValuesToLimit := [][]string{}
	if lmtHeadersIsSet {
		for headerKey, headerValues := range lmtHeaders {
			reqHeaderValue := r.Header.Get(headerKey)
			if reqHeaderValue == "" {
				continue
			}

			if headerValues == nil || len(headerValues) <= 0 {
				// If header values are empty, rate-limit all request containing headerKey.
				headerValuesToLimit = append(headerValuesToLimit, []string{headerKey, reqHeaderValue})

			} else if len(headerValues) > 0 {
				// If header values are not empty, rate-limit all request with headerKey and headerValues.
				for _, headerValue := range headerValues {
					if r.Header.Get(headerKey) == headerValue {
						headerValuesToLimit = append(headerValuesToLimit, []string{headerKey, headerValue})
						break
					}
				}
			}
		}
	} else {
		headerValuesToLimit = append(headerValuesToLimit, []string{"", ""})
	}

	for _, header := range headerValuesToLimit {
		sliceKeys = append(sliceKeys, []string{remoteIP, path, method, header[0], header[1], usernameToLimit})
	}

	return sliceKeys
}

// LimitByRequest builds keys based on http.Request struct,
// loops through all the keys, and check if any one of them returns HTTPError.
func LimitByRequest(lmt *limiter.Limiter, w http.ResponseWriter, r *http.Request) *errors.HTTPError {
	setResponseHeaders(lmt, w, r)

	sliceKeys := BuildKeys(lmt, r)

	// Loop sliceKeys and check if one of them has error.
	for _, keys := range sliceKeys {
		httpError := LimitByKeys(lmt, keys)
		if httpError != nil {
			return httpError
		}
	}

	return nil
}

// LimitHandler is a middleware that performs rate-limiting given http.Handler struct.
func LimitHandler(lmt *limiter.Limiter, next http.Handler) http.Handler {
	middle := func(w http.ResponseWriter, r *http.Request) {
		httpError := LimitByRequest(lmt, w, r)
		if httpError != nil {
			lmt.ExecOnLimitReached(w, r)
			w.Header().Add("Content-Type", lmt.GetMessageContentType())
			w.WriteHeader(httpError.StatusCode)
			w.Write([]byte(httpError.Message))
			return
		}

		// There's no rate-limit error, serve the next handler.
		next.ServeHTTP(w, r)
	}

	return http.HandlerFunc(middle)
}

// LimitFuncHandler is a middleware that performs rate-limiting given request handler function.
func LimitFuncHandler(lmt *limiter.Limiter, nextFunc func(http.ResponseWriter, *http.Request)) http.Handler {
	return LimitHandler(lmt, http.HandlerFunc(nextFunc))
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"crypto/tls"
	"flag"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/coreos/go-semver/semver"
	"github.com/fnproject/fn/fnlb/lb"
	"github.com/sirupsen/logrus"
)

const VERSION = "0.0.223"

func main() {
	// XXX (reed): normalize
	level, err := logrus.ParseLevel(os.Getenv("LOG_LEVEL"))
	if err != nil {
		level = logrus.InfoLevel
	}
	logrus.SetLevel(level)

	fnodes := flag.String("nodes", "", "comma separated list of functions nodes")
	minAPIVersion := flag.String("min-api-version", "0.0.190", "minimal node API to accept")

	var conf lb.Config
	flag.StringVar(&conf.DBurl, "db", "sqlite3://:memory:", "backend to store nodes, default to in memory; use k8s for kuberneted")
	flag.StringVar(&conf.Listen, "listen", ":8081", "port to run on")
	flag.StringVar(&conf.MgmtListen, "mgmt-listen", ":8081", "management port to run on")
	flag.IntVar(&conf.ShutdownTimeout, "shutdown-timeout", 0, "graceful shutdown timeout")
	flag.IntVar(&conf.HealthcheckInterval, "hc-interval", 3, "how often to check f(x) nodes, in seconds")
	flag.StringVar(&conf.HealthcheckEndpoint, "hc-path", "/version", "endpoint to determine node health")
	flag.IntVar(&conf.HealthcheckUnhealthy, "hc-unhealthy", 2, "threshold of failed checks to declare node unhealthy")
	flag.IntVar(&conf.HealthcheckHealthy, "hc-healthy", 1, "threshold of success checks to declare node healthy")
	flag.IntVar(&conf.HealthcheckTimeout, "hc-timeout", 5, "timeout of healthcheck endpoint, in seconds")
	flag.StringVar(&conf.ZipkinURL, "zipkin", "", "zipkin endpoint to send traces")
	flag.StringVar(&conf.Namespace, "namespace", "", "kubernetes namespace to monitor")
	flag.StringVar(&conf.LabelSelector, "label-selector", "", "kubernetes label selector to monitor")
	flag.IntVar(&conf.TargetPort, "target-port", 8080, "kubernetes port to target on selected pods")

	flag.Parse()

	conf.MinAPIVersion = semver.New(*minAPIVersion)

	if len(*fnodes) > 0 {
		// starting w/o nodes is fine too
		conf.Nodes = strings.Split(*fnodes, ",")
	}

	conf.Transport = &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   10 * time.Second,
			KeepAlive: 120 * time.Second,
		}).Dial,
		MaxIdleConnsPerHost: 512,
		TLSHandshakeTimeout: 10 * time.Second,
		TLSClientConfig: &tls.Config{
			ClientSessionCache: tls.NewLRUClientSessionCache(4096),
		},
	}

	db, err := lb.NewDB(conf) // Handles case where DBurl == "k8s"
	if err != nil {
		logrus.WithError(err).Fatal("error setting up database")
	}
	defer db.Close()

	g, err := lb.NewAllGrouper(conf, db)
	if err != nil {
		logrus.WithError(err).Fatal("error setting up grouper")
	}

	r := lb.NewConsistentRouter(conf)
	k := func(r *http.Request) (string, error) {
		return r.URL.Path, nil
	}

	servers := make([]*http.Server, 0, 1)
	handler := lb.NewProxy(k, g, r, conf)

	// a separate mgmt listener is requested? then let's create a LB traffic only server
	if conf.Listen != conf.MgmtListen {
		servers = append(servers, &http.Server{Addr: conf.Listen, Handler: handler})
		handler = lb.NullHandler()
	}

	// add mgmt endpoints to the handler
	handler = g.Wrap(handler) // add/del/list endpoints
	handler = r.Wrap(handler) // stats / dash endpoint

	servers = append(servers, &http.Server{Addr: conf.MgmtListen, Handler: handler})
	serve(servers, &conf)
}

func serve(servers []*http.Server, conf *lb.Config) {

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGQUIT, syscall.SIGINT)

	for i := 0; i < len(servers); i++ {
		go func(idx int) {
			err := servers[idx].ListenAndServe()
			if err != nil && err != http.ErrServerClosed {
				logrus.WithFields(logrus.Fields{"server_id": idx}).WithError(err).Fatal("server error")
			} else {
				logrus.WithFields(logrus.Fields{"server_id": idx}).Info("server stopped")
			}
		}(i)
	}

	sig := <-ch
	logrus.WithFields(logrus.Fields{"signal": sig}).Info("received signal")

	for i := 0; i < len(servers); i++ {

		ctx := context.Background()

		if conf.ShutdownTimeout > 0 {
			tmpCtx, cancel := context.WithTimeout(context.Background(), time.Duration(conf.ShutdownTimeout)*time.Second)
			ctx = tmpCtx
			defer cancel()
		}

		err := servers[i].Shutdown(ctx) // safe shutdown
		if err != nil {
			logrus.WithFields(logrus.Fields{"server_id": i}).WithError(err).Fatal("server shutdown error")
		} else {
			logrus.WithFields(logrus.Fields{"server_id": i}).Info("server shutdown")
		}
	}
}
<commit_msg>fnlb: 0.0.224 release [skip ci]<commit_after>package main

import (
	"context"
	"crypto/tls"
	"flag"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/coreos/go-semver/semver"
	"github.com/fnproject/fn/fnlb/lb"
	"github.com/sirupsen/logrus"
)

const VERSION = "0.0.224"

func main() {
	// XXX (reed): normalize
	level, err := logrus.ParseLevel(os.Getenv("LOG_LEVEL"))
	if err != nil {
		level = logrus.InfoLevel
	}
	logrus.SetLevel(level)

	fnodes := flag.String("nodes", "", "comma separated list of functions nodes")
	minAPIVersion := flag.String("min-api-version", "0.0.191", "minimal node API to accept")

	var conf lb.Config
	flag.StringVar(&conf.DBurl, "db", "sqlite3://:memory:", "backend to store nodes, default to in memory; use k8s for kuberneted")
	flag.StringVar(&conf.Listen, "listen", ":8081", "port to run on")
	flag.StringVar(&conf.MgmtListen, "mgmt-listen", ":8081", "management port to run on")
	flag.IntVar(&conf.ShutdownTimeout, "shutdown-timeout", 0, "graceful shutdown timeout")
	flag.IntVar(&conf.HealthcheckInterval, "hc-interval", 3, "how often to check f(x) nodes, in seconds")
	flag.StringVar(&conf.HealthcheckEndpoint, "hc-path", "/version", "endpoint to determine node health")
	flag.IntVar(&conf.HealthcheckUnhealthy, "hc-unhealthy", 2, "threshold of failed checks to declare node unhealthy")
	flag.IntVar(&conf.HealthcheckHealthy, "hc-healthy", 1, "threshold of success checks to declare node healthy")
	flag.IntVar(&conf.HealthcheckTimeout, "hc-timeout", 5, "timeout of healthcheck endpoint, in seconds")
	flag.StringVar(&conf.ZipkinURL, "zipkin", "", "zipkin endpoint to send traces")
	flag.StringVar(&conf.Namespace, "namespace", "", "kubernetes namespace to monitor")
	flag.StringVar(&conf.LabelSelector, "label-selector", "", "kubernetes label selector to monitor")
	flag.IntVar(&conf.TargetPort, "target-port", 8080, "kubernetes port to target on selected pods")

	flag.Parse()

	conf.MinAPIVersion = semver.New(*minAPIVersion)

	if len(*fnodes) > 0 {
		// starting w/o nodes is fine too
		conf.Nodes = strings.Split(*fnodes, ",")
	}

	conf.Transport = &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   10 * time.Second,
			KeepAlive: 120 * time.Second,
		}).Dial,
		MaxIdleConnsPerHost: 512,
		TLSHandshakeTimeout: 10 * time.Second,
		TLSClientConfig: &tls.Config{
			ClientSessionCache: tls.NewLRUClientSessionCache(4096),
		},
	}

	db, err := lb.NewDB(conf) // Handles case where DBurl == "k8s"
	if err != nil {
		logrus.WithError(err).Fatal("error setting up database")
	}
	defer db.Close()

	g, err := lb.NewAllGrouper(conf, db)
	if err != nil {
		logrus.WithError(err).Fatal("error setting up grouper")
	}

	r := lb.NewConsistentRouter(conf)
	k := func(r *http.Request) (string, error) {
		return r.URL.Path, nil
	}

	servers := make([]*http.Server, 0, 1)
	handler := lb.NewProxy(k, g, r, conf)

	// a separate mgmt listener is requested? then let's create a LB traffic only server
	if conf.Listen != conf.MgmtListen {
		servers = append(servers, &http.Server{Addr: conf.Listen, Handler: handler})
		handler = lb.NullHandler()
	}

	// add mgmt endpoints to the handler
	handler = g.Wrap(handler) // add/del/list endpoints
	handler = r.Wrap(handler) // stats / dash endpoint

	servers = append(servers, &http.Server{Addr: conf.MgmtListen, Handler: handler})
	serve(servers, &conf)
}

func serve(servers []*http.Server, conf *lb.Config) {

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGQUIT, syscall.SIGINT)

	for i := 0; i < len(servers); i++ {
		go func(idx int) {
			err := servers[idx].ListenAndServe()
			if err != nil && err != http.ErrServerClosed {
				logrus.WithFields(logrus.Fields{"server_id": idx}).WithError(err).Fatal("server error")
			} else {
				logrus.WithFields(logrus.Fields{"server_id": idx}).Info("server stopped")
			}
		}(i)
	}

	sig := <-ch
	logrus.WithFields(logrus.Fields{"signal": sig}).Info("received signal")

	for i := 0; i < len(servers); i++ {

		ctx := context.Background()

		if conf.ShutdownTimeout > 0 {
			tmpCtx, cancel := context.WithTimeout(context.Background(), time.Duration(conf.ShutdownTimeout)*time.Second)
			ctx = tmpCtx
			defer cancel()
		}

		err := servers[i].Shutdown(ctx) // safe shutdown
		if err != nil {
			logrus.WithFields(logrus.Fields{"server_id": i}).WithError(err).Fatal("server shutdown error")
		} else {
			logrus.WithFields(logrus.Fields{"server_id": i}).Info("server shutdown")
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"crypto/tls"
	"flag"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/coreos/go-semver/semver"
	"github.com/fnproject/fn/fnlb/lb"
	"github.com/sirupsen/logrus"
)

const VERSION = "0.0.70"

func main() {
	// XXX (reed): normalize
	fnodes := flag.String("nodes", "", "comma separated list of functions nodes")
	minAPIVersion := flag.String("min-api-version", "0.0.37", "minimal node API to accept")

	var conf lb.Config
	flag.StringVar(&conf.DBurl, "db", "sqlite3://:memory:", "backend to store nodes, default to in memory")
	flag.StringVar(&conf.Listen, "listen", ":8081", "port to run on")
	flag.IntVar(&conf.HealthcheckInterval, "hc-interval", 3, "how often to check f(x) nodes, in seconds")
	flag.StringVar(&conf.HealthcheckEndpoint, "hc-path", "/version", "endpoint to determine node health")
	flag.IntVar(&conf.HealthcheckUnhealthy, "hc-unhealthy", 2, "threshold of failed checks to declare node unhealthy")
	flag.IntVar(&conf.HealthcheckTimeout, "hc-timeout", 5, "timeout of healthcheck endpoint, in seconds")
	flag.StringVar(&conf.ZipkinURL, "zipkin", "", "zipkin endpoint to send traces")
	flag.Parse()

	conf.MinAPIVersion = semver.New(*minAPIVersion)

	if len(*fnodes) > 0 {
		// starting w/o nodes is fine too
		conf.Nodes = strings.Split(*fnodes, ",")
	}

	conf.Transport = &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   10 * time.Second,
			KeepAlive: 120 * time.Second,
		}).Dial,
		MaxIdleConnsPerHost: 512,
		TLSHandshakeTimeout: 10 * time.Second,
		TLSClientConfig: &tls.Config{
			ClientSessionCache: tls.NewLRUClientSessionCache(4096),
		},
	}

	g, err := lb.NewAllGrouper(conf)
	if err != nil {
		logrus.WithError(err).Fatal("error setting up grouper")
	}

	r := lb.NewConsistentRouter(conf)
	k := func(r *http.Request) (string, error) {
		return r.URL.Path, nil
	}

	h := lb.NewProxy(k, g, r, conf)
	h = g.Wrap(h) // add/del/list endpoints
	h = r.Wrap(h) // stats / dash endpoint

	err = serve(conf.Listen, h)
	if err != nil {
		logrus.WithError(err).Fatal("server error")
	}
}

func serve(addr string, handler http.Handler) error {
	server := &http.Server{Addr: addr, Handler: handler}

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGQUIT, syscall.SIGINT)
	go func() {
		for sig := range ch {
			logrus.WithFields(logrus.Fields{"signal": sig}).Info("received signal")
			server.Shutdown(context.Background()) // safe shutdown
			return
		}
	}()
	return server.ListenAndServe()
}
<commit_msg>fnlb: 0.0.71 release [skip ci]<commit_after>package main

import (
	"context"
	"crypto/tls"
	"flag"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/coreos/go-semver/semver"
	"github.com/fnproject/fn/fnlb/lb"
	"github.com/sirupsen/logrus"
)

const VERSION = "0.0.71"

func main() {
	// XXX (reed): normalize
	fnodes := flag.String("nodes", "", "comma separated list of functions nodes")
	minAPIVersion := flag.String("min-api-version", "0.0.38", "minimal node API to accept")

	var conf lb.Config
	flag.StringVar(&conf.DBurl, "db", "sqlite3://:memory:", "backend to store nodes, default to in memory")
	flag.StringVar(&conf.Listen, "listen", ":8081", "port to run on")
	flag.IntVar(&conf.HealthcheckInterval, "hc-interval", 3, "how often to check f(x) nodes, in seconds")
	flag.StringVar(&conf.HealthcheckEndpoint, "hc-path", "/version", "endpoint to determine node health")
	flag.IntVar(&conf.HealthcheckUnhealthy, "hc-unhealthy", 2, "threshold of failed checks to declare node unhealthy")
	flag.IntVar(&conf.HealthcheckTimeout, "hc-timeout", 5, "timeout of healthcheck endpoint, in seconds")
	flag.StringVar(&conf.ZipkinURL, "zipkin", "", "zipkin endpoint to send traces")
	flag.Parse()

	conf.MinAPIVersion = semver.New(*minAPIVersion)

	if len(*fnodes) > 0 {
		// starting w/o nodes is fine too
		conf.Nodes = strings.Split(*fnodes, ",")
	}

	conf.Transport = &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   10 * time.Second,
			KeepAlive: 120 * time.Second,
		}).Dial,
		MaxIdleConnsPerHost: 512,
		TLSHandshakeTimeout: 10 * time.Second,
		TLSClientConfig: &tls.Config{
			ClientSessionCache: tls.NewLRUClientSessionCache(4096),
		},
	}

	g, err := lb.NewAllGrouper(conf)
	if err != nil {
		logrus.WithError(err).Fatal("error setting up grouper")
	}

	r := lb.NewConsistentRouter(conf)
	k := func(r *http.Request) (string, error) {
		return r.URL.Path, nil
	}

	h := lb.NewProxy(k, g, r, conf)
	h = g.Wrap(h) // add/del/list endpoints
	h = r.Wrap(h) // stats / dash endpoint

	err = serve(conf.Listen, h)
	if err != nil {
		logrus.WithError(err).Fatal("server error")
	}
}

func serve(addr string, handler http.Handler) error {
	server := &http.Server{Addr: addr, Handler: handler}

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGQUIT, syscall.SIGINT)
	go func() {
		for sig := range ch {
			logrus.WithFields(logrus.Fields{"signal": sig}).Info("received signal")
			server.Shutdown(context.Background()) // safe shutdown
			return
		}
	}()
	return server.ListenAndServe()
}
<|endoftext|>