122 lines
4.2 KiB
Go
122 lines
4.2 KiB
Go
package handler
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/go-sql-driver/mysql"
|
|
"github.com/wmms/wmms/internal/model"
|
|
"github.com/wmms/wmms/internal/service"
|
|
"github.com/wmms/wmms/pkg/database"
|
|
)
|
|
|
|
// This one request-level integration scenario covers database revocation and
|
|
// stale JWT roles. It uses a disposable schema, never a production database.
|
|
func TestOrganizationContextRevocationMySQL(t *testing.T) {
|
|
dsn := os.Getenv("TEST_DB_DSN")
|
|
if dsn == "" {
|
|
t.Skip("real MySQL integration requires TEST_DB_DSN")
|
|
}
|
|
cfg, err := mysql.ParseDSN(dsn)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !strings.HasPrefix(cfg.DBName, "company_isolation_test") {
|
|
t.Fatal("TEST_DB_DSN must name a company_isolation_test schema")
|
|
}
|
|
cfg.DBName = ""
|
|
cfg.MultiStatements = true
|
|
cfg.ParseTime = true
|
|
admin, err := sql.Open("mysql", cfg.FormatDSN())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
name := fmt.Sprintf("company_isolation_test_context_%d", time.Now().UnixNano())
|
|
if _, err = admin.Exec("CREATE DATABASE `" + name + "` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"); err != nil {
|
|
admin.Close()
|
|
t.Fatal(err)
|
|
}
|
|
cfg.DBName = name
|
|
db, err := sql.Open("mysql", cfg.FormatDSN())
|
|
if err != nil {
|
|
admin.Exec("DROP DATABASE `" + name + "`")
|
|
admin.Close()
|
|
t.Fatal(err)
|
|
}
|
|
previous := database.DB
|
|
database.DB = db
|
|
t.Cleanup(func() { database.DB = previous; db.Close(); admin.Exec("DROP DATABASE `" + name + "`"); admin.Close() })
|
|
content, err := os.ReadFile(filepath.Join("..", "..", "scripts", "init.sql"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
lines := []string{}
|
|
for _, line := range strings.Split(strings.TrimPrefix(string(content), "\ufeff"), "\n") {
|
|
if strings.HasPrefix(line, "CREATE DATABASE") || strings.HasPrefix(line, "USE ") {
|
|
continue
|
|
}
|
|
lines = append(lines, line)
|
|
}
|
|
exec := func(query string) {
|
|
t.Helper()
|
|
if _, err := db.Exec(query); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
exec(strings.Join(lines, "\n"))
|
|
exec(`INSERT INTO organizations(id,code,name,status) VALUES(2,'CONTEXT_TWO','公司二','active');
|
|
INSERT INTO users(id,username,password) VALUES(7,'context-user','x'),(8,'context-admin','x');
|
|
INSERT INTO user_roles(user_id,role_id) SELECT 7,id FROM roles WHERE code='warehouse_admin';
|
|
INSERT INTO user_roles(user_id,role_id) SELECT 8,id FROM roles WHERE code='super_admin';
|
|
INSERT INTO user_organizations(user_id,organization_id,is_default) VALUES(7,2,1);`)
|
|
token := contextToken(t)
|
|
adminToken, err := service.NewAuthService().GenerateToken(&model.User{ID: 8, Roles: []model.Role{{Code: "super_admin", Enabled: true}}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
gin.SetMode(gin.TestMode)
|
|
router := gin.New()
|
|
NewHandler().RegisterRoutes(router)
|
|
request := func(token string, status int) {
|
|
t.Helper()
|
|
req := httptest.NewRequest("GET", "/api/units", nil)
|
|
req.Header.Set("Authorization", "Bearer "+token)
|
|
req.Header.Set("X-Organization-ID", "2")
|
|
w := httptest.NewRecorder()
|
|
router.ServeHTTP(w, req)
|
|
if w.Code != status {
|
|
t.Fatalf("status=%d want=%d body=%s", w.Code, status, w.Body.String())
|
|
}
|
|
}
|
|
request(token, 200)
|
|
exec(`DELETE FROM user_organizations WHERE user_id=7 AND organization_id=2`)
|
|
request(token, 403)
|
|
exec(`INSERT INTO user_organizations(user_id,organization_id,is_default) VALUES(7,2,1)`)
|
|
request(token, 200)
|
|
exec(`UPDATE organizations SET status='inactive' WHERE id=2`)
|
|
request(token, 403)
|
|
exec(`UPDATE organizations SET status='active' WHERE id=2`)
|
|
request(adminToken, 200)
|
|
exec(`DELETE FROM user_roles WHERE user_id=8;
|
|
INSERT INTO user_roles(user_id,role_id) SELECT 8,id FROM roles WHERE code='warehouse_admin';`)
|
|
request(adminToken, 403)
|
|
// Explicit authorization still works after downgrading the role, while the
|
|
// same old super-admin JWT can no longer pass global administration guards.
|
|
exec(`INSERT INTO user_organizations(user_id,organization_id,is_default) VALUES(8,2,1)`)
|
|
request(adminToken, 200)
|
|
req := httptest.NewRequest("GET", "/api/organizations", nil)
|
|
req.Header.Set("Authorization", "Bearer "+adminToken)
|
|
w := httptest.NewRecorder()
|
|
router.ServeHTTP(w, req)
|
|
if w.Code != 403 {
|
|
t.Fatalf("stale JWT retained admin privileges: %d %s", w.Code, w.Body.String())
|
|
}
|
|
}
|