types: make pre auth key use bcrypt (#2853)

hscontrol/db/users_test.go

@@ -1,138 +1,276 @@
 package db
 
 import (
-	"strings"
+	"testing"
 
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/util"
-	"gopkg.in/check.v1"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"gorm.io/gorm"
 	"tailscale.com/types/ptr"
 )
 
-func (s *Suite) TestCreateAndDestroyUser(c *check.C) {
+func TestCreateAndDestroyUser(t *testing.T) {
+	db, err := newSQLiteTestDB()
+	require.NoError(t, err)
+
 	user := db.CreateUserForTest("test")
-	c.Assert(user.Name, check.Equals, "test")
+	assert.Equal(t, "test", user.Name)
 
 	users, err := db.ListUsers()
-	c.Assert(err, check.IsNil)
-	c.Assert(len(users), check.Equals, 1)
+	require.NoError(t, err)
+	assert.Len(t, users, 1)
 
 	err = db.DestroyUser(types.UserID(user.ID))
-	c.Assert(err, check.IsNil)
+	require.NoError(t, err)
 
 	_, err = db.GetUserByID(types.UserID(user.ID))
-	c.Assert(err, check.NotNil)
+	assert.Error(t, err)
 }
 
-func (s *Suite) TestDestroyUserErrors(c *check.C) {
-	err := db.DestroyUser(9998)
-	c.Assert(err, check.Equals, ErrUserNotFound)
+func TestDestroyUserErrors(t *testing.T) {
+	tests := []struct {
+		name string
+		test func(*testing.T, *HSDatabase)
+	}{
+		{
+			name: "error_user_not_found",
+			test: func(t *testing.T, db *HSDatabase) {
+				t.Helper()
 
-	user := db.CreateUserForTest("test")
+				err := db.DestroyUser(9998)
+				assert.ErrorIs(t, err, ErrUserNotFound)
+			},
+		},
+		{
+			name: "success_deletes_preauthkeys",
+			test: func(t *testing.T, db *HSDatabase) {
+				t.Helper()
 
-	pak, err := db.CreatePreAuthKey(types.UserID(user.ID), false, false, nil, nil)
-	c.Assert(err, check.IsNil)
+				user := db.CreateUserForTest("test")
 
-	err = db.DestroyUser(types.UserID(user.ID))
-	c.Assert(err, check.IsNil)
+				pak, err := db.CreatePreAuthKey(types.UserID(user.ID), false, false, nil, nil)
+				require.NoError(t, err)
 
-	result := db.DB.Preload("User").First(&pak, "key = ?", pak.Key)
-	// destroying a user also deletes all associated preauthkeys
-	c.Assert(result.Error, check.Equals, gorm.ErrRecordNotFound)
+				err = db.DestroyUser(types.UserID(user.ID))
+				require.NoError(t, err)
 
-	user, err = db.CreateUser(types.User{Name: "test"})
-	c.Assert(err, check.IsNil)
+				// Verify preauth key was deleted (need to search by prefix for new keys)
+				var foundPak types.PreAuthKey
 
-	pak, err = db.CreatePreAuthKey(types.UserID(user.ID), false, false, nil, nil)
-	c.Assert(err, check.IsNil)
+				result := db.DB.First(&foundPak, "id = ?", pak.ID)
+				assert.ErrorIs(t, result.Error, gorm.ErrRecordNotFound)
+			},
+		},
+		{
+			name: "error_user_has_nodes",
+			test: func(t *testing.T, db *HSDatabase) {
+				t.Helper()
 
-	node := types.Node{
-		ID:             0,
-		Hostname:       "testnode",
-		UserID:         user.ID,
-		RegisterMethod: util.RegisterMethodAuthKey,
-		AuthKeyID:      ptr.To(pak.ID),
+				user, err := db.CreateUser(types.User{Name: "test"})
+				require.NoError(t, err)
+
+				pak, err := db.CreatePreAuthKey(types.UserID(user.ID), false, false, nil, nil)
+				require.NoError(t, err)
+
+				node := types.Node{
+					ID:             0,
+					Hostname:       "testnode",
+					UserID:         user.ID,
+					RegisterMethod: util.RegisterMethodAuthKey,
+					AuthKeyID:      ptr.To(pak.ID),
+				}
+				trx := db.DB.Save(&node)
+				require.NoError(t, trx.Error)
+
+				err = db.DestroyUser(types.UserID(user.ID))
+				assert.ErrorIs(t, err, ErrUserStillHasNodes)
+			},
+		},
 	}
-	trx := db.DB.Save(&node)
-	c.Assert(trx.Error, check.IsNil)
 
-	err = db.DestroyUser(types.UserID(user.ID))
-	c.Assert(err, check.Equals, ErrUserStillHasNodes)
-}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			db, err := newSQLiteTestDB()
+			require.NoError(t, err)
 
-func (s *Suite) TestRenameUser(c *check.C) {
-	userTest := db.CreateUserForTest("test")
-	c.Assert(userTest.Name, check.Equals, "test")
-
-	users, err := db.ListUsers()
-	c.Assert(err, check.IsNil)
-	c.Assert(len(users), check.Equals, 1)
-
-	err = db.RenameUser(types.UserID(userTest.ID), "test-renamed")
-	c.Assert(err, check.IsNil)
-
-	users, err = db.ListUsers(&types.User{Name: "test"})
-	c.Assert(err, check.Equals, nil)
-	c.Assert(len(users), check.Equals, 0)
-
-	users, err = db.ListUsers(&types.User{Name: "test-renamed"})
-	c.Assert(err, check.IsNil)
-	c.Assert(len(users), check.Equals, 1)
-
-	err = db.RenameUser(99988, "test")
-	c.Assert(err, check.Equals, ErrUserNotFound)
-
-	userTest2 := db.CreateUserForTest("test2")
-	c.Assert(userTest2.Name, check.Equals, "test2")
-
-	want := "UNIQUE constraint failed"
-	err = db.RenameUser(types.UserID(userTest2.ID), "test-renamed")
-	if err == nil || !strings.Contains(err.Error(), want) {
-		c.Fatalf("expected failure with unique constraint, want: %q got: %q", want, err)
+			tt.test(t, db)
+		})
 	}
 }
 
-func (s *Suite) TestSetMachineUser(c *check.C) {
-	oldUser := db.CreateUserForTest("old")
-	newUser := db.CreateUserForTest("new")
+func TestRenameUser(t *testing.T) {
+	tests := []struct {
+		name string
+		test func(*testing.T, *HSDatabase)
+	}{
+		{
+			name: "success_rename",
+			test: func(t *testing.T, db *HSDatabase) {
+				t.Helper()
 
-	pak, err := db.CreatePreAuthKey(types.UserID(oldUser.ID), false, false, nil, nil)
-	c.Assert(err, check.IsNil)
+				userTest := db.CreateUserForTest("test")
+				assert.Equal(t, "test", userTest.Name)
 
-	node := types.Node{
-		ID:             12,
-		Hostname:       "testnode",
-		UserID:         oldUser.ID,
-		RegisterMethod: util.RegisterMethodAuthKey,
-		AuthKeyID:      ptr.To(pak.ID),
+				users, err := db.ListUsers()
+				require.NoError(t, err)
+				assert.Len(t, users, 1)
+
+				err = db.RenameUser(types.UserID(userTest.ID), "test-renamed")
+				require.NoError(t, err)
+
+				users, err = db.ListUsers(&types.User{Name: "test"})
+				require.NoError(t, err)
+				assert.Empty(t, users)
+
+				users, err = db.ListUsers(&types.User{Name: "test-renamed"})
+				require.NoError(t, err)
+				assert.Len(t, users, 1)
+			},
+		},
+		{
+			name: "error_user_not_found",
+			test: func(t *testing.T, db *HSDatabase) {
+				t.Helper()
+
+				err := db.RenameUser(99988, "test")
+				assert.ErrorIs(t, err, ErrUserNotFound)
+			},
+		},
+		{
+			name: "error_duplicate_name",
+			test: func(t *testing.T, db *HSDatabase) {
+				t.Helper()
+
+				userTest := db.CreateUserForTest("test")
+				userTest2 := db.CreateUserForTest("test2")
+
+				assert.Equal(t, "test", userTest.Name)
+				assert.Equal(t, "test2", userTest2.Name)
+
+				err := db.RenameUser(types.UserID(userTest2.ID), "test")
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), "UNIQUE constraint failed")
+			},
+		},
 	}
-	trx := db.DB.Save(&node)
-	c.Assert(trx.Error, check.IsNil)
-	c.Assert(node.UserID, check.Equals, oldUser.ID)
 
-	err = db.Write(func(tx *gorm.DB) error {
-		return AssignNodeToUser(tx, 12, types.UserID(newUser.ID))
-	})
-	c.Assert(err, check.IsNil)
-	// Reload node from database to see updated values
-	updatedNode, err := db.GetNodeByID(12)
-	c.Assert(err, check.IsNil)
-	c.Assert(updatedNode.UserID, check.Equals, newUser.ID)
-	c.Assert(updatedNode.User.Name, check.Equals, newUser.Name)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			db, err := newSQLiteTestDB()
+			require.NoError(t, err)
 
-	err = db.Write(func(tx *gorm.DB) error {
-		return AssignNodeToUser(tx, 12, 9584849)
-	})
-	c.Assert(err, check.Equals, ErrUserNotFound)
-
-	err = db.Write(func(tx *gorm.DB) error {
-		return AssignNodeToUser(tx, 12, types.UserID(newUser.ID))
-	})
-	c.Assert(err, check.IsNil)
-	// Reload node from database again to see updated values
-	finalNode, err := db.GetNodeByID(12)
-	c.Assert(err, check.IsNil)
-	c.Assert(finalNode.UserID, check.Equals, newUser.ID)
-	c.Assert(finalNode.User.Name, check.Equals, newUser.Name)
+			tt.test(t, db)
+		})
+	}
+}
+
+func TestAssignNodeToUser(t *testing.T) {
+	tests := []struct {
+		name string
+		test func(*testing.T, *HSDatabase)
+	}{
+		{
+			name: "success_reassign_node",
+			test: func(t *testing.T, db *HSDatabase) {
+				t.Helper()
+
+				oldUser := db.CreateUserForTest("old")
+				newUser := db.CreateUserForTest("new")
+
+				pak, err := db.CreatePreAuthKey(types.UserID(oldUser.ID), false, false, nil, nil)
+				require.NoError(t, err)
+
+				node := types.Node{
+					ID:             12,
+					Hostname:       "testnode",
+					UserID:         oldUser.ID,
+					RegisterMethod: util.RegisterMethodAuthKey,
+					AuthKeyID:      ptr.To(pak.ID),
+				}
+				trx := db.DB.Save(&node)
+				require.NoError(t, trx.Error)
+				assert.Equal(t, oldUser.ID, node.UserID)
+
+				err = db.Write(func(tx *gorm.DB) error {
+					return AssignNodeToUser(tx, 12, types.UserID(newUser.ID))
+				})
+				require.NoError(t, err)
+
+				// Reload node from database to see updated values
+				updatedNode, err := db.GetNodeByID(12)
+				require.NoError(t, err)
+				assert.Equal(t, newUser.ID, updatedNode.UserID)
+				assert.Equal(t, newUser.Name, updatedNode.User.Name)
+			},
+		},
+		{
+			name: "error_user_not_found",
+			test: func(t *testing.T, db *HSDatabase) {
+				t.Helper()
+
+				oldUser := db.CreateUserForTest("old")
+
+				pak, err := db.CreatePreAuthKey(types.UserID(oldUser.ID), false, false, nil, nil)
+				require.NoError(t, err)
+
+				node := types.Node{
+					ID:             12,
+					Hostname:       "testnode",
+					UserID:         oldUser.ID,
+					RegisterMethod: util.RegisterMethodAuthKey,
+					AuthKeyID:      ptr.To(pak.ID),
+				}
+				trx := db.DB.Save(&node)
+				require.NoError(t, trx.Error)
+
+				err = db.Write(func(tx *gorm.DB) error {
+					return AssignNodeToUser(tx, 12, 9584849)
+				})
+				assert.ErrorIs(t, err, ErrUserNotFound)
+			},
+		},
+		{
+			name: "success_reassign_to_same_user",
+			test: func(t *testing.T, db *HSDatabase) {
+				t.Helper()
+
+				user := db.CreateUserForTest("user")
+
+				pak, err := db.CreatePreAuthKey(types.UserID(user.ID), false, false, nil, nil)
+				require.NoError(t, err)
+
+				node := types.Node{
+					ID:             12,
+					Hostname:       "testnode",
+					UserID:         user.ID,
+					RegisterMethod: util.RegisterMethodAuthKey,
+					AuthKeyID:      ptr.To(pak.ID),
+				}
+				trx := db.DB.Save(&node)
+				require.NoError(t, trx.Error)
+
+				err = db.Write(func(tx *gorm.DB) error {
+					return AssignNodeToUser(tx, 12, types.UserID(user.ID))
+				})
+				require.NoError(t, err)
+
+				// Reload node from database again to see updated values
+				finalNode, err := db.GetNodeByID(12)
+				require.NoError(t, err)
+				assert.Equal(t, user.ID, finalNode.UserID)
+				assert.Equal(t, user.Name, finalNode.User.Name)
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			db, err := newSQLiteTestDB()
+			require.NoError(t, err)
+
+			tt.test(t, db)
+		})
+	}
 }