package main import ( library "git.ailur.dev/ailur/fg-library/v3" nucleusLibrary "git.ailur.dev/ailur/fg-nucleus-library" "bytes" "os" "database/sql" "errors" "path/filepath" "github.com/google/uuid" ) var ServiceInformation = library.Service{ Name: "Storage", Permissions: library.Permissions{ Authenticate: false, // This service does not require authentication Database: true, // This service requires database access to store quotas BlobStorage: false, // This service *is* the blob storage InterServiceCommunication: true, // This service does require inter-service communication }, ServiceID: uuid.MustParse("00000000-0000-0000-0000-000000000003"), } var ( conn library.Database loggerService = uuid.MustParse("00000000-0000-0000-0000-000000000002") ) func logFunc(message string, messageType library.MessageCode, information *library.ServiceInitializationInformation) { // Log the message to the logger service information.SendISMessage(loggerService, messageType, message) } func respondError(message library.InterServiceMessage, err error, information *library.ServiceInitializationInformation, myFault bool) { // Respond with an error message var errCode = library.BadRequest if myFault { // Log the error message to the logger service logFunc(err.Error(), 2, information) errCode = library.InternalError } message.Respond(errCode, err, information) } func checkUserExists(userID uuid.UUID) bool { // Check if a user exists in the database var userCheck []byte err := conn.DB.QueryRow("SELECT id FROM users WHERE id = $1", userID[:]).Scan(&userCheck) if err != nil { if errors.Is(err, sql.ErrNoRows) { return false } else { return false } } else { return bytes.Equal(userCheck, userID[:]) } } // addQuota can be used with a negative quota to remove quota from a user func addQuota(information *library.ServiceInitializationInformation, message library.InterServiceMessage) { // Add more quota to a user userID := message.Message.(nucleusLibrary.Quota).User if checkUserExists(userID) { _, err := conn.DB.Exec("UPDATE users SET quota = quota + $1 WHERE id = $2", message.Message.(nucleusLibrary.Quota).Bytes, message.Message.(nucleusLibrary.Quota).User) if err != nil { respondError(message, err, information, true) } } else { _, err := conn.DB.Exec("INSERT INTO users (id, quota, reserved) VALUES ($1, $2, 0)", userID[:], int64(information.Configuration["defaultQuota"].(int))+message.Message.(nucleusLibrary.Quota).Bytes) if err != nil { respondError(message, err, information, true) } } // Success message.Respond(library.Success, nil, information) } // And so does addReserved func addReserved(information *library.ServiceInitializationInformation, message library.InterServiceMessage) { // Add more reserved space to a user userID := message.Message.(nucleusLibrary.Quota).User if checkUserExists(userID) { // Check if the user has enough space quota, err := getQuota(userID) if err != nil { respondError(message, err, information, true) } used, err := getUsed(userID, information) if err != nil { respondError(message, err, information, true) } if used+message.Message.(nucleusLibrary.Quota).Bytes > quota { respondError(message, errors.New("insufficient storage"), information, false) return } _, err = conn.DB.Exec("UPDATE users SET reserved = reserved + $1 WHERE id = $2", message.Message.(nucleusLibrary.Quota).Bytes, userID[:]) if err != nil { respondError(message, err, information, true) } } else { // Check if the user has enough space if int64(information.Configuration["defaultQuota"].(int)) < message.Message.(nucleusLibrary.Quota).Bytes { respondError(message, errors.New("insufficient storage"), information, false) return } _, err := conn.DB.Exec("INSERT INTO users (id, quota, reserved) VALUES ($1, $2, $3)", userID[:], int64(information.Configuration["defaultQuota"].(int)), message.Message.(nucleusLibrary.Quota).Bytes) if err != nil { respondError(message, err, information, true) } } // Success message.Respond(library.Success, nil, information) } func getQuota(userID uuid.UUID) (int64, error) { // Get the quota for a user var quota int64 err := conn.DB.QueryRow("SELECT quota FROM users WHERE id = $1", userID[:]).Scan("a) if err != nil { return 0, err } return quota, nil } func getUsed(userID uuid.UUID, information *library.ServiceInitializationInformation) (int64, error) { // Get the used space for a user by first getting the reserved space from file storage _, err := os.Stat(filepath.Join(information.Configuration["path"].(string), userID.String())) if os.IsNotExist(err) { // Create the directory err = os.Mkdir(filepath.Join(information.Configuration["path"].(string), userID.String()), 0755) if err != nil { return 0, err } } var used int64 err = filepath.Walk(filepath.Join(information.Configuration["path"].(string), userID.String()), func(path string, entry os.FileInfo, err error) error { if err != nil { return err } used += entry.Size() return nil }) if err != nil { return 0, err } // Then add the reserved space from the database var reserved int64 err = conn.DB.QueryRow("SELECT reserved FROM users WHERE id = $1", userID[:]).Scan(&reserved) if err != nil { return 0, err } return used + reserved, nil } func modifyFile(information *library.ServiceInitializationInformation, message library.InterServiceMessage) { // Check if the file already exists path := filepath.Join(information.Configuration["path"].(string), message.Message.(nucleusLibrary.File).User.String(), message.Message.(nucleusLibrary.File).Name) logFunc(path, 0, information) _, err := os.Stat(path) if err == nil { // Delete the file err = os.Remove(path) if err != nil { respondError(message, err, information, true) } } else if !os.IsNotExist(err) { respondError(message, err, information, true) } // Check if the user has enough space quota, err := getQuota(message.Message.(nucleusLibrary.File).User) if err != nil { respondError(message, err, information, true) } used, err := getUsed(message.Message.(nucleusLibrary.File).User, information) if err != nil { respondError(message, err, information, true) } if used+int64(len(message.Message.(nucleusLibrary.File).Bytes)) > quota { respondError(message, errors.New("insufficient storage"), information, false) return } // Add a file to the user's storage file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0644) if err != nil { respondError(message, err, information, true) } // Write the file _, err = file.Write(message.Message.(nucleusLibrary.File).Bytes) if err != nil { respondError(message, err, information, true) } // Close the file err = file.Close() if err != nil { respondError(message, err, information, true) } // Success message.Respond(library.Success, nil, information) } func getFile(information *library.ServiceInitializationInformation, message library.InterServiceMessage) { // Check if the file exists path := filepath.Join(information.Configuration["path"].(string), message.Message.(nucleusLibrary.File).User.String(), message.Message.(nucleusLibrary.File).Name) _, err := os.Stat(path) if os.IsNotExist(err) { respondError(message, errors.New("file not found"), information, false) return } // Open the file file, err := os.Open(path) if err != nil { respondError(message, err, information, true) } // Respond with the file // It's their responsibility to close the file message.Respond(library.Success, file, information) } func deleteFile(information *library.ServiceInitializationInformation, message library.InterServiceMessage) { // Check if the file exists path := filepath.Join(information.Configuration["path"].(string), message.Message.(nucleusLibrary.File).User.String(), message.Message.(nucleusLibrary.File).Name) _, err := os.Stat(path) if os.IsNotExist(err) { respondError(message, errors.New("file not found"), information, false) return } // Delete the file err = os.Remove(path) if err != nil { respondError(message, err, information, true) } // Success message.Respond(library.Success, nil, information) } // processInterServiceMessages listens for incoming messages and processes them func processInterServiceMessages(information *library.ServiceInitializationInformation) { // Listen for incoming messages for { message := information.AcceptMessage() switch message.MessageType { case 1: // Add quota addQuota(information, message) case 2: // Add reserved addReserved(information, message) case 3: // Modify file modifyFile(information, message) case 4: // Get file getFile(information, message) case 5: deleteFile(information, message) default: // Respond with an error message respondError(message, errors.New("invalid message type"), information, false) } } } func Main(information *library.ServiceInitializationInformation) { // Start up the ISM processor information.StartISProcessor() // Get the database connection conn, err := information.GetDatabase() if err != nil { logFunc(err.Error(), 3, information) } // Create the quotas table if it doesn't exist if conn.DBType == library.Sqlite { _, err := conn.DB.Exec("CREATE TABLE IF NOT EXISTS users (id BLOB PRIMARY KEY, quota BIGINT, reserved BIGINT)") if err != nil { logFunc(err.Error(), 3, information) } } else { _, err := conn.DB.Exec("CREATE TABLE IF NOT EXISTS users (id BYTEA PRIMARY KEY, quota BIGINT, reserved BIGINT)") if err != nil { logFunc(err.Error(), 3, information) } } // Listen for incoming messages go processInterServiceMessages(information) }