1095 lines
32 KiB
Go
1095 lines
32 KiB
Go
package client
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"testing"
|
|
"time"
|
|
|
|
"tunnel_pls/internal/registry"
|
|
"tunnel_pls/session/interaction"
|
|
"tunnel_pls/session/lifecycle"
|
|
"tunnel_pls/session/slug"
|
|
"tunnel_pls/types"
|
|
|
|
proto "git.fossy.my.id/bagas/tunnel-please-grpc/gen"
|
|
"github.com/stretchr/testify/mock"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/health/grpc_health_v1"
|
|
"google.golang.org/grpc/status"
|
|
)
|
|
|
|
func TestClient_ClientConn(t *testing.T) {
|
|
conn := &grpc.ClientConn{}
|
|
c := &client{conn: conn}
|
|
if c.ClientConn() != conn {
|
|
t.Errorf("ClientConn() did not return expected connection")
|
|
}
|
|
}
|
|
|
|
func TestClient_Close(t *testing.T) {
|
|
c := &client{}
|
|
if err := c.Close(); err != nil {
|
|
t.Errorf("Close() on nil connection returned error: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestAuthorizeConn(t *testing.T) {
|
|
mockUserSvc := &mockUserServiceClient{}
|
|
c := &client{authorizeConnectionService: mockUserSvc}
|
|
|
|
tests := []struct {
|
|
name string
|
|
token string
|
|
mockResp *proto.CheckResponse
|
|
mockErr error
|
|
wantAuth bool
|
|
wantUser string
|
|
wantErr bool
|
|
}{
|
|
{
|
|
name: "Success",
|
|
token: "valid",
|
|
mockResp: &proto.CheckResponse{Response: proto.AuthorizationResponse_MESSAGE_TYPE_AUTHORIZED, User: "mas-fuad"},
|
|
wantAuth: true,
|
|
wantUser: "mas-fuad",
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Unauthorized",
|
|
token: "invalid",
|
|
mockResp: &proto.CheckResponse{Response: proto.AuthorizationResponse_MESSAGE_TYPE_UNAUTHORIZED},
|
|
wantAuth: false,
|
|
wantUser: "UNAUTHORIZED",
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "Error",
|
|
token: "error",
|
|
mockErr: errors.New("grpc error"),
|
|
wantAuth: false,
|
|
wantUser: "UNAUTHORIZED",
|
|
wantErr: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
mockUserSvc.checkFunc = func(ctx context.Context, in *proto.CheckRequest, opts ...grpc.CallOption) (*proto.CheckResponse, error) {
|
|
if in.AuthToken != tt.token {
|
|
t.Errorf("expected token %s, got %s", tt.token, in.AuthToken)
|
|
}
|
|
return tt.mockResp, tt.mockErr
|
|
}
|
|
|
|
auth, user, err := c.AuthorizeConn(context.Background(), tt.token)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Errorf("AuthorizeConn() error = %v, wantErr %v", err, tt.wantErr)
|
|
}
|
|
if auth != tt.wantAuth {
|
|
t.Errorf("AuthorizeConn() auth = %v, wantAuth %v", auth, tt.wantAuth)
|
|
}
|
|
if user != tt.wantUser {
|
|
t.Errorf("AuthorizeConn() user = %s, wantUser %s", user, tt.wantUser)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestHandleSubscribeError(t *testing.T) {
|
|
c := &client{}
|
|
ctx := context.Background()
|
|
canceledCtx, cancel := context.WithCancel(ctx)
|
|
cancel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
ctx context.Context
|
|
err error
|
|
backoff time.Duration
|
|
wantErr bool
|
|
wantB time.Duration
|
|
}{
|
|
{
|
|
name: "ContextCanceled",
|
|
ctx: canceledCtx,
|
|
err: context.Canceled,
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "GrpcCanceled",
|
|
ctx: ctx,
|
|
err: status.Error(codes.Canceled, "canceled"),
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "CtxErrSet",
|
|
ctx: canceledCtx,
|
|
err: errors.New("other error"),
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "Unauthenticated",
|
|
ctx: ctx,
|
|
err: status.Error(codes.Unauthenticated, "unauth"),
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "ConnectionError",
|
|
ctx: ctx,
|
|
err: status.Error(codes.Unavailable, "unavailable"),
|
|
backoff: time.Second,
|
|
wantErr: false,
|
|
wantB: 2 * time.Second,
|
|
},
|
|
{
|
|
name: "NonConnectionError",
|
|
ctx: ctx,
|
|
err: status.Error(codes.Internal, "internal"),
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "WaitCanceled",
|
|
ctx: canceledCtx,
|
|
err: status.Error(codes.Unavailable, "unavailable"),
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
backoff := tt.backoff
|
|
err := c.handleSubscribeError(tt.ctx, tt.err, &backoff)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Errorf("handleSubscribeError() error = %v, wantErr %v", err, tt.wantErr)
|
|
}
|
|
if !tt.wantErr && backoff != tt.wantB {
|
|
t.Errorf("handleSubscribeError() backoff = %v, want %v", backoff, tt.wantB)
|
|
}
|
|
})
|
|
}
|
|
|
|
t.Run("WaitCanceledReal", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
backoff := 50 * time.Millisecond
|
|
go func() {
|
|
time.Sleep(10 * time.Millisecond)
|
|
cancel()
|
|
}()
|
|
err := c.handleSubscribeError(ctx, status.Error(codes.Unavailable, "unavailable"), &backoff)
|
|
if err == nil {
|
|
t.Errorf("expected error from wait")
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestHandleStreamError(t *testing.T) {
|
|
c := &client{}
|
|
ctx := context.Background()
|
|
canceledCtx, cancel := context.WithCancel(ctx)
|
|
cancel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
ctx context.Context
|
|
err error
|
|
backoff time.Duration
|
|
wantErr bool
|
|
wantB time.Duration
|
|
}{
|
|
{
|
|
name: "ContextCanceled",
|
|
ctx: canceledCtx,
|
|
err: context.Canceled,
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "GrpcCanceled",
|
|
ctx: ctx,
|
|
err: status.Error(codes.Canceled, "canceled"),
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "CtxErrSet",
|
|
ctx: canceledCtx,
|
|
err: errors.New("other error"),
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "ConnectionError",
|
|
ctx: ctx,
|
|
err: status.Error(codes.Unavailable, "unavailable"),
|
|
backoff: time.Second,
|
|
wantErr: false,
|
|
wantB: 2 * time.Second,
|
|
},
|
|
{
|
|
name: "NonConnectionError",
|
|
ctx: ctx,
|
|
err: status.Error(codes.Internal, "internal"),
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "WaitCanceled",
|
|
ctx: canceledCtx,
|
|
err: status.Error(codes.Unavailable, "unavailable"),
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
backoff := tt.backoff
|
|
err := c.handleStreamError(tt.ctx, tt.err, &backoff)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Errorf("handleStreamError() error = %v, wantErr %v", err, tt.wantErr)
|
|
}
|
|
if !tt.wantErr && backoff != tt.wantB {
|
|
t.Errorf("handleStreamError() backoff = %v, want %v", backoff, tt.wantB)
|
|
}
|
|
})
|
|
}
|
|
|
|
t.Run("WaitCanceledReal", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
backoff := 50 * time.Millisecond
|
|
go func() {
|
|
time.Sleep(10 * time.Millisecond)
|
|
cancel()
|
|
}()
|
|
err := c.handleStreamError(ctx, status.Error(codes.Unavailable, "unavailable"), &backoff)
|
|
if err == nil {
|
|
t.Errorf("expected error from wait")
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestHandleAuthError(t *testing.T) {
|
|
c := &client{}
|
|
ctx := context.Background()
|
|
|
|
tests := []struct {
|
|
name string
|
|
err error
|
|
backoff time.Duration
|
|
wantErr bool
|
|
wantB time.Duration
|
|
}{
|
|
{
|
|
name: "ConnectionError",
|
|
err: status.Error(codes.Unavailable, "unavailable"),
|
|
backoff: time.Second,
|
|
wantErr: false,
|
|
wantB: 2 * time.Second,
|
|
},
|
|
{
|
|
name: "NonConnectionError",
|
|
err: status.Error(codes.Internal, "internal"),
|
|
backoff: time.Second,
|
|
wantErr: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
backoff := tt.backoff
|
|
err := c.handleAuthError(ctx, tt.err, &backoff)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Errorf("handleAuthError() error = %v, wantErr %v", err, tt.wantErr)
|
|
}
|
|
if !tt.wantErr && backoff != tt.wantB {
|
|
t.Errorf("handleAuthError() backoff = %v, want %v", backoff, tt.wantB)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestHandleAuthError_WaitFail(t *testing.T) {
|
|
c := &client{}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
backoff := time.Second
|
|
err := c.handleAuthError(ctx, status.Error(codes.Unavailable, "unavailable"), &backoff)
|
|
if err == nil {
|
|
t.Errorf("expected error when wait fails")
|
|
}
|
|
}
|
|
|
|
func TestProcessEventStream(t *testing.T) {
|
|
mockStream := &mockSubscribeClient{}
|
|
c := &client{}
|
|
|
|
t.Run("UnknownEventType", func(t *testing.T) {
|
|
mockStream.recvFunc = func() (*proto.Events, error) {
|
|
return &proto.Events{Type: proto.EventType(999)}, nil
|
|
}
|
|
first := true
|
|
mockStream.recvFunc = func() (*proto.Events, error) {
|
|
if first {
|
|
first = false
|
|
return &proto.Events{Type: proto.EventType(999)}, nil
|
|
}
|
|
return nil, io.EOF
|
|
}
|
|
err := c.processEventStream(mockStream)
|
|
if !errors.Is(err, io.EOF) {
|
|
t.Errorf("expected EOF, got %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("DispatchSuccess", func(t *testing.T) {
|
|
events := []proto.EventType{
|
|
proto.EventType_SLUG_CHANGE,
|
|
proto.EventType_GET_SESSIONS,
|
|
proto.EventType_TERMINATE_SESSION,
|
|
}
|
|
|
|
for _, et := range events {
|
|
t.Run(et.String(), func(t *testing.T) {
|
|
first := true
|
|
mockStream.recvFunc = func() (*proto.Events, error) {
|
|
if first {
|
|
first = false
|
|
payload := &proto.Events{Type: et}
|
|
switch et {
|
|
case proto.EventType_SLUG_CHANGE:
|
|
payload.Payload = &proto.Events_SlugEvent{SlugEvent: &proto.SlugChangeEvent{}}
|
|
case proto.EventType_GET_SESSIONS:
|
|
payload.Payload = &proto.Events_GetSessionsEvent{GetSessionsEvent: &proto.GetSessionsEvent{}}
|
|
case proto.EventType_TERMINATE_SESSION:
|
|
payload.Payload = &proto.Events_TerminateSessionEvent{TerminateSessionEvent: &proto.TerminateSessionEvent{}}
|
|
}
|
|
return payload, nil
|
|
}
|
|
return nil, io.EOF
|
|
}
|
|
mockReg := &mockRegistry{}
|
|
mockReg.getAllSessionFromUserFunc = func(user string) []registry.Session { return nil }
|
|
mockReg.getFunc = func(key registry.Key) (registry.Session, error) { return nil, errors.New("fail") }
|
|
mockReg.getWithUserFunc = func(user string, key registry.Key) (registry.Session, error) { return nil, errors.New("fail") }
|
|
c.sessionRegistry = mockReg
|
|
c.config = &MockConfig{}
|
|
c.config.(*MockConfig).On("Domain").Return("test.com")
|
|
mockStream.sendFunc = func(n *proto.Node) error { return nil }
|
|
|
|
err := c.processEventStream(mockStream)
|
|
if !errors.Is(err, io.EOF) {
|
|
t.Errorf("expected EOF, got %v", err)
|
|
}
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("HandlerError", func(t *testing.T) {
|
|
first := true
|
|
mockStream.recvFunc = func() (*proto.Events, error) {
|
|
if first {
|
|
first = false
|
|
return &proto.Events{Type: proto.EventType_SLUG_CHANGE, Payload: &proto.Events_SlugEvent{SlugEvent: &proto.SlugChangeEvent{}}}, nil
|
|
}
|
|
return nil, io.EOF
|
|
}
|
|
mockReg := &mockRegistry{}
|
|
mockReg.getFunc = func(key registry.Key) (registry.Session, error) { return nil, errors.New("fail") }
|
|
c.sessionRegistry = mockReg
|
|
mockStream.sendFunc = func(n *proto.Node) error { return status.Error(codes.Unavailable, "send fail") }
|
|
|
|
err := c.processEventStream(mockStream)
|
|
if !errors.Is(err, status.Error(codes.Unavailable, "send fail")) {
|
|
t.Errorf("expected send fail error, got %v", err)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestSendNode(t *testing.T) {
|
|
c := &client{}
|
|
mockStream := &mockSubscribeClient{}
|
|
|
|
t.Run("Success", func(t *testing.T) {
|
|
mockStream.sendFunc = func(n *proto.Node) error { return nil }
|
|
err := c.sendNode(mockStream, &proto.Node{}, "context")
|
|
if err != nil {
|
|
t.Errorf("sendNode error = %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("ConnectionError", func(t *testing.T) {
|
|
mockStream.sendFunc = func(n *proto.Node) error { return status.Error(codes.Unavailable, "fail") }
|
|
err := c.sendNode(mockStream, &proto.Node{}, "context")
|
|
if err == nil {
|
|
t.Errorf("expected error")
|
|
}
|
|
})
|
|
|
|
t.Run("OtherError", func(t *testing.T) {
|
|
mockStream.sendFunc = func(n *proto.Node) error { return status.Error(codes.Internal, "fail") }
|
|
err := c.sendNode(mockStream, &proto.Node{}, "context")
|
|
if err != nil {
|
|
t.Errorf("expected nil error for non-connection error (logged only)")
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestHandleSlugChange(t *testing.T) {
|
|
mockReg := &mockRegistry{}
|
|
mockStream := &mockSubscribeClient{}
|
|
c := &client{sessionRegistry: mockReg}
|
|
|
|
evt := &proto.Events{
|
|
Payload: &proto.Events_SlugEvent{
|
|
SlugEvent: &proto.SlugChangeEvent{
|
|
User: "mas-fuad",
|
|
Old: "old-slug",
|
|
New: "new-slug",
|
|
},
|
|
},
|
|
}
|
|
|
|
t.Run("Success", func(t *testing.T) {
|
|
mockSess := &mockSession{}
|
|
mockInter := &mockInteraction{}
|
|
mockSess.interactionFunc = func() interaction.Interaction { return mockInter }
|
|
|
|
mockReg.getFunc = func(key registry.Key) (registry.Session, error) {
|
|
if key.Id != "old-slug" {
|
|
t.Errorf("expected old-slug, got %s", key.Id)
|
|
}
|
|
return mockSess, nil
|
|
}
|
|
mockReg.updateFunc = func(user string, oldKey, newKey registry.Key) error {
|
|
if user != "mas-fuad" || oldKey.Id != "old-slug" || newKey.Id != "new-slug" {
|
|
t.Errorf("unexpected update args")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
sent := false
|
|
mockStream.sendFunc = func(n *proto.Node) error {
|
|
sent = true
|
|
if n.Type != proto.EventType_SLUG_CHANGE_RESPONSE {
|
|
t.Errorf("expected slug change response")
|
|
}
|
|
resp := n.GetSlugEventResponse()
|
|
if !resp.Success {
|
|
t.Errorf("expected success")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
err := c.handleSlugChange(mockStream, evt)
|
|
if err != nil {
|
|
t.Errorf("handleSlugChange error = %v", err)
|
|
}
|
|
if !mockInter.redrawCalled {
|
|
t.Errorf("redraw was not called")
|
|
}
|
|
if !sent {
|
|
t.Errorf("response not sent")
|
|
}
|
|
})
|
|
|
|
t.Run("SessionNotFound", func(t *testing.T) {
|
|
mockReg.getFunc = func(key registry.Key) (registry.Session, error) {
|
|
return nil, errors.New("not found")
|
|
}
|
|
mockStream.sendFunc = func(n *proto.Node) error {
|
|
resp := n.GetSlugEventResponse()
|
|
if resp.Success || resp.Message != "not found" {
|
|
t.Errorf("unexpected failure response: %v", resp)
|
|
}
|
|
return nil
|
|
}
|
|
err := c.handleSlugChange(mockStream, evt)
|
|
if err != nil {
|
|
t.Errorf("handleSlugChange should return nil if error is handled via response, but it currently returns whatever sendNode returns")
|
|
}
|
|
})
|
|
|
|
t.Run("UpdateError", func(t *testing.T) {
|
|
mockSess := &mockSession{}
|
|
mockReg.getFunc = func(key registry.Key) (registry.Session, error) { return mockSess, nil }
|
|
mockReg.updateFunc = func(user string, oldKey, newKey registry.Key) error {
|
|
return errors.New("update fail")
|
|
}
|
|
mockStream.sendFunc = func(n *proto.Node) error {
|
|
resp := n.GetSlugEventResponse()
|
|
if resp.Success || resp.Message != "update fail" {
|
|
t.Errorf("unexpected failure response: %v", resp)
|
|
}
|
|
return nil
|
|
}
|
|
err := c.handleSlugChange(mockStream, evt)
|
|
if err != nil {
|
|
t.Errorf("handleSlugChange error = %v", err)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestHandleGetSessions(t *testing.T) {
|
|
mockReg := &mockRegistry{}
|
|
mockStream := &mockSubscribeClient{}
|
|
mockCfg := &MockConfig{}
|
|
mockCfg.On("Domain").Return("test.com")
|
|
c := &client{sessionRegistry: mockReg, config: mockCfg}
|
|
|
|
evt := &proto.Events{
|
|
Payload: &proto.Events_GetSessionsEvent{
|
|
GetSessionsEvent: &proto.GetSessionsEvent{
|
|
Identity: "mas-fuad",
|
|
},
|
|
},
|
|
}
|
|
|
|
t.Run("Success", func(t *testing.T) {
|
|
now := time.Now()
|
|
mockSess := &mockSession{}
|
|
mockSess.detailFunc = func() *types.Detail {
|
|
return &types.Detail{
|
|
ForwardingType: "http",
|
|
Slug: "myslug",
|
|
UserID: "mas-fuad",
|
|
Active: true,
|
|
StartedAt: now,
|
|
}
|
|
}
|
|
|
|
mockReg.getAllSessionFromUserFunc = func(user string) []registry.Session {
|
|
if user != "mas-fuad" {
|
|
t.Errorf("expected mas-fuad, got %s", user)
|
|
}
|
|
return []registry.Session{mockSess}
|
|
}
|
|
|
|
sent := false
|
|
mockStream.sendFunc = func(n *proto.Node) error {
|
|
sent = true
|
|
if n.Type != proto.EventType_GET_SESSIONS {
|
|
t.Errorf("expected get sessions response type")
|
|
}
|
|
resp := n.GetGetSessionsEvent()
|
|
if len(resp.Details) != 1 || resp.Details[0].Slug != "myslug" {
|
|
t.Errorf("unexpected details: %v", resp.Details)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
err := c.handleGetSessions(mockStream, evt)
|
|
if err != nil {
|
|
t.Errorf("handleGetSessions error = %v", err)
|
|
}
|
|
if !sent {
|
|
t.Errorf("response not sent")
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestHandleTerminateSession(t *testing.T) {
|
|
mockReg := &mockRegistry{}
|
|
mockStream := &mockSubscribeClient{}
|
|
c := &client{sessionRegistry: mockReg}
|
|
|
|
evt := &proto.Events{
|
|
Payload: &proto.Events_TerminateSessionEvent{
|
|
TerminateSessionEvent: &proto.TerminateSessionEvent{
|
|
User: "mas-fuad",
|
|
Slug: "myslug",
|
|
TunnelType: proto.TunnelType_HTTP,
|
|
},
|
|
},
|
|
}
|
|
|
|
t.Run("Success", func(t *testing.T) {
|
|
mockSess := &mockSession{}
|
|
mockLife := &mockLifecycle{}
|
|
mockSess.lifecycleFunc = func() lifecycle.Lifecycle { return mockLife }
|
|
|
|
closed := false
|
|
mockLife.closeFunc = func() error {
|
|
closed = true
|
|
return nil
|
|
}
|
|
|
|
mockReg.getWithUserFunc = func(user string, key registry.Key) (registry.Session, error) {
|
|
if user != "mas-fuad" || key.Id != "myslug" || key.Type != types.TunnelTypeHTTP {
|
|
t.Errorf("unexpected get args")
|
|
}
|
|
return mockSess, nil
|
|
}
|
|
|
|
sent := false
|
|
mockStream.sendFunc = func(n *proto.Node) error {
|
|
sent = true
|
|
resp := n.GetTerminateSessionEventResponse()
|
|
if !resp.Success {
|
|
t.Errorf("expected success")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
err := c.handleTerminateSession(mockStream, evt)
|
|
if err != nil {
|
|
t.Errorf("handleTerminateSession error = %v", err)
|
|
}
|
|
if !closed {
|
|
t.Errorf("close was not called")
|
|
}
|
|
if !sent {
|
|
t.Errorf("response not sent")
|
|
}
|
|
})
|
|
|
|
t.Run("TunnelTypeUnknown", func(t *testing.T) {
|
|
badEvt := &proto.Events{
|
|
Payload: &proto.Events_TerminateSessionEvent{
|
|
TerminateSessionEvent: &proto.TerminateSessionEvent{
|
|
TunnelType: proto.TunnelType(999),
|
|
},
|
|
},
|
|
}
|
|
mockStream.sendFunc = func(n *proto.Node) error {
|
|
resp := n.GetTerminateSessionEventResponse()
|
|
if resp.Success || resp.Message == "" {
|
|
t.Errorf("expected failure response")
|
|
}
|
|
return nil
|
|
}
|
|
err := c.handleTerminateSession(mockStream, badEvt)
|
|
if err != nil {
|
|
t.Errorf("handleTerminateSession error = %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("SessionNotFound", func(t *testing.T) {
|
|
mockReg.getWithUserFunc = func(user string, key registry.Key) (registry.Session, error) {
|
|
return nil, errors.New("not found")
|
|
}
|
|
mockStream.sendFunc = func(n *proto.Node) error {
|
|
resp := n.GetTerminateSessionEventResponse()
|
|
if resp.Success || resp.Message != "not found" {
|
|
t.Errorf("unexpected failure response: %v", resp)
|
|
}
|
|
return nil
|
|
}
|
|
err := c.handleTerminateSession(mockStream, evt)
|
|
if err != nil {
|
|
t.Errorf("handleTerminateSession error = %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("CloseError", func(t *testing.T) {
|
|
mockSess := &mockSession{}
|
|
mockLife := &mockLifecycle{}
|
|
mockSess.lifecycleFunc = func() lifecycle.Lifecycle { return mockLife }
|
|
mockLife.closeFunc = func() error { return errors.New("close fail") }
|
|
mockReg.getWithUserFunc = func(user string, key registry.Key) (registry.Session, error) { return mockSess, nil }
|
|
|
|
mockStream.sendFunc = func(n *proto.Node) error {
|
|
resp := n.GetTerminateSessionEventResponse()
|
|
if resp.Success || resp.Message != "close fail" {
|
|
t.Errorf("expected failure response: %v", resp)
|
|
}
|
|
return nil
|
|
}
|
|
err := c.handleTerminateSession(mockStream, evt)
|
|
if err != nil {
|
|
t.Errorf("handleTerminateSession error = %v", err)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestSubscribeAndProcess(t *testing.T) {
|
|
mockEventSvc := &mockEventServiceClient{}
|
|
c := &client{eventService: mockEventSvc}
|
|
ctx := context.Background()
|
|
backoff := time.Second
|
|
|
|
t.Run("SubscribeError", func(t *testing.T) {
|
|
mockEventSvc.subscribeFunc = func(ctx context.Context, opts ...grpc.CallOption) (proto.EventService_SubscribeClient, error) {
|
|
return nil, status.Error(codes.Unauthenticated, "unauth")
|
|
}
|
|
err := c.subscribeAndProcess(ctx, "id", "token", &backoff)
|
|
if !errors.Is(err, status.Error(codes.Unauthenticated, "unauth")) {
|
|
t.Errorf("expected unauth error, got %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("AuthSendError", func(t *testing.T) {
|
|
mockStream := &mockSubscribeClient{}
|
|
mockEventSvc.subscribeFunc = func(ctx context.Context, opts ...grpc.CallOption) (proto.EventService_SubscribeClient, error) {
|
|
return mockStream, nil
|
|
}
|
|
mockStream.sendFunc = func(n *proto.Node) error {
|
|
return status.Error(codes.Internal, "send fail")
|
|
}
|
|
err := c.subscribeAndProcess(ctx, "id", "token", &backoff)
|
|
if !errors.Is(err, status.Error(codes.Internal, "send fail")) {
|
|
t.Errorf("expected send fail, got %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("StreamError", func(t *testing.T) {
|
|
mockStream := &mockSubscribeClient{}
|
|
mockEventSvc.subscribeFunc = func(ctx context.Context, opts ...grpc.CallOption) (proto.EventService_SubscribeClient, error) {
|
|
return mockStream, nil
|
|
}
|
|
mockStream.sendFunc = func(n *proto.Node) error { return nil }
|
|
mockStream.recvFunc = func() (*proto.Events, error) {
|
|
return nil, status.Error(codes.Internal, "stream fail")
|
|
}
|
|
err := c.subscribeAndProcess(ctx, "id", "token", &backoff)
|
|
if !errors.Is(err, status.Error(codes.Internal, "stream fail")) {
|
|
t.Errorf("expected stream fail, got %v", err)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestSubscribeEvents(t *testing.T) {
|
|
mockEventSvc := &mockEventServiceClient{}
|
|
c := &client{eventService: mockEventSvc}
|
|
|
|
t.Run("ReturnsOnError", func(t *testing.T) {
|
|
mockEventSvc.subscribeFunc = func(ctx context.Context, opts ...grpc.CallOption) (proto.EventService_SubscribeClient, error) {
|
|
return nil, errors.New("fatal error")
|
|
}
|
|
err := c.SubscribeEvents(context.Background(), "id", "token")
|
|
if err == nil || err.Error() != "fatal error" {
|
|
t.Errorf("expected fatal error, got %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("RetryLoop", func(t *testing.T) {
|
|
oldB := initialBackoff
|
|
initialBackoff = 5 * time.Millisecond
|
|
defer func() { initialBackoff = oldB }()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
|
defer cancel()
|
|
|
|
callCount := 0
|
|
mockEventSvc.subscribeFunc = func(ctx context.Context, opts ...grpc.CallOption) (proto.EventService_SubscribeClient, error) {
|
|
callCount++
|
|
return nil, status.Error(codes.Unavailable, "unavailable")
|
|
}
|
|
|
|
err := c.SubscribeEvents(ctx, "id", "token")
|
|
if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) {
|
|
t.Errorf("expected timeout/canceled error, got %v", err)
|
|
}
|
|
if callCount <= 1 {
|
|
t.Errorf("expected multiple calls due to retry, got %d", callCount)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestCheckServerHealth(t *testing.T) {
|
|
mockHealth := &mockHealthClient{}
|
|
old := healthNewHealthClient
|
|
healthNewHealthClient = func(cc grpc.ClientConnInterface) grpc_health_v1.HealthClient {
|
|
return mockHealth
|
|
}
|
|
defer func() { healthNewHealthClient = old }()
|
|
|
|
c := &client{}
|
|
|
|
t.Run("Success", func(t *testing.T) {
|
|
mockHealth.checkFunc = func(ctx context.Context, in *grpc_health_v1.HealthCheckRequest, opts ...grpc.CallOption) (*grpc_health_v1.HealthCheckResponse, error) {
|
|
return &grpc_health_v1.HealthCheckResponse{Status: grpc_health_v1.HealthCheckResponse_SERVING}, nil
|
|
}
|
|
err := c.CheckServerHealth(context.Background())
|
|
if err != nil {
|
|
t.Errorf("expected nil error, got %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("Error", func(t *testing.T) {
|
|
mockHealth.checkFunc = func(ctx context.Context, in *grpc_health_v1.HealthCheckRequest, opts ...grpc.CallOption) (*grpc_health_v1.HealthCheckResponse, error) {
|
|
return nil, errors.New("health fail")
|
|
}
|
|
err := c.CheckServerHealth(context.Background())
|
|
if err == nil || err.Error() != "health check failed: health fail" {
|
|
t.Errorf("expected health fail error, got %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("NotServing", func(t *testing.T) {
|
|
mockHealth.checkFunc = func(ctx context.Context, in *grpc_health_v1.HealthCheckRequest, opts ...grpc.CallOption) (*grpc_health_v1.HealthCheckResponse, error) {
|
|
return &grpc_health_v1.HealthCheckResponse{Status: grpc_health_v1.HealthCheckResponse_NOT_SERVING}, nil
|
|
}
|
|
err := c.CheckServerHealth(context.Background())
|
|
if err == nil || err.Error() != "server not serving: NOT_SERVING" {
|
|
t.Errorf("expected not serving error, got %v", err)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestNew_Error(t *testing.T) {
|
|
old := grpcNewClient
|
|
grpcNewClient = func(target string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
|
|
return nil, errors.New("dial fail")
|
|
}
|
|
defer func() { grpcNewClient = old }()
|
|
mockConfig := &MockConfig{}
|
|
|
|
mockConfig.On("GRPCAddress").Return("localhost")
|
|
mockConfig.On("GRPCPort").Return("1234")
|
|
cli, err := New(mockConfig, &mockRegistry{})
|
|
if err == nil || err.Error() != "failed to connect to gRPC server at localhost:1234: dial fail" {
|
|
t.Errorf("expected dial fail error, got %v", err)
|
|
}
|
|
if cli != nil {
|
|
t.Errorf("expected nil client")
|
|
}
|
|
}
|
|
|
|
func TestNew(t *testing.T) {
|
|
mockConfig := &MockConfig{}
|
|
mockReg := &mockRegistry{}
|
|
mockConfig.On("GRPCAddress").Return("localhost")
|
|
mockConfig.On("GRPCPort").Return("1234")
|
|
cli, err := New(mockConfig, mockReg)
|
|
if err != nil {
|
|
t.Errorf("New() error = %v", err)
|
|
}
|
|
if cli == nil {
|
|
t.Fatal("New() returned nil client")
|
|
}
|
|
defer cli.Close()
|
|
}
|
|
|
|
type MockConfig struct {
|
|
mock.Mock
|
|
}
|
|
|
|
func (m *MockConfig) Domain() string { return m.Called().String(0) }
|
|
func (m *MockConfig) SSHPort() string { return m.Called().String(0) }
|
|
func (m *MockConfig) HTTPPort() string { return m.Called().String(0) }
|
|
func (m *MockConfig) HTTPSPort() string { return m.Called().String(0) }
|
|
func (m *MockConfig) TLSEnabled() bool { return m.Called().Bool(0) }
|
|
func (m *MockConfig) TLSRedirect() bool { return m.Called().Bool(0) }
|
|
func (m *MockConfig) TLSStoragePath() string { return m.Called().String(0) }
|
|
func (m *MockConfig) ACMEEmail() string { return m.Called().String(0) }
|
|
func (m *MockConfig) CFAPIToken() string { return m.Called().String(0) }
|
|
func (m *MockConfig) ACMEStaging() bool { return m.Called().Bool(0) }
|
|
func (m *MockConfig) AllowedPortsStart() uint16 { return uint16(m.Called().Int(0)) }
|
|
func (m *MockConfig) AllowedPortsEnd() uint16 { return uint16(m.Called().Int(0)) }
|
|
func (m *MockConfig) BufferSize() int { return m.Called().Int(0) }
|
|
func (m *MockConfig) PprofEnabled() bool { return m.Called().Bool(0) }
|
|
func (m *MockConfig) PprofPort() string { return m.Called().String(0) }
|
|
func (m *MockConfig) Mode() types.ServerMode { return m.Called().Get(0).(types.ServerMode) }
|
|
func (m *MockConfig) GRPCAddress() string { return m.Called().String(0) }
|
|
func (m *MockConfig) GRPCPort() string { return m.Called().String(0) }
|
|
func (m *MockConfig) NodeToken() string { return m.Called().String(0) }
|
|
func (m *MockConfig) KeyLoc() string { return m.Called().String(0) }
|
|
|
|
type mockRegistry struct {
|
|
registry.Registry
|
|
getFunc func(key registry.Key) (registry.Session, error)
|
|
getWithUserFunc func(user string, key registry.Key) (registry.Session, error)
|
|
updateFunc func(user string, oldKey, newKey registry.Key) error
|
|
getAllSessionFromUserFunc func(user string) []registry.Session
|
|
}
|
|
|
|
func (m *mockRegistry) Get(key registry.Key) (registry.Session, error) {
|
|
return m.getFunc(key)
|
|
}
|
|
func (m *mockRegistry) GetWithUser(user string, key registry.Key) (registry.Session, error) {
|
|
return m.getWithUserFunc(user, key)
|
|
}
|
|
func (m *mockRegistry) Update(user string, oldKey, newKey registry.Key) error {
|
|
return m.updateFunc(user, oldKey, newKey)
|
|
}
|
|
func (m *mockRegistry) GetAllSessionFromUser(user string) []registry.Session {
|
|
return m.getAllSessionFromUserFunc(user)
|
|
}
|
|
|
|
type mockSession struct {
|
|
registry.Session
|
|
lifecycleFunc func() lifecycle.Lifecycle
|
|
interactionFunc func() interaction.Interaction
|
|
detailFunc func() *types.Detail
|
|
slugFunc func() slug.Slug
|
|
}
|
|
|
|
func (m *mockSession) Lifecycle() lifecycle.Lifecycle { return m.lifecycleFunc() }
|
|
func (m *mockSession) Interaction() interaction.Interaction { return m.interactionFunc() }
|
|
func (m *mockSession) Detail() *types.Detail { return m.detailFunc() }
|
|
func (m *mockSession) Slug() slug.Slug { return m.slugFunc() }
|
|
|
|
type mockInteraction struct {
|
|
interaction.Interaction
|
|
redrawCalled bool
|
|
}
|
|
|
|
func (m *mockInteraction) Redraw() { m.redrawCalled = true }
|
|
|
|
type mockLifecycle struct {
|
|
lifecycle.Lifecycle
|
|
closeFunc func() error
|
|
}
|
|
|
|
func (m *mockLifecycle) Close() error { return m.closeFunc() }
|
|
|
|
type mockEventServiceClient struct {
|
|
proto.EventServiceClient
|
|
subscribeFunc func(ctx context.Context, opts ...grpc.CallOption) (proto.EventService_SubscribeClient, error)
|
|
}
|
|
|
|
func (m *mockEventServiceClient) Subscribe(ctx context.Context, opts ...grpc.CallOption) (proto.EventService_SubscribeClient, error) {
|
|
return m.subscribeFunc(ctx, opts...)
|
|
}
|
|
|
|
type mockSubscribeClient struct {
|
|
grpc.ClientStream
|
|
sendFunc func(*proto.Node) error
|
|
recvFunc func() (*proto.Events, error)
|
|
}
|
|
|
|
func (m *mockSubscribeClient) Send(n *proto.Node) error { return m.sendFunc(n) }
|
|
func (m *mockSubscribeClient) Recv() (*proto.Events, error) { return m.recvFunc() }
|
|
func (m *mockSubscribeClient) Context() context.Context { return context.Background() }
|
|
|
|
type mockUserServiceClient struct {
|
|
proto.UserServiceClient
|
|
checkFunc func(ctx context.Context, in *proto.CheckRequest, opts ...grpc.CallOption) (*proto.CheckResponse, error)
|
|
}
|
|
|
|
func (m *mockUserServiceClient) Check(ctx context.Context, in *proto.CheckRequest, opts ...grpc.CallOption) (*proto.CheckResponse, error) {
|
|
return m.checkFunc(ctx, in, opts...)
|
|
}
|
|
|
|
type mockHealthClient struct {
|
|
grpc_health_v1.HealthClient
|
|
checkFunc func(ctx context.Context, in *grpc_health_v1.HealthCheckRequest, opts ...grpc.CallOption) (*grpc_health_v1.HealthCheckResponse, error)
|
|
}
|
|
|
|
func (m *mockHealthClient) Check(ctx context.Context, in *grpc_health_v1.HealthCheckRequest, opts ...grpc.CallOption) (*grpc_health_v1.HealthCheckResponse, error) {
|
|
return m.checkFunc(ctx, in, opts...)
|
|
}
|
|
|
|
func TestProtoToTunnelType(t *testing.T) {
|
|
c := &client{}
|
|
tests := []struct {
|
|
name string
|
|
input proto.TunnelType
|
|
want types.TunnelType
|
|
wantErr bool
|
|
}{
|
|
{"HTTP", proto.TunnelType_HTTP, types.TunnelTypeHTTP, false},
|
|
{"TCP", proto.TunnelType_TCP, types.TunnelTypeTCP, false},
|
|
{"Unknown", proto.TunnelType(999), types.TunnelTypeUNKNOWN, true},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got, err := c.protoToTunnelType(tt.input)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Errorf("protoToTunnelType() error = %v, wantErr %v", err, tt.wantErr)
|
|
return
|
|
}
|
|
if got != tt.want {
|
|
t.Errorf("protoToTunnelType() got = %v, want %v", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIsConnectionError(t *testing.T) {
|
|
c := &client{}
|
|
tests := []struct {
|
|
name string
|
|
closing bool
|
|
err error
|
|
want bool
|
|
}{
|
|
{"NilError", false, nil, false},
|
|
{"Closing", true, io.EOF, false},
|
|
{"EOF", false, io.EOF, true},
|
|
{"Unavailable", false, status.Error(codes.Unavailable, "unavailable"), true},
|
|
{"Canceled", false, status.Error(codes.Canceled, "canceled"), true},
|
|
{"DeadlineExceeded", false, status.Error(codes.DeadlineExceeded, "deadline"), true},
|
|
{"Internal", false, status.Error(codes.Internal, "internal"), false},
|
|
{"WrappedEOF", false, errors.New("wrapped: " + io.EOF.Error()), false},
|
|
}
|
|
|
|
tests[7].err = fmt.Errorf("wrapped: %w", io.EOF)
|
|
tests[7].want = true
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
c.closing = tt.closing
|
|
if got := c.isConnectionError(tt.err); got != tt.want {
|
|
t.Errorf("isConnectionError() = %v, want %v for error %v", got, tt.want, tt.err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGrowBackoff(t *testing.T) {
|
|
c := &client{}
|
|
tests := []struct {
|
|
name string
|
|
initial time.Duration
|
|
want time.Duration
|
|
}{
|
|
{"NormalGrow", time.Second, 2 * time.Second},
|
|
{"MaxLimit", 20 * time.Second, 30 * time.Second},
|
|
{"AlreadyAtMax", 30 * time.Second, 30 * time.Second},
|
|
{"OverMax", 40 * time.Second, 30 * time.Second},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
backoff := tt.initial
|
|
c.growBackoff(&backoff)
|
|
if backoff != tt.want {
|
|
t.Errorf("growBackoff() = %v, want %v", backoff, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestWait(t *testing.T) {
|
|
c := &client{}
|
|
|
|
t.Run("ZeroDuration", func(t *testing.T) {
|
|
err := c.wait(context.Background(), 0)
|
|
if err != nil {
|
|
t.Errorf("wait() zero duration error = %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("ContextCanceled", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
err := c.wait(ctx, time.Minute)
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Errorf("wait() context canceled error = %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("Timeout", func(t *testing.T) {
|
|
start := time.Now()
|
|
err := c.wait(context.Background(), 10*time.Millisecond)
|
|
if err != nil {
|
|
t.Errorf("wait() timeout error = %v", err)
|
|
}
|
|
if time.Since(start) < 10*time.Millisecond {
|
|
t.Errorf("wait() returned too early")
|
|
}
|
|
})
|
|
}
|