package middleware import ( "errors" "net/http" "net/http/httptest" "testing" "github.com/DATA-DOG/go-sqlmock" "github.com/gin-gonic/gin" "github.com/wmms/wmms/internal/model" "github.com/wmms/wmms/internal/service" "github.com/wmms/wmms/pkg/database" ) func organizationMock(t *testing.T) sqlmock.Sqlmock { t.Helper() db, mock, err := sqlmock.New() if err != nil { t.Fatal(err) } previous := database.DB database.DB = db t.Cleanup(func() { database.DB = previous; db.Close() }) return mock } func TestOrganizationMiddlewareHeadersAndAccess(t *testing.T) { gin.SetMode(gin.TestMode) for _, tc := range []struct { name, header string allowed bool status int }{ {"missing", "", false, 400}, {"zero", "0", false, 400}, {"negative", "-1", false, 400}, {"signed", "+1", false, 400}, {"space", " 1", false, 400}, {"fraction", "1.2", false, 400}, {"text", "abc", false, 400}, {"overflow", "9223372036854775808", false, 400}, {"ungranted", "2", false, 403}, {"nonexistent", "999", false, 403}, {"inactive", "3", false, 403}, {"allowed", "2", true, 204}, } { t.Run(tc.name, func(t *testing.T) { mock := organizationMock(t) if tc.status != 400 { mock.ExpectQuery("SELECT EXISTS").WithArgs(sqlmock.AnyArg(), int64(7), int64(7)).WillReturnRows(sqlmock.NewRows([]string{"allowed"}).AddRow(tc.allowed)) } router := gin.New() router.GET("/", func(c *gin.Context) { c.Set(UserIDKey, int64(7)) }, OrganizationMiddleware(), func(c *gin.Context) { if GetCurrentOrganizationID(c) != 2 { t.Fatal("trusted organization missing") } c.Request.Header.Set("X-Organization-ID", "999") if GetCurrentOrganizationID(c) != 2 { t.Fatal("header mutated trusted context") } c.Status(204) }) req := httptest.NewRequest("GET", "/", nil) req.Header.Set("X-Organization-ID", tc.header) w := httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != tc.status { t.Fatalf("status=%d body=%s", w.Code, w.Body.String()) } if err := mock.ExpectationsWereMet(); err != nil { t.Fatal(err) } }) } } func TestOrganizationMiddlewareRechecksRevocationOnNextRequest(t *testing.T) { for _, name := range []string{"grant revoked", "company disabled", "super admin revoked"} { t.Run(name, func(t *testing.T) { mock := organizationMock(t) token, err := service.NewAuthService().GenerateToken(&model.User{ID: 7, Roles: []model.Role{{Code: "super_admin", Enabled: true}}}) if err != nil { t.Fatal(err) } router := gin.New() router.GET("/", AuthMiddleware(), OrganizationMiddleware(), func(c *gin.Context) { c.Status(204) }) for i, allowed := range []bool{true, false} { role := "warehouse_admin" if name == "super admin revoked" && i == 0 { role = "super_admin" } mock.ExpectQuery("SELECT u.id, r.code FROM users u").WithArgs(int64(7)).WillReturnRows(sqlmock.NewRows([]string{"id", "code"}).AddRow(7, role)) mock.ExpectQuery("SELECT EXISTS").WithArgs(int64(2), int64(7), int64(7)).WillReturnRows(sqlmock.NewRows([]string{"allowed"}).AddRow(allowed)) req := httptest.NewRequest("GET", "/", nil) req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("X-Organization-ID", "2") w := httptest.NewRecorder() router.ServeHTTP(w, req) want := 204 if !allowed { want = 403 } if w.Code != want { t.Fatalf("request %d status=%d", i, w.Code) } } if err := mock.ExpectationsWereMet(); err != nil { t.Fatal(err) } }) } } func TestAuthMiddlewareUsesCurrentRolesAndFailsClosed(t *testing.T) { for _, tc := range []struct { name string role *string queryErr error missing bool status int }{ {name: "stale admin", role: stringPointer("warehouse_admin"), status: 403}, {name: "disabled roles", status: 403}, {name: "deleted user", missing: true, status: 401}, {name: "database failure", queryErr: errors.New("offline"), status: 500}, {name: "current admin", role: stringPointer("super_admin"), status: 204}, } { t.Run(tc.name, func(t *testing.T) { mock := organizationMock(t) query := mock.ExpectQuery("SELECT u.id, r.code FROM users u").WithArgs(int64(7)) rows := sqlmock.NewRows([]string{"id", "code"}) if !tc.missing { rows.AddRow(7, tc.role) } if tc.queryErr != nil { query.WillReturnError(tc.queryErr) } else { query.WillReturnRows(rows) } token, _ := service.NewAuthService().GenerateToken(&model.User{ID: 7, Roles: []model.Role{{Code: "super_admin", Enabled: true}}}) router := gin.New() router.GET("/", AuthMiddleware(), AdminOnly(), func(c *gin.Context) { c.Status(http.StatusNoContent) }) req := httptest.NewRequest("GET", "/", nil) req.Header.Set("Authorization", "Bearer "+token) w := httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != tc.status { t.Fatalf("status=%d body=%s", w.Code, w.Body.String()) } if err := mock.ExpectationsWereMet(); err != nil { t.Fatal(err) } }) } } func stringPointer(s string) *string { return &s }