<commit_before>//
// Copyright 2021, Sander van Harmelen
//
// 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 gitlab

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

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestBlockUser(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/block", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusCreated)
	})

	err := client.Users.BlockUser(1)
	if err != nil {
		t.Errorf("Users.BlockUser returned error: %v", err)
	}
}

func TestBlockUser_UserNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/block", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusNotFound)
	})

	err := client.Users.BlockUser(1)
	if !errors.Is(err, ErrUserNotFound) {
		t.Errorf("Users.BlockUser error.\nExpected: %+v\nGot: %+v", ErrUserNotFound, err)
	}
}

func TestBlockUser_BlockPrevented(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/block", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusForbidden)
	})

	err := client.Users.BlockUser(1)
	if !errors.Is(err, ErrUserBlockPrevented) {
		t.Errorf("Users.BlockUser error.\nExpected: %+v\nGot: %+v", ErrUserBlockPrevented, err)
	}
}

func TestBlockUser_UnknownError(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/block", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusTeapot)
	})

	want := fmt.Sprintf("Received unexpected result code: %d", http.StatusTeapot)

	err := client.Users.BlockUser(1)
	if err.Error() != want {
		t.Errorf("Users.BlockUser error.\nExpected: %s\nGot: %v", want, err)
	}
}

//  ------------------------  Unblock user -------------------------
func TestUnblockUser(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/unblock", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusCreated)
	})

	err := client.Users.UnblockUser(1)
	if err != nil {
		t.Errorf("Users.UnblockUser returned error: %v", err)
	}
}

func TestUnblockUser_UserNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/unblock", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusNotFound)
	})

	err := client.Users.UnblockUser(1)
	if !errors.Is(err, ErrUserNotFound) {
		t.Errorf("Users.UnblockUser error.\nExpected: %v\nGot: %v", ErrUserNotFound, err)
	}
}

func TestUnblockUser_UnblockPrevented(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/unblock", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusForbidden)
	})

	err := client.Users.UnblockUser(1)
	if !errors.Is(err, ErrUserUnblockPrevented) {
		t.Errorf("Users.UnblockUser error.\nExpected: %v\nGot: %v", ErrUserUnblockPrevented, err)
	}
}

func TestUnblockUser_UnknownError(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/unblock", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusTeapot)
	})

	want := fmt.Sprintf("Received unexpected result code: %d", http.StatusTeapot)

	err := client.Users.UnblockUser(1)
	if err.Error() != want {
		t.Errorf("Users.UnblockUser error.\nExpected: %s\n\tGot: %v", want, err)
	}
}

func TestDeactivateUser(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/deactivate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusCreated)
	})

	err := client.Users.DeactivateUser(1)
	if err != nil {
		t.Errorf("Users.DeactivateUser returned error: %v", err)
	}
}

func TestDeactivateUser_UserNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/deactivate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusNotFound)
	})

	err := client.Users.DeactivateUser(1)
	if !errors.Is(err, ErrUserNotFound) {
		t.Errorf("Users.DeactivateUser error.\nExpected: %+v\n\tGot: %+v", ErrUserNotFound, err)
	}
}

func TestDeactivateUser_DeactivatePrevented(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/deactivate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusForbidden)
	})

	err := client.Users.DeactivateUser(1)
	if !errors.Is(err, ErrUserDeactivatePrevented) {
		t.Errorf("Users.DeactivateUser error.\nExpected: %+v\n\tGot: %+v", ErrUserDeactivatePrevented, err)
	}
}

func TestActivateUser(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/activate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusCreated)
	})

	err := client.Users.ActivateUser(1)
	if err != nil {
		t.Errorf("Users.ActivateUser returned error: %v", err)
	}
}

func TestActivateUser_ActivatePrevented(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/activate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusForbidden)
	})

	err := client.Users.ActivateUser(1)
	if !errors.Is(err, ErrUserActivatePrevented) {
		t.Errorf("Users.ActivateUser error.\nExpected: %+v\n\tGot: %+v", ErrUserActivatePrevented, err)
	}
}

func TestActivateUser_UserNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/activate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusNotFound)
	})

	err := client.Users.ActivateUser(1)
	if !errors.Is(err, ErrUserNotFound) {
		t.Errorf("Users.ActivateUser error.\nExpected: %+v\n\tGot: %+v", ErrUserNotFound, err)
	}
}

func TestGetMemberships(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/memberships", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		mustWriteHTTPResponse(t, w, "testdata/get_user_memberships.json")
	})

	opt := new(GetUserMembershipOptions)

	memberships, _, err := client.Users.GetUserMemberships(1, opt)
	require.NoError(t, err)

	want := []*UserMembership{{SourceID: 1, SourceName: "Project one", SourceType: "Project", AccessLevel: 20}, {SourceID: 3, SourceName: "Group three", SourceType: "Namespace", AccessLevel: 20}}
	assert.Equal(t, want, memberships)
}
<commit_msg>Add tests for Approve and Reject functions<commit_after>//
// Copyright 2021, Sander van Harmelen
//
// 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 gitlab

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

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestBlockUser(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/block", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusCreated)
	})

	err := client.Users.BlockUser(1)
	if err != nil {
		t.Errorf("Users.BlockUser returned error: %v", err)
	}
}

func TestBlockUser_UserNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/block", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusNotFound)
	})

	err := client.Users.BlockUser(1)
	if !errors.Is(err, ErrUserNotFound) {
		t.Errorf("Users.BlockUser error.\nExpected: %+v\nGot: %+v", ErrUserNotFound, err)
	}
}

func TestBlockUser_BlockPrevented(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/block", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusForbidden)
	})

	err := client.Users.BlockUser(1)
	if !errors.Is(err, ErrUserBlockPrevented) {
		t.Errorf("Users.BlockUser error.\nExpected: %+v\nGot: %+v", ErrUserBlockPrevented, err)
	}
}

func TestBlockUser_UnknownError(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/block", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusTeapot)
	})

	want := fmt.Sprintf("Received unexpected result code: %d", http.StatusTeapot)

	err := client.Users.BlockUser(1)
	if err.Error() != want {
		t.Errorf("Users.BlockUser error.\nExpected: %s\nGot: %v", want, err)
	}
}

//  ------------------------  Unblock user -------------------------
func TestUnblockUser(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/unblock", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusCreated)
	})

	err := client.Users.UnblockUser(1)
	if err != nil {
		t.Errorf("Users.UnblockUser returned error: %v", err)
	}
}

func TestUnblockUser_UserNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/unblock", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusNotFound)
	})

	err := client.Users.UnblockUser(1)
	if !errors.Is(err, ErrUserNotFound) {
		t.Errorf("Users.UnblockUser error.\nExpected: %v\nGot: %v", ErrUserNotFound, err)
	}
}

func TestUnblockUser_UnblockPrevented(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/unblock", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusForbidden)
	})

	err := client.Users.UnblockUser(1)
	if !errors.Is(err, ErrUserUnblockPrevented) {
		t.Errorf("Users.UnblockUser error.\nExpected: %v\nGot: %v", ErrUserUnblockPrevented, err)
	}
}

func TestUnblockUser_UnknownError(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/unblock", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusTeapot)
	})

	want := fmt.Sprintf("Received unexpected result code: %d", http.StatusTeapot)

	err := client.Users.UnblockUser(1)
	if err.Error() != want {
		t.Errorf("Users.UnblockUser error.\nExpected: %s\n\tGot: %v", want, err)
	}
}

func TestDeactivateUser(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/deactivate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusCreated)
	})

	err := client.Users.DeactivateUser(1)
	if err != nil {
		t.Errorf("Users.DeactivateUser returned error: %v", err)
	}
}

func TestDeactivateUser_UserNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/deactivate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusNotFound)
	})

	err := client.Users.DeactivateUser(1)
	if !errors.Is(err, ErrUserNotFound) {
		t.Errorf("Users.DeactivateUser error.\nExpected: %+v\n\tGot: %+v", ErrUserNotFound, err)
	}
}

func TestDeactivateUser_DeactivatePrevented(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/deactivate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusForbidden)
	})

	err := client.Users.DeactivateUser(1)
	if !errors.Is(err, ErrUserDeactivatePrevented) {
		t.Errorf("Users.DeactivateUser error.\nExpected: %+v\n\tGot: %+v", ErrUserDeactivatePrevented, err)
	}
}

func TestActivateUser(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/activate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusCreated)
	})

	err := client.Users.ActivateUser(1)
	if err != nil {
		t.Errorf("Users.ActivateUser returned error: %v", err)
	}
}

func TestActivateUser_ActivatePrevented(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/activate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusForbidden)
	})

	err := client.Users.ActivateUser(1)
	if !errors.Is(err, ErrUserActivatePrevented) {
		t.Errorf("Users.ActivateUser error.\nExpected: %+v\n\tGot: %+v", ErrUserActivatePrevented, err)
	}
}

func TestActivateUser_UserNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/activate", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusNotFound)
	})

	err := client.Users.ActivateUser(1)
	if !errors.Is(err, ErrUserNotFound) {
		t.Errorf("Users.ActivateUser error.\nExpected: %+v\n\tGot: %+v", ErrUserNotFound, err)
	}
}

func TestGetMemberships(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/memberships", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		mustWriteHTTPResponse(t, w, "testdata/get_user_memberships.json")
	})

	opt := new(GetUserMembershipOptions)

	memberships, _, err := client.Users.GetUserMemberships(1, opt)
	require.NoError(t, err)

	want := []*UserMembership{{SourceID: 1, SourceName: "Project one", SourceType: "Project", AccessLevel: 20}, {SourceID: 3, SourceName: "Group three", SourceType: "Namespace", AccessLevel: 20}}
	assert.Equal(t, want, memberships)
}

//  ------------------------  Approve user -------------------------
func TestApproveUser(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/approve", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusCreated)
	})

	err := client.Users.ApproveUser(1)
	if err != nil {
		t.Errorf("Users.ApproveUser returned error: %v", err)
	}
}

func TestApproveUser_UserNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/approve", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusNotFound)
	})

	err := client.Users.ApproveUser(1)
	if !errors.Is(err, ErrUserNotFound) {
		t.Errorf("Users.ApproveUser error.\nExpected: %v\nGot: %v", ErrUserNotFound, err)
	}
}

func TestApproveUser_ApprovePrevented(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/approve", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusForbidden)
	})

	err := client.Users.ApproveUser(1)
	if !errors.Is(err, ErrUserApprovePrevented) {
		t.Errorf("Users.ApproveUser error.\nExpected: %v\nGot: %v", ErrUserApprovePrevented, err)
	}
}

func TestApproveUser_UnknownError(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/approve", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusTeapot)
	})

	want := fmt.Sprintf("Received unexpected result code: %d", http.StatusTeapot)

	err := client.Users.ApproveUser(1)
	if err.Error() != want {
		t.Errorf("Users.ApproveUser error.\nExpected: %s\n\tGot: %v", want, err)
	}
}

//  ------------------------  Reject user -------------------------
func TestRejectUser(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/reject", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusOK)
	})

	err := client.Users.RejectUser(1)
	if err != nil {
		t.Errorf("Users.RejectUser returned error: %v", err)
	}
}

func TestRejectUser_UserNotFound(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/reject", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusNotFound)
	})

	err := client.Users.RejectUser(1)
	if !errors.Is(err, ErrUserNotFound) {
		t.Errorf("Users.RejectUser error.\nExpected: %v\nGot: %v", ErrUserNotFound, err)
	}
}

func TestRejectUser_RejectPrevented(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/reject", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusForbidden)
	})

	err := client.Users.RejectUser(1)
	if !errors.Is(err, ErrUserRejectPrevented) {
		t.Errorf("Users.RejectUser error.\nExpected: %v\nGot: %v", ErrUserRejectPrevented, err)
	}
}

func TestRejectUser_Conflict(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/reject", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusConflict)
	})

	err := client.Users.RejectUser(1)
	if !errors.Is(err, ErrUserConflict) {
		t.Errorf("Users.RejectUser error.\nExpected: %v\nGot: %v", ErrUserConflict, err)
	}
}

func TestRejectUser_UnknownError(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	path := fmt.Sprintf("/%susers/1/reject", apiVersionPath)
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		w.WriteHeader(http.StatusTeapot)
	})

	want := fmt.Sprintf("Received unexpected result code: %d", http.StatusTeapot)

	err := client.Users.RejectUser(1)
	if err.Error() != want {
		t.Errorf("Users.RejectUser error.\nExpected: %s\n\tGot: %v", want, err)
	}
}
<|endoftext|>


<commit_before>package checks

import (
	"fmt"
	"sync"
	"time"

	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/consul/types"
)

// Constants related to alias check backoff.
const (
	checkAliasBackoffMin     = 3               // 3 attempts before backing off
	checkAliasBackoffMaxWait = 1 * time.Minute // maximum backoff wait time
)

// CheckAlias is a check type that aliases the health of another service
// instance or node. If the service aliased has any critical health checks, then
// this check is critical. If the service has no critical but warnings,
// then this check is warning, and if a service has only passing checks, then
// this check is passing.
type CheckAlias struct {
	Node      string // Node name of the service. If empty, assumed to be this node.
	ServiceID string // ID (not name) of the service to alias

	CheckID types.CheckID               // ID of this check
	RPC     RPC                         // Used to query remote server if necessary
	RPCReq  structs.NodeSpecificRequest // Base request
	Notify  AliasNotifier               // For updating the check state

	stop     bool
	stopCh   chan struct{}
	stopLock sync.Mutex
}

// AliasNotifier is a CheckNotifier specifically for the Alias check.
// This requires additional methods that are satisfied by the agent
// local state.
type AliasNotifier interface {
	CheckNotifier

	AddAliasCheck(types.CheckID, string, chan<- struct{}) error
	RemoveAliasCheck(types.CheckID, string)
	Checks() map[types.CheckID]*structs.HealthCheck
}

// Start is used to start the check, runs until Stop() func (c *CheckAlias) Start() {
func (c *CheckAlias) Start() {
	c.stopLock.Lock()
	defer c.stopLock.Unlock()
	c.stop = false
	c.stopCh = make(chan struct{})
	go c.run(c.stopCh)
}

// Stop is used to stop the check.
func (c *CheckAlias) Stop() {
	c.stopLock.Lock()
	defer c.stopLock.Unlock()
	if !c.stop {
		c.stop = true
		close(c.stopCh)
	}
}

// run is invoked in a goroutine until Stop() is called.
func (c *CheckAlias) run(stopCh chan struct{}) {
	// If we have a specific node set, then use a blocking query
	if c.Node != "" {
		c.runQuery(stopCh)
		return
	}

	// Use the local state to match the service.
	c.runLocal(stopCh)
}

func (c *CheckAlias) runLocal(stopCh chan struct{}) {
	// Very important this is buffered as 1 so that we do not lose any
	// queued updates. This only has to be exactly 1 since the existence
	// of any update triggers us to load the full health check state.
	notifyCh := make(chan struct{}, 1)
	c.Notify.AddAliasCheck(c.CheckID, c.ServiceID, notifyCh)
	defer c.Notify.RemoveAliasCheck(c.CheckID, c.ServiceID)

	for {
		select {
		case <-notifyCh:
			checks := c.Notify.Checks()
			checksList := make([]*structs.HealthCheck, 0, len(checks))
			for _, chk := range checks {
				checksList = append(checksList, chk)
			}
			c.processChecks(checksList)

		case <-stopCh:
			return
		}
	}
}

func (c *CheckAlias) runQuery(stopCh chan struct{}) {
	args := c.RPCReq
	args.Node = c.Node
	args.AllowStale = true
	args.MaxQueryTime = 1 * time.Minute

	var attempt uint
	for {
		// Check if we're stopped. We fallthrough and block otherwise,
		// which has a maximum time set above so we'll always check for
		// stop within a reasonable amount of time.
		select {
		case <-stopCh:
			return
		default:
		}

		// Backoff if we have to
		if attempt > checkAliasBackoffMin {
			waitTime := (1 << (attempt - checkAliasBackoffMin)) * time.Second
			if waitTime > checkAliasBackoffMaxWait {
				waitTime = checkAliasBackoffMaxWait
			}
			time.Sleep(waitTime)
		}

		// Get the current health checks for the specified node.
		//
		// NOTE(mitchellh): This currently returns ALL health checks for
		// a node even though we also have the service ID. This can be
		// optimized if we introduce a new RPC endpoint to filter both,
		// but for blocking queries isn't that much more efficient since the checks
		// index is global to the cluster.
		var out structs.IndexedHealthChecks
		if err := c.RPC.RPC("Health.NodeChecks", &args, &out); err != nil {
			attempt++
			if attempt > 1 {
				c.Notify.UpdateCheck(c.CheckID, api.HealthCritical,
					fmt.Sprintf("Failure checking aliased node or service: %s", err))
			}

			continue
		}

		attempt = 0 // Reset the attempts so we don't backoff the next

		// Set our index for the next request
		args.MinQueryIndex = out.Index

		// We want to ensure that we're always blocking on subsequent requests
		// to avoid hot loops. Index 1 is always safe since the min raft index
		// is at least 5. Note this shouldn't happen but protecting against this
		// case is safer than a 100% CPU loop.
		if args.MinQueryIndex < 1 {
			args.MinQueryIndex = 1
		}

		c.processChecks(out.HealthChecks)
	}
}

// processChecks is a common helper for taking a set of health checks and
// using them to update our alias. This is abstracted since the checks can
// come from both the remote server as well as local state.
func (c *CheckAlias) processChecks(checks []*structs.HealthCheck) {
	health := api.HealthPassing
	msg := "No checks found."
	for _, chk := range checks {
		if c.Node != "" && chk.Node != c.Node {
			continue
		}

		// We allow ServiceID == "" so that we also check node checks
		if chk.ServiceID != "" && chk.ServiceID != c.ServiceID {
			continue
		}

		if chk.Status == api.HealthCritical || chk.Status == api.HealthWarning {
			health = chk.Status
			msg = fmt.Sprintf("Aliased check %q failing: %s", chk.Name, chk.Output)

			// Critical checks exit the for loop immediately since we
			// know that this is the health state. Warnings do not since
			// there may still be a critical check.
			if chk.Status == api.HealthCritical {
				break
			}
		}

		msg = "All checks passing."
	}

	// Update our check value
	c.Notify.UpdateCheck(c.CheckID, health, msg)
}
<commit_msg>agent/checks: prevent overflow of backoff<commit_after>package checks

import (
	"fmt"
	"sync"
	"time"

	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/consul/types"
)

// Constants related to alias check backoff.
const (
	checkAliasBackoffMin     = 3               // 3 attempts before backing off
	checkAliasBackoffMaxWait = 1 * time.Minute // maximum backoff wait time
)

// CheckAlias is a check type that aliases the health of another service
// instance or node. If the service aliased has any critical health checks, then
// this check is critical. If the service has no critical but warnings,
// then this check is warning, and if a service has only passing checks, then
// this check is passing.
type CheckAlias struct {
	Node      string // Node name of the service. If empty, assumed to be this node.
	ServiceID string // ID (not name) of the service to alias

	CheckID types.CheckID               // ID of this check
	RPC     RPC                         // Used to query remote server if necessary
	RPCReq  structs.NodeSpecificRequest // Base request
	Notify  AliasNotifier               // For updating the check state

	stop     bool
	stopCh   chan struct{}
	stopLock sync.Mutex
}

// AliasNotifier is a CheckNotifier specifically for the Alias check.
// This requires additional methods that are satisfied by the agent
// local state.
type AliasNotifier interface {
	CheckNotifier

	AddAliasCheck(types.CheckID, string, chan<- struct{}) error
	RemoveAliasCheck(types.CheckID, string)
	Checks() map[types.CheckID]*structs.HealthCheck
}

// Start is used to start the check, runs until Stop() func (c *CheckAlias) Start() {
func (c *CheckAlias) Start() {
	c.stopLock.Lock()
	defer c.stopLock.Unlock()
	c.stop = false
	c.stopCh = make(chan struct{})
	go c.run(c.stopCh)
}

// Stop is used to stop the check.
func (c *CheckAlias) Stop() {
	c.stopLock.Lock()
	defer c.stopLock.Unlock()
	if !c.stop {
		c.stop = true
		close(c.stopCh)
	}
}

// run is invoked in a goroutine until Stop() is called.
func (c *CheckAlias) run(stopCh chan struct{}) {
	// If we have a specific node set, then use a blocking query
	if c.Node != "" {
		c.runQuery(stopCh)
		return
	}

	// Use the local state to match the service.
	c.runLocal(stopCh)
}

func (c *CheckAlias) runLocal(stopCh chan struct{}) {
	// Very important this is buffered as 1 so that we do not lose any
	// queued updates. This only has to be exactly 1 since the existence
	// of any update triggers us to load the full health check state.
	notifyCh := make(chan struct{}, 1)
	c.Notify.AddAliasCheck(c.CheckID, c.ServiceID, notifyCh)
	defer c.Notify.RemoveAliasCheck(c.CheckID, c.ServiceID)

	for {
		select {
		case <-notifyCh:
			checks := c.Notify.Checks()
			checksList := make([]*structs.HealthCheck, 0, len(checks))
			for _, chk := range checks {
				checksList = append(checksList, chk)
			}
			c.processChecks(checksList)

		case <-stopCh:
			return
		}
	}
}

func (c *CheckAlias) runQuery(stopCh chan struct{}) {
	args := c.RPCReq
	args.Node = c.Node
	args.AllowStale = true
	args.MaxQueryTime = 1 * time.Minute

	var attempt uint
	for {
		// Check if we're stopped. We fallthrough and block otherwise,
		// which has a maximum time set above so we'll always check for
		// stop within a reasonable amount of time.
		select {
		case <-stopCh:
			return
		default:
		}

		// Backoff if we have to
		if attempt > checkAliasBackoffMin {
			shift := attempt - checkAliasBackoffMin
			if shift > 31 {
				shift = 31 // so we don't overflow to 0
			}
			waitTime := (1 << shift) * time.Second
			if waitTime > checkAliasBackoffMaxWait {
				waitTime = checkAliasBackoffMaxWait
			}
			time.Sleep(waitTime)
		}

		// Get the current health checks for the specified node.
		//
		// NOTE(mitchellh): This currently returns ALL health checks for
		// a node even though we also have the service ID. This can be
		// optimized if we introduce a new RPC endpoint to filter both,
		// but for blocking queries isn't that much more efficient since the checks
		// index is global to the cluster.
		var out structs.IndexedHealthChecks
		if err := c.RPC.RPC("Health.NodeChecks", &args, &out); err != nil {
			attempt++
			if attempt > 1 {
				c.Notify.UpdateCheck(c.CheckID, api.HealthCritical,
					fmt.Sprintf("Failure checking aliased node or service: %s", err))
			}

			continue
		}

		attempt = 0 // Reset the attempts so we don't backoff the next

		// Set our index for the next request
		args.MinQueryIndex = out.Index

		// We want to ensure that we're always blocking on subsequent requests
		// to avoid hot loops. Index 1 is always safe since the min raft index
		// is at least 5. Note this shouldn't happen but protecting against this
		// case is safer than a 100% CPU loop.
		if args.MinQueryIndex < 1 {
			args.MinQueryIndex = 1
		}

		c.processChecks(out.HealthChecks)
	}
}

// processChecks is a common helper for taking a set of health checks and
// using them to update our alias. This is abstracted since the checks can
// come from both the remote server as well as local state.
func (c *CheckAlias) processChecks(checks []*structs.HealthCheck) {
	health := api.HealthPassing
	msg := "No checks found."
	for _, chk := range checks {
		if c.Node != "" && chk.Node != c.Node {
			continue
		}

		// We allow ServiceID == "" so that we also check node checks
		if chk.ServiceID != "" && chk.ServiceID != c.ServiceID {
			continue
		}

		if chk.Status == api.HealthCritical || chk.Status == api.HealthWarning {
			health = chk.Status
			msg = fmt.Sprintf("Aliased check %q failing: %s", chk.Name, chk.Output)

			// Critical checks exit the for loop immediately since we
			// know that this is the health state. Warnings do not since
			// there may still be a critical check.
			if chk.Status == api.HealthCritical {
				break
			}
		}

		msg = "All checks passing."
	}

	// Update our check value
	c.Notify.UpdateCheck(c.CheckID, health, msg)
}
<|endoftext|>


<commit_before>package api

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"sync"

	"github.com/koding/logging"
)

const maxBodyLen = 2 << 20 // 2 MiB

var defaultLog = logging.NewCustom("api-transport", false)

// contextKey is a value for use with context.WithValue.
type contextKey struct {
	name string
}

// UserContextKey is a context key. It can be used in
// HTTP handlers to attach a client session to a request.
//
// The session is then read by the Transport and used
// for authentication.
//
// The associated value will be of *User type.
var UserContextKey = &contextKey{"client-session"}

// UserFunc is used to extract a user, which requires authentication,
// from a request.
type UserFunc func(*http.Request) *User

// User represents a user, which is going to be authenticated.
type User struct {
	Username string // name of the user
	Team     string // team which the user belongs to
}

// WithRequest attaches a user to the given request.
func (u *User) WithRequest(req *http.Request) *http.Request {
	return req.WithContext(context.WithValue(req.Context(), UserContextKey, u))
}

// Strings implements the fmt.Stringer interface.
func (u *User) String() string {
	return u.Team + "/" + u.Username
}

// Session represents a user session.
type Session struct {
	ClientID string // an ID of an authenticated session.
	User     *User  // authenticated user
}

// Valid implements the stack.Validator interface.
func (s *Session) Valid() error {
	if s == nil {
		return errors.New("session is missing")
	}
	if s.ClientID == "" {
		return errors.New("empty client ID")
	}
	return nil
}

// WriteTo writes the s session to the given requeest.
func (s *Session) WriteTo(req *http.Request) {
	req.Header.Set("Authorization", "Bearer "+s.ClientID)
}

// ReadFrom initializes the s session by reading it from
// the given request.
//
// If no session can be read from req, the method is a nop.
func (s *Session) ReadFrom(req *http.Request) {
	if auth := req.Header.Get("Authorization"); strings.HasPrefix(auth, "Bearer ") {
		s.ClientID = auth[len("Bearer "):]
	}
}

// Match tests whether the other session matches the s one.
//
// Since ClientID is typically auto-generated, we match all the other fields.
func (s *Session) Match(other *Session) error {
	if s.User.Username != other.User.Username {
		return fmt.Errorf("username is  %q, the other one is %q", s.User.Username, other.User.Username)
	}

	if s.User.Team != other.User.Team {
		return fmt.Errorf("team is  %q, the other one is %q", s.User.Team, other.User.Team)
	}

	return nil
}

// AuthFunc is used to fetch session information.
//
// When Refresh is true, AuthFunc must obtain
// a session directly from an authorisation
// endpoint, in case it was cached.
type AuthFunc func(*AuthOptions) (*Session, error)

// AuthOptions represents the arguments for AuthFunc.
type AuthOptions struct {
	// User represents a user, which we are building
	// authorisation for.
	User *User

	// Refresh requests to refresh a session when true.
	Refresh bool

	// Request is a copy of the request, which
	// we're building authorisation for.
	Request *http.Request

	// Error is the error that was the reason
	// of temporary failure received while trying
	// to sending request.
	//
	// Error is non-nil when request failed with
	// temporary error and Transport is retrying.
	Error error
}

type httpTransport interface {
	http.RoundTripper
	httpRequestCanceler
	httpIdleConnectionsCloser
}

type httpRequestCanceler interface {
	CancelRequest(*http.Request)
}

type httpIdleConnectionsCloser interface {
	CloseIdleConnections()
}

// nopCloser is like ioutil.NopCloser, but wraps io.ReadSeeker.
type nopCloser struct{ io.ReadSeeker }

func (nopCloser) Close() (_ error) { return }

var _ io.ReadSeeker = nopCloser{bytes.NewReader(nil)}

// Transport is a signing transport that
// authorizes each request with jSession.clientId.
type Transport struct {
	http.RoundTripper                // a transport to use; optional, by default http.DefaultClient.Transport
	AuthFunc          AuthFunc       // clientID authorisation to use; required
	UserFunc          UserFunc       // used to get a user from request to authenticate for; by default looked up with UserContextKey
	Host              string         // original Host name, overwrites req.Host; optional
	RetryNum          int            // number of retries in case of temporary failures; optional, by default 3
	Log               logging.Logger // logger to use
	Debug             bool           // whether to dump HTTP request/response

	// request mapping implementation stolen from:
	//
	//   https://github.com/golang/oauth2/blob/master/transport.go
	//
	mu     sync.Mutex                      // guards modReq
	modReq map[*http.Request]*http.Request // original -> modified
}

var _ httpTransport = (*Transport)(nil)

// NewSingleUser returns a HTTP transport that authenticates all
// requests for the given user.
//
// This is usefull for api clients that give no control over HTTP requests,
// like the swagger-generated ones (remote.api).
func (t *Transport) NewSingleUser(u *User) http.RoundTripper {
	return &Transport{
		RoundTripper: t.RoundTripper,
		AuthFunc:     t.AuthFunc,
		UserFunc:     func(*http.Request) *User { return u },
		Host:         t.Host,
		RetryNum:     t.RetryNum,
		Log:          t.Log,
		Debug:        t.Debug,
	}
}

// RoundTrip implements the http.RoundTripper interface.
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	user := t.user(req)
	if user == nil {
		return t.roundTripper().RoundTrip(req)
	}

	reqCopy := copyRequest(req) // per RoundTripper contract
	t.setModReq(req, reqCopy)

	switch reqCopy.Body.(type) {
	case nil:
	case io.ReadSeeker:
	default:
		p, err := ioutil.ReadAll(io.LimitReader(reqCopy.Body, maxBodyLen))
		if err != nil {
			return nil, err // non-retryable
		}

		reqCopy.Body = nopCloser{bytes.NewReader(p)}
	}

	var refresh bool
	var lastErr error
	var session *Session

	for attempts := t.retryNum(); attempts >= 0; attempts-- {
		if reqCopy.Body != nil {
			if _, err := reqCopy.Body.(io.ReadSeeker).Seek(0, io.SeekStart); err != nil {
				return nil, err // non-retryable
			}
		}

		opts := &AuthOptions{
			User:    user,
			Refresh: refresh,
			Request: reqCopy,
			Error:   lastErr,
		}

		s, err := t.AuthFunc(opts)

		t.log().Debug("AuthFunc(%+v): session=%+v, error=%v", opts, s, err)

		if err != nil {
			return nil, err // non-retryable
		}

		if err := s.Valid(); err != nil {
			return nil, err // non-retryable
		}

		session = s
		refresh = false

		session.WriteTo(reqCopy)

		if t.Host != "" {
			reqCopy.Host = t.Host
		}

		if t.Debug {
			p, err := httputil.DumpRequestOut(reqCopy, true)
			t.log().Debug("[request] %s (err=%v)", p, err)
		}

		resp, err := t.roundTripper().RoundTrip(reqCopy)

		if t.Debug {
			p, err := httputil.DumpResponse(resp, true)
			t.log().Debug("[response] %s (err=%v)", p, err)
		}

		if e, ok := err.(net.Error); ok && (e.Temporary() || e.Timeout()) {
			lastErr = err
			continue // retry
		}

		if err != nil {
			t.setModReq(req, nil)
			return nil, err // non-retryable
		}

		if resp.StatusCode == http.StatusUnauthorized {
			resp.Body.Close()
			lastErr = errors.New(http.StatusText(resp.StatusCode))
			refresh = true
			continue // retry and force new session
		}

		resp.Body = &onEOFReader{
			rc: resp.Body,
			fn: func() { t.setModReq(req, nil) },
		}

		return resp, nil
	}

	t.setModReq(req, nil)

	return nil, fmt.Errorf("api: error sending request: %v", lastErr)
}

func (t *Transport) CancelRequest(req *http.Request) {
	if rc, ok := t.roundTripper().(httpRequestCanceler); ok {
		t.mu.Lock()
		modReq := t.modReq[req]
		delete(t.modReq, req)
		t.mu.Unlock()

		rc.CancelRequest(modReq)
	}
}

func (t *Transport) CloseIdleConnections() {
	if icl, ok := t.roundTripper().(httpIdleConnectionsCloser); ok {
		icl.CloseIdleConnections()
	}
}

func (t *Transport) user(req *http.Request) *User {
	if t.UserFunc != nil {
		return t.UserFunc(req)
	}
	user, _ := req.Context().Value(UserContextKey).(*User)
	return user
}

func (t *Transport) roundTripper() http.RoundTripper {
	if t.RoundTripper != nil {
		return t.RoundTripper
	}

	if http.DefaultClient.Transport != nil {
		return http.DefaultClient.Transport
	}

	return http.DefaultTransport
}

func (t *Transport) retryNum() int {
	if t.RetryNum > 0 {
		return t.RetryNum
	}
	return 3
}

func (t *Transport) log() logging.Logger {
	if t.Log != nil {
		return t.Log
	}
	return defaultLog
}

func (t *Transport) setModReq(orig, mod *http.Request) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.modReq == nil {
		t.modReq = make(map[*http.Request]*http.Request)
	}
	if mod == nil {
		delete(t.modReq, orig)
	} else {
		t.modReq[orig] = mod
	}
}

func copyURL(u *url.URL) *url.URL {
	uCopy := *u
	if u.User != nil {
		userCopy := *u.User
		uCopy.User = &userCopy
	}

	return &uCopy
}

func copyRequest(req *http.Request) *http.Request {
	reqCopy := new(http.Request)
	*reqCopy = *req
	reqCopy.URL = copyURL(req.URL)
	reqCopy.Header = make(http.Header, len(req.Header))
	for k, s := range req.Header {
		reqCopy.Header[k] = append([]string(nil), s...)
	}
	return reqCopy
}

type onEOFReader struct {
	rc io.ReadCloser
	fn func()
}

func (r *onEOFReader) Read(p []byte) (n int, err error) {
	n, err = r.rc.Read(p)
	if err == io.EOF {
		r.runFunc()
	}
	return
}

func (r *onEOFReader) Close() error {
	err := r.rc.Close()
	r.runFunc()
	return err
}

func (r *onEOFReader) runFunc() {
	if fn := r.fn; fn != nil {
		fn()
		r.fn = nil
	}
}
<commit_msg>api/transport: add workaround for 500 from remote.api<commit_after>package api

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"sync"

	"github.com/koding/logging"
)

const maxBodyLen = 2 << 20 // 2 MiB

var defaultLog = logging.NewCustom("api-transport", false)

// contextKey is a value for use with context.WithValue.
type contextKey struct {
	name string
}

// UserContextKey is a context key. It can be used in
// HTTP handlers to attach a client session to a request.
//
// The session is then read by the Transport and used
// for authentication.
//
// The associated value will be of *User type.
var UserContextKey = &contextKey{"client-session"}

// UserFunc is used to extract a user, which requires authentication,
// from a request.
type UserFunc func(*http.Request) *User

// User represents a user, which is going to be authenticated.
type User struct {
	Username string // name of the user
	Team     string // team which the user belongs to
}

// WithRequest attaches a user to the given request.
func (u *User) WithRequest(req *http.Request) *http.Request {
	return req.WithContext(context.WithValue(req.Context(), UserContextKey, u))
}

// Strings implements the fmt.Stringer interface.
func (u *User) String() string {
	return u.Team + "/" + u.Username
}

// Session represents a user session.
type Session struct {
	ClientID string // an ID of an authenticated session.
	User     *User  // authenticated user
}

// Valid implements the stack.Validator interface.
func (s *Session) Valid() error {
	if s == nil {
		return errors.New("session is missing")
	}
	if s.ClientID == "" {
		return errors.New("empty client ID")
	}
	return nil
}

// WriteTo writes the s session to the given requeest.
func (s *Session) WriteTo(req *http.Request) {
	req.Header.Set("Authorization", "Bearer "+s.ClientID)
}

// ReadFrom initializes the s session by reading it from
// the given request.
//
// If no session can be read from req, the method is a nop.
func (s *Session) ReadFrom(req *http.Request) {
	if auth := req.Header.Get("Authorization"); strings.HasPrefix(auth, "Bearer ") {
		s.ClientID = auth[len("Bearer "):]
	}
}

// Match tests whether the other session matches the s one.
//
// Since ClientID is typically auto-generated, we match all the other fields.
func (s *Session) Match(other *Session) error {
	if s.User.Username != other.User.Username {
		return fmt.Errorf("username is  %q, the other one is %q", s.User.Username, other.User.Username)
	}

	if s.User.Team != other.User.Team {
		return fmt.Errorf("team is  %q, the other one is %q", s.User.Team, other.User.Team)
	}

	return nil
}

// AuthFunc is used to fetch session information.
//
// When Refresh is true, AuthFunc must obtain
// a session directly from an authorisation
// endpoint, in case it was cached.
type AuthFunc func(*AuthOptions) (*Session, error)

// AuthOptions represents the arguments for AuthFunc.
type AuthOptions struct {
	// User represents a user, which we are building
	// authorisation for.
	User *User

	// Refresh requests to refresh a session when true.
	Refresh bool

	// Request is a copy of the request, which
	// we're building authorisation for.
	Request *http.Request

	// Error is the error that was the reason
	// of temporary failure received while trying
	// to sending request.
	//
	// Error is non-nil when request failed with
	// temporary error and Transport is retrying.
	Error error
}

type httpTransport interface {
	http.RoundTripper
	httpRequestCanceler
	httpIdleConnectionsCloser
}

type httpRequestCanceler interface {
	CancelRequest(*http.Request)
}

type httpIdleConnectionsCloser interface {
	CloseIdleConnections()
}

// nopCloser is like ioutil.NopCloser, but wraps io.ReadSeeker.
type nopCloser struct{ io.ReadSeeker }

func (nopCloser) Close() (_ error) { return }

var _ io.ReadSeeker = nopCloser{bytes.NewReader(nil)}

// Transport is a signing transport that
// authorizes each request with jSession.clientId.
type Transport struct {
	http.RoundTripper                // a transport to use; optional, by default http.DefaultClient.Transport
	AuthFunc          AuthFunc       // clientID authorisation to use; required
	UserFunc          UserFunc       // used to get a user from request to authenticate for; by default looked up with UserContextKey
	Host              string         // original Host name, overwrites req.Host; optional
	RetryNum          int            // number of retries in case of temporary failures; optional, by default 3
	Log               logging.Logger // logger to use
	Debug             bool           // whether to dump HTTP request/response

	// request mapping implementation stolen from:
	//
	//   https://github.com/golang/oauth2/blob/master/transport.go
	//
	mu     sync.Mutex                      // guards modReq
	modReq map[*http.Request]*http.Request // original -> modified
}

var _ httpTransport = (*Transport)(nil)

// NewSingleUser returns a HTTP transport that authenticates all
// requests for the given user.
//
// This is usefull for api clients that give no control over HTTP requests,
// like the swagger-generated ones (remote.api).
func (t *Transport) NewSingleUser(u *User) http.RoundTripper {
	return &Transport{
		RoundTripper: t.RoundTripper,
		AuthFunc:     t.AuthFunc,
		UserFunc:     func(*http.Request) *User { return u },
		Host:         t.Host,
		RetryNum:     t.RetryNum,
		Log:          t.Log,
		Debug:        t.Debug,
	}
}

// RoundTrip implements the http.RoundTripper interface.
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	user := t.user(req)
	if user == nil {
		return t.roundTripper().RoundTrip(req)
	}

	reqCopy := copyRequest(req) // per RoundTripper contract
	t.setModReq(req, reqCopy)

	switch reqCopy.Body.(type) {
	case nil:
	case io.ReadSeeker:
	default:
		p, err := ioutil.ReadAll(io.LimitReader(reqCopy.Body, maxBodyLen))
		if err != nil {
			return nil, err // non-retryable
		}

		reqCopy.Body = nopCloser{bytes.NewReader(p)}
	}

	var refresh bool
	var lastErr error
	var session *Session

	for attempts := t.retryNum(); attempts >= 0; attempts-- {
		if reqCopy.Body != nil {
			if _, err := reqCopy.Body.(io.ReadSeeker).Seek(0, io.SeekStart); err != nil {
				return nil, err // non-retryable
			}
		}

		opts := &AuthOptions{
			User:    user,
			Refresh: refresh,
			Request: reqCopy,
			Error:   lastErr,
		}

		s, err := t.AuthFunc(opts)

		t.log().Debug("AuthFunc(%+v): session=%+v, error=%v", opts, s, err)

		if err != nil {
			return nil, err // non-retryable
		}

		if err := s.Valid(); err != nil {
			return nil, err // non-retryable
		}

		session = s
		refresh = false

		session.WriteTo(reqCopy)

		if t.Host != "" {
			reqCopy.Host = t.Host
		}

		if t.Debug {
			p, err := httputil.DumpRequestOut(reqCopy, true)
			t.log().Debug("[request] %s (err=%v)", p, err)
		}

		resp, err := t.roundTripper().RoundTrip(reqCopy)

		if t.Debug {
			p, err := httputil.DumpResponse(resp, true)
			t.log().Debug("[response] %s (err=%v)", p, err)
		}

		if e, ok := err.(net.Error); ok && (e.Temporary() || e.Timeout()) {
			lastErr = err
			continue // retry
		}

		if err != nil {
			t.setModReq(req, nil)
			return nil, err // non-retryable
		}

		switch resp.StatusCode {
		case http.StatusInternalServerError:
			// TODO(rjeczalik): remove, this is a temporary workaround
			// for remote.api, that return 500 when session could not
			// be validated.
			fallthrough
		case http.StatusUnauthorized:
			resp.Body.Close()
			lastErr = errors.New(http.StatusText(resp.StatusCode))
			refresh = true
			continue // retry and force new session
		}

		resp.Body = &onEOFReader{
			rc: resp.Body,
			fn: func() { t.setModReq(req, nil) },
		}

		return resp, nil
	}

	t.setModReq(req, nil)

	return nil, fmt.Errorf("api: error sending request: %v", lastErr)
}

func (t *Transport) CancelRequest(req *http.Request) {
	if rc, ok := t.roundTripper().(httpRequestCanceler); ok {
		t.mu.Lock()
		modReq := t.modReq[req]
		delete(t.modReq, req)
		t.mu.Unlock()

		rc.CancelRequest(modReq)
	}
}

func (t *Transport) CloseIdleConnections() {
	if icl, ok := t.roundTripper().(httpIdleConnectionsCloser); ok {
		icl.CloseIdleConnections()
	}
}

func (t *Transport) user(req *http.Request) *User {
	if t.UserFunc != nil {
		return t.UserFunc(req)
	}
	user, _ := req.Context().Value(UserContextKey).(*User)
	return user
}

func (t *Transport) roundTripper() http.RoundTripper {
	if t.RoundTripper != nil {
		return t.RoundTripper
	}

	if http.DefaultClient.Transport != nil {
		return http.DefaultClient.Transport
	}

	return http.DefaultTransport
}

func (t *Transport) retryNum() int {
	if t.RetryNum > 0 {
		return t.RetryNum
	}
	return 3
}

func (t *Transport) log() logging.Logger {
	if t.Log != nil {
		return t.Log
	}
	return defaultLog
}

func (t *Transport) setModReq(orig, mod *http.Request) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.modReq == nil {
		t.modReq = make(map[*http.Request]*http.Request)
	}
	if mod == nil {
		delete(t.modReq, orig)
	} else {
		t.modReq[orig] = mod
	}
}

func copyURL(u *url.URL) *url.URL {
	uCopy := *u
	if u.User != nil {
		userCopy := *u.User
		uCopy.User = &userCopy
	}

	return &uCopy
}

func copyRequest(req *http.Request) *http.Request {
	reqCopy := new(http.Request)
	*reqCopy = *req
	reqCopy.URL = copyURL(req.URL)
	reqCopy.Header = make(http.Header, len(req.Header))
	for k, s := range req.Header {
		reqCopy.Header[k] = append([]string(nil), s...)
	}
	return reqCopy
}

type onEOFReader struct {
	rc io.ReadCloser
	fn func()
}

func (r *onEOFReader) Read(p []byte) (n int, err error) {
	n, err = r.rc.Read(p)
	if err == io.EOF {
		r.runFunc()
	}
	return
}

func (r *onEOFReader) Close() error {
	err := r.rc.Close()
	r.runFunc()
	return err
}

func (r *onEOFReader) runFunc() {
	if fn := r.fn; fn != nil {
		fn()
		r.fn = nil
	}
}
<|endoftext|>


<commit_before>package hyperbahn

import (
	"fmt"
	"math/rand"
	"time"

	"golang.org/x/net/context"
)

const (
	// maxRegistrationFailures is the number of consecutive registration failures after
	// which we give up and trigger an OnError event.
	maxRegistrationFailures = 5
	// registrationInterval is the (approximate) time interval between re-registrations.
	// The interval is fuzzed to fuzzAmount.
	registrationInterval = 10 * time.Second
	// registrationRetryInterval is the duration to wait on failed registrations before retrying.
	registrationRetryInterval = 1 * time.Second
	// fuzzInterval is used to fuzz the registration interval.
	fuzzInterval = 10 * time.Second
)

// timeSleep is a variable for stubbing in unit tests.
var timeSleep = time.Sleep

// ErrRegistrationFailed is triggered when registration is failed.
type ErrRegistrationFailed struct {
	// WillRetry is set to true if registration will be retried.
	WillRetry bool
	// Cause is the underlying error returned from the register call.
	Cause error
}

func (e ErrRegistrationFailed) Error() string {
	return fmt.Sprintf("registration failed, retry: %v, cause: %v", e.WillRetry, e.Cause)
}

// The following parameters define the request/response for the Hyperbahn 'ad' call.
type service struct {
	Name string `json:"serviceName"`
	Cost int    `json:"cost"`
}

type adRequest struct {
	Services []service `json:"services"`
}

type adResponse struct {
	ConnectionCount int `json:"connectionCount"`
}

func (c *Client) sendRegistration() error {
	ctx, cancel := context.WithTimeout(context.Background(), c.opts.Timeout)
	defer cancel()

	sc := c.tchan.GetSubChannel(hyperbahnServiceName)
	arg := &adRequest{
		Services: []service{{
			Name: c.tchan.PeerInfo().ServiceName,
			Cost: 0,
		}},
	}
	var resp adResponse
	c.opts.Handler.On(RegistrationAttempt)
	if err := makeJSONCall(ctx, sc, "ad", arg, &resp); err != nil {
		return err
	}

	return nil
}

func (c *Client) fuzzedRegistrationInterval() time.Duration {
	// fuzz is a random value between -fuzzInterval and fuzzInterval
	fuzz := time.Duration(rand.Intn(int(fuzzInterval)*2)) - fuzzInterval
	return registrationInterval + fuzz
}

// registrationLoops re-registers the service approximately every minute (with some fuzzing).
func (c *Client) registrationLoop() {
	sleepFor := c.fuzzedRegistrationInterval()
	consecutiveFailures := 0

	for {
		timeSleep(sleepFor)

		if err := c.sendRegistration(); err != nil {
			consecutiveFailures++
			if consecutiveFailures >= maxRegistrationFailures {
				c.opts.Handler.OnError(ErrRegistrationFailed{Cause: err, WillRetry: false})
				if c.opts.FailStrategy == FailStrategyFatal {
					c.tchan.Logger().Fatalf("Hyperbahn client registration failed: %v", err)
				}
				return
			}
			c.opts.Handler.OnError(ErrRegistrationFailed{Cause: err, WillRetry: true})
			sleepFor = registrationRetryInterval
		} else {
			c.opts.Handler.On(RegistrationRefreshed)
			sleepFor = c.fuzzedRegistrationInterval()
			consecutiveFailures = 0
		}
	}
}
<commit_msg>switching from int to int64<commit_after>package hyperbahn

import (
	"fmt"
	"math/rand"
	"time"

	"golang.org/x/net/context"
)

const (
	// maxRegistrationFailures is the number of consecutive registration failures after
	// which we give up and trigger an OnError event.
	maxRegistrationFailures = 5
	// registrationInterval is the (approximate) time interval between re-registrations.
	// The interval is fuzzed to fuzzAmount.
	registrationInterval = 10 * time.Second
	// registrationRetryInterval is the duration to wait on failed registrations before retrying.
	registrationRetryInterval = 1 * time.Second
	// fuzzInterval is used to fuzz the registration interval.
	fuzzInterval = 10 * time.Second
)

// timeSleep is a variable for stubbing in unit tests.
var timeSleep = time.Sleep

// ErrRegistrationFailed is triggered when registration is failed.
type ErrRegistrationFailed struct {
	// WillRetry is set to true if registration will be retried.
	WillRetry bool
	// Cause is the underlying error returned from the register call.
	Cause error
}

func (e ErrRegistrationFailed) Error() string {
	return fmt.Sprintf("registration failed, retry: %v, cause: %v", e.WillRetry, e.Cause)
}

// The following parameters define the request/response for the Hyperbahn 'ad' call.
type service struct {
	Name string `json:"serviceName"`
	Cost int    `json:"cost"`
}

type adRequest struct {
	Services []service `json:"services"`
}

type adResponse struct {
	ConnectionCount int `json:"connectionCount"`
}

func (c *Client) sendRegistration() error {
	ctx, cancel := context.WithTimeout(context.Background(), c.opts.Timeout)
	defer cancel()

	sc := c.tchan.GetSubChannel(hyperbahnServiceName)
	arg := &adRequest{
		Services: []service{{
			Name: c.tchan.PeerInfo().ServiceName,
			Cost: 0,
		}},
	}
	var resp adResponse
	c.opts.Handler.On(RegistrationAttempt)
	if err := makeJSONCall(ctx, sc, "ad", arg, &resp); err != nil {
		return err
	}

	return nil
}

func (c *Client) fuzzedRegistrationInterval() time.Duration {
	// fuzz is a random value between -fuzzInterval and fuzzInterval
	fuzz := time.Duration(rand.Int63n(int64(fuzzInterval)*2)) - fuzzInterval
	return registrationInterval + fuzz
}

// registrationLoops re-registers the service approximately every minute (with some fuzzing).
func (c *Client) registrationLoop() {
	sleepFor := c.fuzzedRegistrationInterval()
	consecutiveFailures := 0

	for {
		timeSleep(sleepFor)

		if err := c.sendRegistration(); err != nil {
			consecutiveFailures++
			if consecutiveFailures >= maxRegistrationFailures {
				c.opts.Handler.OnError(ErrRegistrationFailed{Cause: err, WillRetry: false})
				if c.opts.FailStrategy == FailStrategyFatal {
					c.tchan.Logger().Fatalf("Hyperbahn client registration failed: %v", err)
				}
				return
			}
			c.opts.Handler.OnError(ErrRegistrationFailed{Cause: err, WillRetry: true})
			sleepFor = registrationRetryInterval
		} else {
			c.opts.Handler.On(RegistrationRefreshed)
			sleepFor = c.fuzzedRegistrationInterval()
			consecutiveFailures = 0
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"koding/broker/storage"
	"koding/tools/sockjs"
	"strconv"
	"strings"
	"time"

	"github.com/fatih/set"
	"github.com/streadway/amqp"
)

type Client struct {
	Session        *sockjs.Session
	ControlChannel *amqp.Channel
	SocketId       string
	Broker         *Broker
	LastPayload    string
	Subscriptions  storage.Subscriptionable
}

// NewClient retuns a new client that is defined on a given session.
func NewClient(session *sockjs.Session, broker *Broker) (*Client, error) {
	socketId := randomString()
	session.Tag = socketId

	controlChannel, err := broker.PublishConn.Channel()
	if err != nil {
		return nil, fmt.Errorf("Couldnt create publish channel %v", err)
	}

	subscriptions, err := createSubscriptionStorage(socketId)
	if err != nil {
		return nil, err
	}

	globalMapMutex.Lock()
	sessionsMap[socketId] = session
	globalMapMutex.Unlock()

	return &Client{
		Session:        session,
		SocketId:       socketId,
		ControlChannel: controlChannel,
		Broker:         broker,
		Subscriptions:  subscriptions,
	}, nil
}

func createSubscriptionStorage(socketId string) (s storage.Subscriptionable, errc error) {
	defer func() {
		log.Critical("Couldnt access to redis/create a key for client %v: Error: %v", socketId, errc)
	}()

	redisSubscriptions, err := storage.NewStorage(conf, storage.REDIS, socketId)
	if err == nil {
		return redisSubscriptions, nil
	}

	setSubscriptions, err := storage.NewStorage(conf, storage.SET, socketId)
	if err == nil {
		return setSubscriptions, nil
	}

	// this will never fail to here
	return nil, fmt.Errorf("Couldnt create subscription storage %v", err)
}

// Close should be called whenever a client disconnects.
func (c *Client) Close() {
	log.Debug("Client Close Request for socketID: %v", c.SocketId)
	c.Subscriptions.Each(func(routingKeyPrefix interface{}) bool {
		c.RemoveFromRoute(routingKeyPrefix.(string))
		return true
	})

	c.Subscriptions.ClearWithTimeout(time.Minute * 5)

	for {
		err := c.ControlChannel.Publish(c.Broker.Config.AuthAllExchange, "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(c.SocketId)})
		if err == nil {
			break
		}
		if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != amqp.ChannelError {
			log.Critical("Error while publising -not rabbitmq error- %v", err)
		}
		c.resetControlChannel()
	}

	log.Debug("Closing control channel for socketID: %v", c.SocketId)
	c.ControlChannel.Close()

	globalMapMutex.Lock()
	defer globalMapMutex.Unlock()
	delete(sessionsMap, c.SocketId)
}

// handleSessionMessage handles the received message from the client. It
// passes a response back to the client or publish the received message to a
// rabbitmq exchange for further process.
func (c *Client) handleSessionMessage(data interface{}) {
	message := data.(map[string]interface{})
	log.Debug("Received message: %v", message)

	action := message["action"]
	switch action {
	case "subscribe":
		routingKeyPrefixes := strings.Split(message["routingKeyPrefix"].(string), " ")
		if err := c.Subscribe(routingKeyPrefixes...); err != nil {
			log.Error(err.Error())
		}

		sendToClient(c.Session, "broker.subscribed", message["routingKeyPrefix"])

	case "resubscribe":
		clientId := message["socketId"].(string)
		log.Debug("Resubscribe event for clientId: %v SocketId: %v", clientId, c.SocketId)
		found, err := c.Resubscribe(clientId)
		if err != nil {
			log.Error(err.Error())
		}
		log.Debug("Resubscribe found for socketID: %v, %v", clientId, found)
		sendToClient(c.Session, "broker.resubscribed", found)

	case "unsubscribe":
		routingKeyPrefix := message["routingKeyPrefix"].(string)
		log.Debug("Unsubscribe event for socketID: %v, and prefixes", c.SocketId, routingKeyPrefix)
		routingKeyPrefixes := strings.Split(routingKeyPrefix, " ")
		c.Unsubscribe(routingKeyPrefixes...)

	case "publish":
		exchange := message["exchange"].(string)
		routingKey := message["routingKey"].(string)
		payload := message["payload"].(string)

		log.Debug("Publish Event: Exchange: %v, RoutingKey %v, Payload %v",
			exchange,
			routingKey,
			payload,
		)

		if err := c.Publish(exchange, routingKey, payload); err != nil {
			log.Error(err.Error())
		}

	case "ping":
		sendToClient(c.Session, "broker.pong", nil)
		if c.Subscriptions.Backend() == storage.REDIS {
			// TOOD - may be we need to revisit this part later about duration and request count
			go c.Subscriptions.ClearWithTimeout(time.Minute * 59)
		}
	default:
		log.Warning("Invalid action. message: %v socketId: %v", message, c.SocketId)

	}
}

// Publish publish the given payload for to the given exchange and routingkey.
func (c *Client) Publish(exchange, routingKey, payload string) error {
	if !strings.HasPrefix(routingKey, "client.") {
		return fmt.Errorf("Invalid routing key: %v socketId: %v", routingKey, c.SocketId)
	}

	for {
		c.LastPayload = ""
		err := c.ControlChannel.Publish(exchange, routingKey, false, false, amqp.Publishing{CorrelationId: c.SocketId, Body: []byte(payload)})
		if err == nil {
			c.LastPayload = payload
			break
		}

		if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != amqp.ChannelError {
			log.Warning("payload: %v routing key: %v exchange: %v err: %v",
				payload, routingKey, exchange, err)
		}

		time.Sleep(time.Second / 4) // penalty for crashing the AMQP channel
		c.resetControlChannel()
	}

	return nil
}

// gaugeStart starts the gauge for a given session. It returns a new
// function which ends the gauge for the given session. Usually one invokes
// gaugeStart and calls the returned function in a defer statement.
func (c *Client) gaugeStart() (gaugeEnd func()) {
	log.Debug("Client connected: %v", c.Session.Tag)
	changeClientsGauge(1)
	changeNewClientsGauge(1)
	if c.Session.IsWebsocket {
		changeWebsocketClientsGauge(1)
	}

	return func() {
		log.Debug("Client disconnected: %v", c.Session.Tag)
		changeClientsGauge(-1)
		if c.Session.IsWebsocket {
			changeWebsocketClientsGauge(-1)
		}
	}
}

// resetControlChannel closes the current client's control channel and creates
// a new channel. It also listens to any server side error and publish back
// the error to the client.
func (c *Client) resetControlChannel() {
	defer log.RecoverAndLog()

	if c.ControlChannel != nil {
		c.ControlChannel.Close()
	}

	var err error
	c.ControlChannel, err = c.Broker.PublishConn.Channel()
	if err != nil {
		log.Critical("Couldnt create publishing channel %v", err)
	}

	go func() {
		defer log.RecoverAndLog()

		for amqpErr := range c.ControlChannel.NotifyClose(make(chan *amqp.Error)) {
			if !(strings.Contains(amqpErr.Error(), "NOT_FOUND") && (strings.Contains(amqpErr.Error(), "koding-social-") || strings.Contains(amqpErr.Error(), "auth-"))) {
				log.Warning("AMQP channel: %v Last publish payload: %v", amqpErr.Error(), c.LastPayload)
			}

			sendToClient(c.Session, "broker.error", map[string]interface{}{
				"code":    amqpErr.Code,
				"reason":  amqpErr.Reason,
				"server":  amqpErr.Server,
				"recover": amqpErr.Recover,
			})
		}
	}()
}

// RemoveFromRoute removes the sessions for the given routingKeyPrefix.
func (c *Client) RemoveFromRoute(routingKeyPrefixes ...string) {
	for _, routingKeyPrefix := range routingKeyPrefixes {
		if _, ok := routeMap[routingKeyPrefix]; !ok {
			continue
		}
		routeMap[routingKeyPrefix].Remove(c.SocketId)

		if routeMap[routingKeyPrefix].Size() == 0 {
			delete(routeMap, routingKeyPrefix)
		}
	}
}

// Add to route
// todo ~ check for multiple subscriptions
func (c *Client) AddToRoute() {
	c.Subscriptions.Each(func(routingKeyPrefix interface{}) bool {
		c.AddToRouteMapNOTS(routingKeyPrefix.(string))
		return true
	})
}

func (c *Client) AddToRouteMapNOTS(routingKeyPrefixes ...string) {
	for _, routingKeyPrefix := range routingKeyPrefixes {
		if _, ok := routeMap[routingKeyPrefix]; !ok {
			routeMap[routingKeyPrefix] = set.New()
		}
		routeMap[routingKeyPrefix].Add(c.SocketId)
	}
}

// Subscribe add the given routingKeyPrefix to the list of subscriptions
// associated with this client.
func (c *Client) Subscribe(routingKeyPrefixes ...string) error {
	if err := c.Subscriptions.Subscribe(routingKeyPrefixes...); err != nil {
		return err
	}

	globalMapMutex.Lock()
	c.AddToRouteMapNOTS(routingKeyPrefixes...)
	globalMapMutex.Unlock()

	// Log some information about the Client
	go func() {
		length, err := c.Subscriptions.Len()
		if err != nil {
			log.Warning("Error while trying to get Subscriptions.Len() for: %v Error: %v", c.Session.Tag, err)
		}

		if length > 0 && length%2000 == 0 {
			log.Warning("Client with more than %v subscriptions %v", strconv.Itoa(length), c.Session.Tag)
		}
	}()

	return nil
}

func (c *Client) Resubscribe(sessionId string) (bool, error) {
	found, err := c.Subscriptions.Resubscribe(sessionId)
	if err != nil {
		return false, err
	}

	if !found {
		return false, nil
	}

	c.AddToRoute()
	return true, nil
}

// Unsubscribe deletes the given routingKey prefix from the subscription list
// and removes it from the global route map
func (c *Client) Unsubscribe(routingKeyPrefixes ...string) {
	c.RemoveFromRoute(routingKeyPrefixes...)
	if err := c.Subscriptions.Unsubscribe(routingKeyPrefixes...); err != nil {
		fmt.Errorf("%v", err)
	}
}

// randomString() returns a new 16 char length random string
func randomString() string {
	r := make([]byte, 128/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}
<commit_msg>Broker: fix error problem for broker<commit_after>package main

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"koding/broker/storage"
	"koding/tools/sockjs"
	"strconv"
	"strings"
	"time"

	"github.com/fatih/set"
	"github.com/streadway/amqp"
)

type Client struct {
	Session        *sockjs.Session
	ControlChannel *amqp.Channel
	SocketId       string
	Broker         *Broker
	LastPayload    string
	Subscriptions  storage.Subscriptionable
}

// NewClient retuns a new client that is defined on a given session.
func NewClient(session *sockjs.Session, broker *Broker) (*Client, error) {
	socketId := randomString()
	session.Tag = socketId

	controlChannel, err := broker.PublishConn.Channel()
	if err != nil {
		return nil, fmt.Errorf("Couldnt create publish channel %v", err)
	}

	subscriptions, err := createSubscriptionStorage(socketId)
	if err != nil {
		return nil, err
	}

	globalMapMutex.Lock()
	sessionsMap[socketId] = session
	globalMapMutex.Unlock()

	return &Client{
		Session:        session,
		SocketId:       socketId,
		ControlChannel: controlChannel,
		Broker:         broker,
		Subscriptions:  subscriptions,
	}, nil
}

func createSubscriptionStorage(socketId string) (storage.Subscriptionable, error) {

	if subscriptions, err := storage.NewStorage(conf, storage.REDIS, socketId); err == nil {
		return subscriptions, nil
	} else {
		log.Critical("Couldnt access to redis/create a key for client %v: Error: %v", socketId, err)
	}

	if subscriptions, err := storage.NewStorage(conf, storage.SET, socketId); err != nil {
		return subscriptions, nil
	}

	// this will never fail to here
	return nil, fmt.Errorf("Couldnt create subscription storage for Client: %v", socketId)
}

// Close should be called whenever a client disconnects.
func (c *Client) Close() {
	log.Debug("Client Close Request for socketID: %v", c.SocketId)
	c.Subscriptions.Each(func(routingKeyPrefix interface{}) bool {
		c.RemoveFromRoute(routingKeyPrefix.(string))
		return true
	})

	c.Subscriptions.ClearWithTimeout(time.Minute * 5)

	for {
		err := c.ControlChannel.Publish(c.Broker.Config.AuthAllExchange, "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(c.SocketId)})
		if err == nil {
			break
		}
		if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != amqp.ChannelError {
			log.Critical("Error while publising -not rabbitmq error- %v", err)
		}
		c.resetControlChannel()
	}

	log.Debug("Closing control channel for socketID: %v", c.SocketId)
	c.ControlChannel.Close()

	globalMapMutex.Lock()
	defer globalMapMutex.Unlock()
	delete(sessionsMap, c.SocketId)
}

// handleSessionMessage handles the received message from the client. It
// passes a response back to the client or publish the received message to a
// rabbitmq exchange for further process.
func (c *Client) handleSessionMessage(data interface{}) {
	message := data.(map[string]interface{})
	log.Debug("Received message: %v", message)

	action := message["action"]
	switch action {
	case "subscribe":
		routingKeyPrefixes := strings.Split(message["routingKeyPrefix"].(string), " ")
		if err := c.Subscribe(routingKeyPrefixes...); err != nil {
			log.Error(err.Error())
		}

		sendToClient(c.Session, "broker.subscribed", message["routingKeyPrefix"])

	case "resubscribe":
		clientId := message["socketId"].(string)
		log.Debug("Resubscribe event for clientId: %v SocketId: %v", clientId, c.SocketId)
		found, err := c.Resubscribe(clientId)
		if err != nil {
			log.Error(err.Error())
		}
		log.Debug("Resubscribe found for socketID: %v, %v", clientId, found)
		sendToClient(c.Session, "broker.resubscribed", found)

	case "unsubscribe":
		routingKeyPrefix := message["routingKeyPrefix"].(string)
		log.Debug("Unsubscribe event for socketID: %v, and prefixes", c.SocketId, routingKeyPrefix)
		routingKeyPrefixes := strings.Split(routingKeyPrefix, " ")
		c.Unsubscribe(routingKeyPrefixes...)

	case "publish":
		exchange := message["exchange"].(string)
		routingKey := message["routingKey"].(string)
		payload := message["payload"].(string)

		log.Debug("Publish Event: Exchange: %v, RoutingKey %v, Payload %v",
			exchange,
			routingKey,
			payload,
		)

		if err := c.Publish(exchange, routingKey, payload); err != nil {
			log.Error(err.Error())
		}

	case "ping":
		sendToClient(c.Session, "broker.pong", nil)
		if c.Subscriptions.Backend() == storage.REDIS {
			// TOOD - may be we need to revisit this part later about duration and request count
			go c.Subscriptions.ClearWithTimeout(time.Minute * 59)
		}
	default:
		log.Warning("Invalid action. message: %v socketId: %v", message, c.SocketId)

	}
}

// Publish publish the given payload for to the given exchange and routingkey.
func (c *Client) Publish(exchange, routingKey, payload string) error {
	if !strings.HasPrefix(routingKey, "client.") {
		return fmt.Errorf("Invalid routing key: %v socketId: %v", routingKey, c.SocketId)
	}

	for {
		c.LastPayload = ""
		err := c.ControlChannel.Publish(exchange, routingKey, false, false, amqp.Publishing{CorrelationId: c.SocketId, Body: []byte(payload)})
		if err == nil {
			c.LastPayload = payload
			break
		}

		if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != amqp.ChannelError {
			log.Warning("payload: %v routing key: %v exchange: %v err: %v",
				payload, routingKey, exchange, err)
		}

		time.Sleep(time.Second / 4) // penalty for crashing the AMQP channel
		c.resetControlChannel()
	}

	return nil
}

// gaugeStart starts the gauge for a given session. It returns a new
// function which ends the gauge for the given session. Usually one invokes
// gaugeStart and calls the returned function in a defer statement.
func (c *Client) gaugeStart() (gaugeEnd func()) {
	log.Debug("Client connected: %v", c.Session.Tag)
	changeClientsGauge(1)
	changeNewClientsGauge(1)
	if c.Session.IsWebsocket {
		changeWebsocketClientsGauge(1)
	}

	return func() {
		log.Debug("Client disconnected: %v", c.Session.Tag)
		changeClientsGauge(-1)
		if c.Session.IsWebsocket {
			changeWebsocketClientsGauge(-1)
		}
	}
}

// resetControlChannel closes the current client's control channel and creates
// a new channel. It also listens to any server side error and publish back
// the error to the client.
func (c *Client) resetControlChannel() {
	defer log.RecoverAndLog()

	if c.ControlChannel != nil {
		c.ControlChannel.Close()
	}

	var err error
	c.ControlChannel, err = c.Broker.PublishConn.Channel()
	if err != nil {
		log.Critical("Couldnt create publishing channel %v", err)
	}

	go func() {
		defer log.RecoverAndLog()

		for amqpErr := range c.ControlChannel.NotifyClose(make(chan *amqp.Error)) {
			if !(strings.Contains(amqpErr.Error(), "NOT_FOUND") && (strings.Contains(amqpErr.Error(), "koding-social-") || strings.Contains(amqpErr.Error(), "auth-"))) {
				log.Warning("AMQP channel: %v Last publish payload: %v", amqpErr.Error(), c.LastPayload)
			}

			sendToClient(c.Session, "broker.error", map[string]interface{}{
				"code":    amqpErr.Code,
				"reason":  amqpErr.Reason,
				"server":  amqpErr.Server,
				"recover": amqpErr.Recover,
			})
		}
	}()
}

// RemoveFromRoute removes the sessions for the given routingKeyPrefix.
func (c *Client) RemoveFromRoute(routingKeyPrefixes ...string) {
	for _, routingKeyPrefix := range routingKeyPrefixes {
		if _, ok := routeMap[routingKeyPrefix]; !ok {
			continue
		}
		routeMap[routingKeyPrefix].Remove(c.SocketId)

		if routeMap[routingKeyPrefix].Size() == 0 {
			delete(routeMap, routingKeyPrefix)
		}
	}
}

// Add to route
// todo ~ check for multiple subscriptions
func (c *Client) AddToRoute() {
	c.Subscriptions.Each(func(routingKeyPrefix interface{}) bool {
		c.AddToRouteMapNOTS(routingKeyPrefix.(string))
		return true
	})
}

func (c *Client) AddToRouteMapNOTS(routingKeyPrefixes ...string) {
	for _, routingKeyPrefix := range routingKeyPrefixes {
		if _, ok := routeMap[routingKeyPrefix]; !ok {
			routeMap[routingKeyPrefix] = set.New()
		}
		routeMap[routingKeyPrefix].Add(c.SocketId)
	}
}

// Subscribe add the given routingKeyPrefix to the list of subscriptions
// associated with this client.
func (c *Client) Subscribe(routingKeyPrefixes ...string) error {
	if err := c.Subscriptions.Subscribe(routingKeyPrefixes...); err != nil {
		return err
	}

	globalMapMutex.Lock()
	c.AddToRouteMapNOTS(routingKeyPrefixes...)
	globalMapMutex.Unlock()

	// Log some information about the Client
	go func() {
		length, err := c.Subscriptions.Len()
		if err != nil {
			log.Warning("Error while trying to get Subscriptions.Len() for: %v Error: %v", c.Session.Tag, err)
		}

		if length > 0 && length%2000 == 0 {
			log.Warning("Client with more than %v subscriptions %v", strconv.Itoa(length), c.Session.Tag)
		}
	}()

	return nil
}

func (c *Client) Resubscribe(sessionId string) (bool, error) {
	found, err := c.Subscriptions.Resubscribe(sessionId)
	if err != nil {
		return false, err
	}

	if !found {
		return false, nil
	}

	c.AddToRoute()
	return true, nil
}

// Unsubscribe deletes the given routingKey prefix from the subscription list
// and removes it from the global route map
func (c *Client) Unsubscribe(routingKeyPrefixes ...string) {
	c.RemoveFromRoute(routingKeyPrefixes...)
	if err := c.Subscriptions.Unsubscribe(routingKeyPrefixes...); err != nil {
		fmt.Errorf("%v", err)
	}
}

// randomString() returns a new 16 char length random string
func randomString() string {
	r := make([]byte, 128/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}
<|endoftext|>


<commit_before>package systests

import (
	"fmt"
	client "github.com/keybase/client/go/client"
	libkb "github.com/keybase/client/go/libkb"
	chat1 "github.com/keybase/client/go/protocol/chat1"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/stretchr/testify/require"
	"testing"
)

func sendChat(t smuTeam, u *smuUser, msg string) {
	tctx := u.primaryDevice().popClone()
	runner := client.NewCmdChatSendRunner(tctx.G)
	runner.SetTeamChatForTest(t.name)
	runner.SetMessage(msg)
	err := runner.Run()
	if err != nil {
		u.ctx.t.Fatal(err)
	}
}

func divDebug(ctx *smuContext, fmt string, arg ...interface{}) {
	div := "------------"
	ctx.log.Debug(div+" "+fmt+" "+div, arg...)
}

func readChatsWithError(team smuTeam, u *smuUser) ([]chat1.MessageUnboxed, error) {
	tctx := u.primaryDevice().popClone()
	runner := client.NewCmdChatReadRunner(tctx.G)
	runner.SetTeamChatForTest(team.name)
	_, messages, err := runner.Fetch()
	return messages, err
}

func readChats(team smuTeam, u *smuUser, nMessages int) {
	messages, err := readChatsWithError(team, u)
	t := u.ctx.t
	if err != nil {
		u.ctx.t.Fatal(err)
	}
	require.Equal(t, nMessages, len(messages))
	for i, msg := range messages {
		require.Equal(t, msg.Valid().MessageBody.Text().Body, fmt.Sprintf("%d", len(messages)-i-1))
	}
	divDebug(u.ctx, "readChat success for %s", u.username)
}

func pollForMembershipUpdate(team smuTeam, ann *smuUser, bob *smuUser, cam *smuUser) {

	details := ann.pollForMembershipUpdate(team, keybase1.PerTeamKeyGeneration(2))
	for _, member := range details.Members.Admins {
		switch member.Username {
		case ann.username:
			require.True(ann.ctx.t, member.Active)
		default:
			ann.ctx.t.Fatalf("unknown admin: %s", member.Username)
		}
	}
	for _, member := range details.Members.Writers {
		switch member.Username {
		case bob.username:
			require.False(ann.ctx.t, member.Active)
		case cam.username:
			require.True(ann.ctx.t, member.Active)
		default:
			ann.ctx.t.Fatalf("unknown writer: %s (%+v)", member.Username, details)
		}
	}
	ann.ctx.log.Debug("team details checked out: %+v", details)
}

func TestTeamReset(t *testing.T) {
	ctx := newSMUContext(t)
	defer ctx.cleanup()

	ann := ctx.installKeybaseForUser("ann", 10)
	ann.signup()
	divDebug(ctx, "Signed up ann (%s)", ann.username)
	bob := ctx.installKeybaseForUser("bob", 10)
	bob.signup()
	divDebug(ctx, "Signed up bob (%s)", bob.username)
	cam := ctx.installKeybaseForUser("cam", 10)
	cam.signup()
	divDebug(ctx, "Signed up cam (%s)", cam.username)

	team := ann.createTeam([]*smuUser{bob, cam})
	divDebug(ctx, "team created (%s)", team.name)

	sendChat(team, ann, "0")
	divDebug(ctx, "Sent chat '2' (%s via %s)", team.name, ann.username)

	readChats(team, ann, 1)
	readChats(team, bob, 1)

	kickTeamRekeyd(bob.getPrimaryGlobalContext(), t)
	bob.reset()
	divDebug(ctx, "Reset bob (%s)", bob.username)

	pollForMembershipUpdate(team, ann, bob, cam)
	divDebug(ctx, "Polled for rekey")

	bob.loginAfterReset(10)
	divDebug(ctx, "Bob logged in after reset")

	_, err := bob.teamGet(team)
	require.Error(t, err)
	ae, ok := err.(libkb.AppStatusError)
	require.True(t, ok)
	require.Equal(t, ae.Code, int(keybase1.StatusCode_SCTeamReadError))
	divDebug(ctx, "Bob failed to read the team")

	// Make sure that ann can still send even though bob is ousted
	sendChat(team, ann, "1")
	divDebug(ctx, "Sent chat '1' (%s via %s)", team.name, ann.username)
	readChats(team, ann, 2)
	// Same goes for cam --- note that she never read before, so nothing
	// is cached for her.
	readChats(team, cam, 2)

	_, err = readChatsWithError(team, bob)
	require.Error(t, err)
	ae, ok = err.(libkb.AppStatusError)
	require.True(t, ok)
	require.Equal(t, ae.Code, int(keybase1.StatusCode_SCTeamReadError))
	divDebug(ctx, "Bob failed to read the chat")

	ann.addWriter(team, bob)
	divDebug(ctx, "Added bob back as a writer")
	_, err = bob.teamGet(team)
	require.NoError(t, err)
	divDebug(ctx, "Bob could read the team after added back")
	readChats(team, bob, 2)
	divDebug(ctx, "Bob reading chats after added back")
	sendChat(team, ann, "2")
	divDebug(ctx, "Ann sending chat '2'")
	readChats(team, bob, 3)
	divDebug(ctx, "Bob reading chat '2'")
}
<commit_msg>Test more reset scenarios<commit_after>package systests

import (
	"fmt"
	"testing"

	client "github.com/keybase/client/go/client"
	libkb "github.com/keybase/client/go/libkb"
	chat1 "github.com/keybase/client/go/protocol/chat1"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/stretchr/testify/require"
)

func sendChat(t smuTeam, u *smuUser, msg string) {
	tctx := u.primaryDevice().popClone()
	runner := client.NewCmdChatSendRunner(tctx.G)
	runner.SetTeamChatForTest(t.name)
	runner.SetMessage(msg)
	err := runner.Run()
	if err != nil {
		u.ctx.t.Fatal(err)
	}
}

func divDebug(ctx *smuContext, fmt string, arg ...interface{}) {
	div := "------------"
	ctx.log.Debug(div+" "+fmt+" "+div, arg...)
}

func readChatsWithError(team smuTeam, u *smuUser) ([]chat1.MessageUnboxed, error) {
	tctx := u.primaryDevice().popClone()
	runner := client.NewCmdChatReadRunner(tctx.G)
	runner.SetTeamChatForTest(team.name)
	_, messages, err := runner.Fetch()
	return messages, err
}

func readChats(team smuTeam, u *smuUser, nMessages int) {
	messages, err := readChatsWithError(team, u)
	t := u.ctx.t
	if err != nil {
		u.ctx.t.Fatal(err)
	}
	require.Equal(t, nMessages, len(messages))
	for i, msg := range messages {
		require.Equal(t, msg.Valid().MessageBody.Text().Body, fmt.Sprintf("%d", len(messages)-i-1))
	}
	divDebug(u.ctx, "readChat success for %s", u.username)
}

func pollForMembershipUpdate(team smuTeam, ann *smuUser, bob *smuUser, cam *smuUser) {

	details := ann.pollForMembershipUpdate(team, keybase1.PerTeamKeyGeneration(2))
	for _, member := range details.Members.Admins {
		switch member.Username {
		case ann.username:
			require.True(ann.ctx.t, member.Active)
		default:
			ann.ctx.t.Fatalf("unknown admin: %s", member.Username)
		}
	}
	for _, member := range details.Members.Writers {
		switch member.Username {
		case bob.username:
			require.False(ann.ctx.t, member.Active)
		case cam.username:
			require.True(ann.ctx.t, member.Active)
		default:
			ann.ctx.t.Fatalf("unknown writer: %s (%+v)", member.Username, details)
		}
	}
	ann.ctx.log.Debug("team details checked out: %+v", details)
}

func TestTeamReset(t *testing.T) {
	ctx := newSMUContext(t)
	defer ctx.cleanup()

	ann := ctx.installKeybaseForUser("ann", 10)
	ann.signup()
	divDebug(ctx, "Signed up ann (%s)", ann.username)
	bob := ctx.installKeybaseForUser("bob", 10)
	bob.signup()
	divDebug(ctx, "Signed up bob (%s)", bob.username)
	cam := ctx.installKeybaseForUser("cam", 10)
	cam.signup()
	divDebug(ctx, "Signed up cam (%s)", cam.username)

	team := ann.createTeam([]*smuUser{bob, cam})
	divDebug(ctx, "team created (%s)", team.name)

	sendChat(team, ann, "0")
	divDebug(ctx, "Sent chat '2' (%s via %s)", team.name, ann.username)

	readChats(team, ann, 1)
	readChats(team, bob, 1)

	kickTeamRekeyd(bob.getPrimaryGlobalContext(), t)
	bob.reset()
	divDebug(ctx, "Reset bob (%s)", bob.username)

	pollForMembershipUpdate(team, ann, bob, cam)
	divDebug(ctx, "Polled for rekey")

	bob.loginAfterReset(10)
	divDebug(ctx, "Bob logged in after reset")

	_, err := bob.teamGet(team)
	require.Error(t, err)
	ae, ok := err.(libkb.AppStatusError)
	require.True(t, ok)
	require.Equal(t, ae.Code, int(keybase1.StatusCode_SCTeamReadError))
	divDebug(ctx, "Bob failed to read the team")

	// Make sure that ann can still send even though bob is ousted
	sendChat(team, ann, "1")
	divDebug(ctx, "Sent chat '1' (%s via %s)", team.name, ann.username)
	readChats(team, ann, 2)
	// Same goes for cam --- note that she never read before, so nothing
	// is cached for her.
	readChats(team, cam, 2)

	_, err = readChatsWithError(team, bob)
	require.Error(t, err)
	ae, ok = err.(libkb.AppStatusError)
	require.True(t, ok)
	require.Equal(t, ae.Code, int(keybase1.StatusCode_SCTeamReadError))
	divDebug(ctx, "Bob failed to read the chat")

	ann.addWriter(team, bob)
	divDebug(ctx, "Added bob back as a writer")
	_, err = bob.teamGet(team)
	require.NoError(t, err)
	divDebug(ctx, "Bob could read the team after added back")
	readChats(team, bob, 2)
	divDebug(ctx, "Bob reading chats after added back")
	sendChat(team, ann, "2")
	divDebug(ctx, "Ann sending chat '2'")
	readChats(team, bob, 3)
	divDebug(ctx, "Bob reading chat '2'")
}

func TestTeamReset2(t *testing.T) {
	ctx := newSMUContext(t)
	defer ctx.cleanup()

	ann := ctx.installKeybaseForUser("ann", 10)
	ann.signup()
	divDebug(ctx, "Signed up ann (%s)", ann.username)
	bob := ctx.installKeybaseForUser("bob", 10)
	bob.signup()
	divDebug(ctx, "Signed up bob (%s)", bob.username)
	cam := ctx.installKeybaseForUser("cam", 10)
	cam.signup()
	divDebug(ctx, "Signed up cam (%s)", cam.username)

	team := ann.createTeam([]*smuUser{cam})
	divDebug(ctx, "team created (%s)", team.name)

	sendChat(team, ann, "0")
	divDebug(ctx, "Sent chat '2' (%s via %s)", team.name, ann.username)

	readChats(team, ann, 1)

	bob.reset()
	divDebug(ctx, "Reset bob (%s)", bob.username)

	bob.loginAfterReset(10)
	divDebug(ctx, "Bob logged in after reset")

	_, err := bob.teamGet(team)
	require.Error(t, err)
	ae, ok := err.(libkb.AppStatusError)
	require.True(t, ok)
	require.Equal(t, ae.Code, int(keybase1.StatusCode_SCTeamReadError))
	divDebug(ctx, "Bob failed to read the team")

	ann.addWriter(team, bob)
	divDebug(ctx, "Added bob as a writer")
	_, err = bob.teamGet(team)
	require.NoError(t, err)
	divDebug(ctx, "Bob could read the team after added")
	readChats(team, bob, 1)
	divDebug(ctx, "Bob reading chats after added")
	sendChat(team, ann, "1")
	divDebug(ctx, "Ann sending chat '2'")
	readChats(team, bob, 2)
	divDebug(ctx, "Bob reading chat '2'")
}

// bob resets and has no keys
func TestTeamResetNoKeys(t *testing.T) {
	ctx := newSMUContext(t)
	defer ctx.cleanup()

	ann := ctx.installKeybaseForUser("ann", 10)
	ann.signup()
	divDebug(ctx, "Signed up ann (%s)", ann.username)
	bob := ctx.installKeybaseForUser("bob", 10)
	bob.signup()
	divDebug(ctx, "Signed up bob (%s)", bob.username)
	cam := ctx.installKeybaseForUser("cam", 10)
	cam.signup()
	divDebug(ctx, "Signed up cam (%s)", cam.username)

	team := ann.createTeam([]*smuUser{cam})
	divDebug(ctx, "team created (%s)", team.name)

	sendChat(team, ann, "0")
	divDebug(ctx, "Sent chat '2' (%s via %s)", team.name, ann.username)

	readChats(team, ann, 1)

	bob.reset()
	divDebug(ctx, "Reset bob (%s)", bob.username)

	ann.addWriter(team, bob)
	divDebug(ctx, "Added bob as a writer")

	/*
		_, err = bob.teamGet(team)
		require.NoError(t, err)
		divDebug(ctx, "Bob could read the team after added")
		readChats(team, bob, 1)
		divDebug(ctx, "Bob reading chats after added")
		sendChat(team, ann, "1")
		divDebug(ctx, "Ann sending chat '2'")
		readChats(team, bob, 2)
		divDebug(ctx, "Bob reading chat '2'")
	*/
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess Authors.

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

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

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

package timer

import (
	"sync/atomic"
	"time"
)

// SuspendableTicker is similar to time.Ticker, but also offers Suspend() and Resume() functions.
// While the ticker is suspended, nothing comes from the time channel C
type SuspendableTicker struct {
	ticker *time.Ticker
	// C is user facing
	C chan time.Time

	suspended int64
}

// NewSuspendableTicker creates a new suspendable ticker, indicating whether the ticker should start
// suspendable or running
func NewSuspendableTicker(d time.Duration, initiallySuspended bool) *SuspendableTicker {
	s := &SuspendableTicker{
		ticker: time.NewTicker(d),
		C:      make(chan time.Time),
	}
	if initiallySuspended {
		s.suspended = 1
	}
	go s.loop()
	return s
}

// Suspend stops sending time events on the channel C
// time events sent during suspended time are lost
func (s *SuspendableTicker) Suspend() {
	atomic.StoreInt64(&s.suspended, 1)
}

// Resume re-enables time events on channel C
func (s *SuspendableTicker) Resume() {
	atomic.StoreInt64(&s.suspended, 0)
}

// Stop completely stops the timer, like time.Timer
func (s *SuspendableTicker) Stop() {
	s.ticker.Stop()
}

func (s *SuspendableTicker) loop() {
	for t := range s.ticker.C {
		if atomic.LoadInt64(&s.suspended) == 0 {
			// not suspended
			s.C <- t
		}
	}
}
<commit_msg>support 'TickNow()'<commit_after>/*
Copyright 2020 The Vitess Authors.

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

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

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

package timer

import (
	"sync/atomic"
	"time"
)

// SuspendableTicker is similar to time.Ticker, but also offers Suspend() and Resume() functions.
// While the ticker is suspended, nothing comes from the time channel C
type SuspendableTicker struct {
	ticker *time.Ticker
	// C is user facing
	C chan time.Time

	suspended int64
}

// NewSuspendableTicker creates a new suspendable ticker, indicating whether the ticker should start
// suspendable or running
func NewSuspendableTicker(d time.Duration, initiallySuspended bool) *SuspendableTicker {
	s := &SuspendableTicker{
		ticker: time.NewTicker(d),
		C:      make(chan time.Time),
	}
	if initiallySuspended {
		s.suspended = 1
	}
	go s.loop()
	return s
}

// Suspend stops sending time events on the channel C
// time events sent during suspended time are lost
func (s *SuspendableTicker) Suspend() {
	atomic.StoreInt64(&s.suspended, 1)
}

// Resume re-enables time events on channel C
func (s *SuspendableTicker) Resume() {
	atomic.StoreInt64(&s.suspended, 0)
}

// Stop completely stops the timer, like time.Timer
func (s *SuspendableTicker) Stop() {
	s.ticker.Stop()
}

// TickNow generates a tick at this point in time. It may block
// if nothing consumes the tick.
func (s *SuspendableTicker) TickNow() {
	if atomic.LoadInt64(&s.suspended) == 0 {
		// not suspended
		s.C <- time.Now()
	}
}

func (s *SuspendableTicker) loop() {
	for t := range s.ticker.C {
		if atomic.LoadInt64(&s.suspended) == 0 {
			// not suspended
			s.C <- t
		}
	}
}
<|endoftext|>