1 package middleware
2
3 import (
4 "context"
5 "io"
6 "net/http"
7 "net/http/httptest"
8 "testing"
9
10 "edge-infra.dev/pkg/edge/api/bsl/types"
11 "edge-infra.dev/pkg/edge/api/totp"
12 "edge-infra.dev/pkg/edge/client"
13
14 "github.com/DATA-DOG/go-sqlmock"
15 "github.com/gin-gonic/gin"
16 "github.com/stretchr/testify/assert"
17 )
18
19 func TestAuthMiddlewareTotpToken(t *testing.T) {
20 assert := assert.New(t)
21
22 totpSecret := "totp-secret"
23 router := gin.Default()
24
25 db, _, err := sqlmock.New()
26 assert.NoError(err)
27
28 router.Use(AuthMiddleware("", totpSecret, db))
29
30 router.GET("/", func(c *gin.Context) { c.String(http.StatusOK, "OK") })
31
32 ts := httptest.NewServer(router)
33 defer ts.Close()
34 req, err := http.NewRequest("GET", ts.URL+"/", http.NoBody)
35 assert.NoError(err)
36
37 totpToken, err := totp.GenerateTotp(totpSecret)
38 assert.NoError(err)
39
40 req.Header.Set("Authorization", client.TotpToken+" "+totpToken.Code)
41
42 resp, err := http.DefaultClient.Do(req)
43 assert.NoError(err)
44
45 respBody, err := io.ReadAll(resp.Body)
46 assert.NoError(err)
47 assert.Equal(http.StatusOK, resp.StatusCode)
48 assert.Equal("OK", string(respBody))
49
50 totpToken, err = totp.GenerateTotp("invalid-secret")
51 assert.NoError(err)
52
53 req.Header.Set("Authorization", client.TotpToken+" "+totpToken.Code)
54
55 resp, err = http.DefaultClient.Do(req)
56 assert.NoError(err)
57 assert.Equal(http.StatusForbidden, resp.StatusCode)
58
59 respBody, err = io.ReadAll(resp.Body)
60 assert.NoError(err)
61 assert.NotEqual("OK", string(respBody))
62 }
63
64 func TestGetEdgeRoles(t *testing.T) {
65 assert := assert.New(t)
66 c := NewContext(context.Background(), &types.AuthUser{
67 Organization: "test-org",
68 Roles: []string{"EDGE_BANNER_ADMIN", "TEST_ROLE_1", "TEST_ROLE_2"},
69 })
70 edgeRoles, err := GetEdgeRoles(c)
71 assert.NoError(err)
72 assert.Equal(len(edgeRoles), 1)
73 assert.Equal(edgeRoles[0], "EDGE_BANNER_ADMIN")
74
75 c = NewContext(context.Background(), &types.AuthUser{
76 Organization: "test-org",
77 Roles: []string{"TEST_ROLE_1", "TEST_ROLE_2"},
78 })
79 _, err = GetEdgeRoles(c)
80 assert.Error(err)
81 assert.Equal(err.Error(), "edge role not found for context user")
82
83 c = NewContext(context.Background(), &types.AuthUser{
84 Organization: "test-org",
85 Roles: []string{"EDGE_BANNER_ADMIN", "TEST_ROLE_1", "TEST_ROLE_2", "EDGE_ORG_ADMIN"},
86 })
87 edgeRoles, err = GetEdgeRoles(c)
88 assert.NoError(err)
89 assert.Equal(len(edgeRoles), 2)
90 assert.Equal(edgeRoles[0], "EDGE_ORG_ADMIN")
91 assert.Equal(edgeRoles[1], "EDGE_BANNER_ADMIN")
92 }
93
View as plain text