From a82fb3a49c0d34bd184405e60e8ed25513f23f34 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Sun, 18 Aug 2024 14:09:52 +0800
Subject: [PATCH] :sparkles: Attachment has pool

---
 .idea/workspace.xml                        | 20 +++++++-------
 go.mod                                     |  6 ++---
 go.sum                                     |  2 --
 pkg/internal/models/attachments.go         |  7 +++--
 pkg/internal/models/pools.go               |  6 +++--
 pkg/internal/server/api/attachments_api.go | 23 +++++++++++-----
 pkg/internal/services/analyzer.go          |  2 +-
 pkg/internal/services/attachments.go       | 12 +++++++--
 pkg/internal/services/pools.go             |  8 ++++++
 pkg/internal/services/random_id.go         | 31 ++++++++++++++++++++++
 settings.toml                              | 17 ++++++++----
 11 files changed, 102 insertions(+), 32 deletions(-)
 create mode 100644 pkg/internal/services/random_id.go

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index c8bf724..2eeec73 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -4,17 +4,18 @@
     <option name="autoReloadType" value="ALL" />
   </component>
   <component name="ChangeListManager">
-    <list default="true" id="18dd0d68-b4b8-40db-9734-9119b5c848bd" name="更改" comment=":sparkles: List attachment original filter">
-      <change afterPath="$PROJECT_DIR$/pkg/internal/models/pools.go" afterDir="false" />
-      <change afterPath="$PROJECT_DIR$/pkg/internal/server/api/pools_api.go" afterDir="false" />
-      <change afterPath="$PROJECT_DIR$/pkg/internal/services/pools.go" afterDir="false" />
+    <list default="true" id="18dd0d68-b4b8-40db-9734-9119b5c848bd" name="更改" comment=":sparkles: Attachment pool basis">
+      <change afterPath="$PROJECT_DIR$/pkg/internal/services/random_id.go" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/pkg/internal/database/migrator.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/internal/database/migrator.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/pkg/internal/models/accounts.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/internal/models/accounts.go" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/go.mod" beforeDir="false" afterPath="$PROJECT_DIR$/go.mod" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/go.sum" beforeDir="false" afterPath="$PROJECT_DIR$/go.sum" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/pkg/internal/models/attachments.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/internal/models/attachments.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/pkg/internal/models/metadata.go" beforeDir="false" />
+      <change beforePath="$PROJECT_DIR$/pkg/internal/models/pools.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/internal/models/pools.go" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/pkg/internal/server/api/attachments_api.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/internal/server/api/attachments_api.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/pkg/internal/server/api/index.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/internal/server/api/index.go" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/pkg/internal/services/analyzer.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/internal/services/analyzer.go" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/pkg/internal/services/attachments.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/internal/services/attachments.go" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/pkg/internal/services/pools.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/internal/services/pools.go" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/settings.toml" beforeDir="false" afterPath="$PROJECT_DIR$/settings.toml" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -127,7 +128,8 @@
     <MESSAGE value=":recycle: Moved onto dealer" />
     <MESSAGE value=":sparkles: Self reference detection" />
     <MESSAGE value=":sparkles: List attachment original filter" />
-    <option name="LAST_COMMIT_MESSAGE" value=":sparkles: List attachment original filter" />
+    <MESSAGE value=":sparkles: Attachment pool basis" />
+    <option name="LAST_COMMIT_MESSAGE" value=":sparkles: Attachment pool basis" />
   </component>
   <component name="VgoProject">
     <settings-migrated>true</settings-migrated>
diff --git a/go.mod b/go.mod
index 71669be..ad9e58e 100644
--- a/go.mod
+++ b/go.mod
@@ -8,12 +8,15 @@ require (
 	github.com/gofiber/fiber/v2 v2.52.4
 	github.com/google/uuid v1.6.0
 	github.com/json-iterator/go v1.1.12
+	github.com/k0kubun/go-ansi v0.0.0-20180517002512-3bf9e2903213
 	github.com/minio/minio-go/v7 v7.0.70
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/rs/zerolog v1.31.0
 	github.com/samber/lo v1.39.0
+	github.com/schollz/progressbar/v3 v3.14.4
 	github.com/spf13/viper v1.18.2
 	google.golang.org/grpc v1.64.0
+	gopkg.in/vansante/go-ffprobe.v2 v2.2.0
 	gorm.io/datatypes v1.2.0
 	gorm.io/driver/postgres v1.5.4
 	gorm.io/gorm v1.25.6
@@ -48,7 +51,6 @@ require (
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
 	github.com/jpillora/backoff v1.0.0 // indirect
-	github.com/k0kubun/go-ansi v0.0.0-20180517002512-3bf9e2903213 // indirect
 	github.com/klauspost/compress v1.17.8 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.7 // indirect
 	github.com/leodido/go-urn v1.2.4 // indirect
@@ -70,7 +72,6 @@ require (
 	github.com/rs/xid v1.5.0 // indirect
 	github.com/sagikazarmark/locafero v0.4.0 // indirect
 	github.com/sagikazarmark/slog-shim v0.1.0 // indirect
-	github.com/schollz/progressbar/v3 v3.14.4 // indirect
 	github.com/sourcegraph/conc v0.3.0 // indirect
 	github.com/spf13/afero v1.11.0 // indirect
 	github.com/spf13/cast v1.6.0 // indirect
@@ -91,7 +92,6 @@ require (
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240604185151-ef581f913117 // indirect
 	google.golang.org/protobuf v1.34.2 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
-	gopkg.in/vansante/go-ffprobe.v2 v2.2.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	gorm.io/driver/mysql v1.5.2 // indirect
 )
diff --git a/go.sum b/go.sum
index e70313c..07166e3 100644
--- a/go.sum
+++ b/go.sum
@@ -362,8 +362,6 @@ golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
-golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI=
 golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
diff --git a/pkg/internal/models/attachments.go b/pkg/internal/models/attachments.go
index 29ce7bc..223816b 100644
--- a/pkg/internal/models/attachments.go
+++ b/pkg/internal/models/attachments.go
@@ -15,11 +15,14 @@ const (
 type Attachment struct {
 	BaseModel
 
-	Uuid        string        `json:"uuid"`
+	// Random ID is for accessing (appear in URL)
+	Rid string `json:"rid" gorm:"uniqueIndex"`
+	// Unique ID is for storing (appear in local file name or object name)
+	Uuid string `json:"uuid"`
+
 	Size        int64         `json:"size"`
 	Name        string        `json:"name"`
 	Alternative string        `json:"alt"`
-	Usage       string        `json:"usage"`
 	MimeType    string        `json:"mimetype"`
 	HashCode    string        `json:"hash"`
 	Destination AttachmentDst `json:"destination"`
diff --git a/pkg/internal/models/pools.go b/pkg/internal/models/pools.go
index 6440891..ff3d438 100644
--- a/pkg/internal/models/pools.go
+++ b/pkg/internal/models/pools.go
@@ -17,6 +17,8 @@ type AttachmentPool struct {
 }
 
 type AttachmentPoolConfig struct {
-	ExistLifecycle     *int `json:"exist_lifecycle"`
-	IsPublicAccessible bool `json:"is_public_accessible"`
+	MaxFileSize           *int64 `json:"max_file_size"`
+	ExistLifecycle        *int64 `json:"exist_lifecycle"`
+	AllowCrossPoolIngress bool   `json:"allow_cross_pool_ingress"`
+	AllowCrossPoolEgress  bool   `json:"allow_cross_pool_egress"`
 }
diff --git a/pkg/internal/server/api/attachments_api.go b/pkg/internal/server/api/attachments_api.go
index 34cc183..1f408bb 100644
--- a/pkg/internal/server/api/attachments_api.go
+++ b/pkg/internal/server/api/attachments_api.go
@@ -85,9 +85,18 @@ func createAttachment(c *fiber.Ctx) error {
 	}
 	user := c.Locals("user").(models.Account)
 
-	usage := c.FormValue("usage")
-	if !lo.Contains(viper.GetStringSlice("accepts_usage"), usage) {
-		return fiber.NewError(fiber.StatusBadRequest, fmt.Sprintf("disallowed usage: %s", usage))
+	poolAlias := c.FormValue("pool")
+	if len(poolAlias) == 0 {
+		poolAlias = c.FormValue("usage")
+	}
+	aliasingMap := viper.GetStringMapString("pools.aliases")
+	if val, ok := aliasingMap[poolAlias]; ok {
+		poolAlias = val
+	}
+
+	pool, err := services.GetAttachmentPoolByAlias(poolAlias)
+	if err != nil {
+		return fiber.NewError(fiber.StatusBadRequest, fmt.Sprintf("unable to get attachment pool info: %v", err))
 	}
 
 	file, err := c.FormFile("file")
@@ -97,6 +106,8 @@ func createAttachment(c *fiber.Ctx) error {
 
 	if err = gap.H.EnsureGrantedPerm(c, "CreateAttachments", file.Size); err != nil {
 		return err
+	} else if pool.Config.Data().MaxFileSize != nil && file.Size > *pool.Config.Data().MaxFileSize {
+		return fiber.NewError(fiber.StatusBadRequest, fmt.Sprintf("attachment pool %s doesn't allow file larger than %d", pool.Alias, *pool.Config.Data().MaxFileSize))
 	}
 
 	usermeta := make(map[string]any)
@@ -105,13 +116,14 @@ func createAttachment(c *fiber.Ctx) error {
 	tx := database.C.Begin()
 
 	metadata, err := services.NewAttachmentMetadata(tx, user, file, models.Attachment{
-		Usage:       usage,
 		Alternative: c.FormValue("alt"),
 		MimeType:    c.FormValue("mimetype"),
 		Metadata:    usermeta,
 		IsMature:    len(c.FormValue("mature")) > 0,
 		IsAnalyzed:  false,
 		Destination: models.AttachmentDstTemporary,
+		Pool:        &pool,
+		PoolID:      &pool.ID,
 	})
 	if err != nil {
 		tx.Rollback()
@@ -126,6 +138,7 @@ func createAttachment(c *fiber.Ctx) error {
 	tx.Commit()
 
 	metadata.Account = user
+	metadata.Pool = &pool
 	services.PublishAnalyzeTask(metadata)
 
 	return c.JSON(metadata)
@@ -141,7 +154,6 @@ func updateAttachmentMeta(c *fiber.Ctx) error {
 
 	var data struct {
 		Alternative string `json:"alt"`
-		Usage       string `json:"usage"`
 		IsMature    bool   `json:"is_mature"`
 	}
 
@@ -155,7 +167,6 @@ func updateAttachmentMeta(c *fiber.Ctx) error {
 	}
 
 	attachment.Alternative = data.Alternative
-	attachment.Usage = data.Usage
 	attachment.IsMature = data.IsMature
 
 	if attachment, err := services.UpdateAttachment(attachment); err != nil {
diff --git a/pkg/internal/services/analyzer.go b/pkg/internal/services/analyzer.go
index e69413a..caa28ca 100644
--- a/pkg/internal/services/analyzer.go
+++ b/pkg/internal/services/analyzer.go
@@ -99,7 +99,7 @@ func AnalyzeAttachment(file models.Attachment) error {
 
 	var start time.Time
 
-	// Do analyze job
+	// Do analyze jobs
 	if !file.IsAnalyzed || len(file.HashCode) == 0 {
 		destMap := viper.GetStringMap("destinations.temporary")
 
diff --git a/pkg/internal/services/attachments.go b/pkg/internal/services/attachments.go
index 5119dd2..00cf60c 100644
--- a/pkg/internal/services/attachments.go
+++ b/pkg/internal/services/attachments.go
@@ -28,7 +28,7 @@ func GetAttachmentByID(id uint) (models.Attachment, error) {
 	var attachment models.Attachment
 	if err := database.C.Where(models.Attachment{
 		BaseModel: models.BaseModel{ID: id},
-	}).Preload("Account").First(&attachment).Error; err != nil {
+	}).Preload("Pool").Preload("Account").First(&attachment).Error; err != nil {
 		return attachment, err
 	} else {
 		MaintainAttachmentCache()
@@ -42,7 +42,7 @@ func GetAttachmentByHash(hash string) (models.Attachment, error) {
 	var attachment models.Attachment
 	if err := database.C.Where(models.Attachment{
 		HashCode: hash,
-	}).First(&attachment).Error; err != nil {
+	}).Preload("Pool").First(&attachment).Error; err != nil {
 		return attachment, err
 	}
 	return attachment, nil
@@ -60,6 +60,7 @@ func CacheAttachment(id uint, item models.Attachment) {
 }
 
 func NewAttachmentMetadata(tx *gorm.DB, user models.Account, file *multipart.FileHeader, attachment models.Attachment) (models.Attachment, error) {
+	attachment.Rid = RandString(16)
 	attachment.Uuid = uuid.NewString()
 	attachment.Size = file.Size
 	attachment.Name = file.Filename
@@ -104,6 +105,13 @@ func TryLinkAttachment(tx *gorm.DB, og models.Attachment, hash string) (bool, er
 		return false, err
 	}
 
+	if prev.PoolID != nil && og.PoolID != nil && prev.PoolID != og.PoolID {
+		if !prev.Pool.Config.Data().AllowCrossPoolEgress || !og.Pool.Config.Data().AllowCrossPoolIngress {
+			// Pool config doesn't allow reference
+			return false, nil
+		}
+	}
+
 	prev.RefCount++
 	og.RefID = &prev.ID
 	og.Uuid = prev.Uuid
diff --git a/pkg/internal/services/pools.go b/pkg/internal/services/pools.go
index e5989b7..ef9a49d 100644
--- a/pkg/internal/services/pools.go
+++ b/pkg/internal/services/pools.go
@@ -21,6 +21,14 @@ func GetAttachmentPool(id uint) (models.AttachmentPool, error) {
 	return pool, nil
 }
 
+func GetAttachmentPoolByAlias(alias string) (models.AttachmentPool, error) {
+	var pool models.AttachmentPool
+	if err := database.C.Where("alias = ?", alias).First(&pool).Error; err != nil {
+		return pool, err
+	}
+	return pool, nil
+}
+
 func GetAttachmentPoolWithUser(id uint, userId uint) (models.AttachmentPool, error) {
 	var pool models.AttachmentPool
 	if err := database.C.Where("id = ? AND account_id = ?", id, userId).First(&pool).Error; err != nil {
diff --git a/pkg/internal/services/random_id.go b/pkg/internal/services/random_id.go
new file mode 100644
index 0000000..b341c30
--- /dev/null
+++ b/pkg/internal/services/random_id.go
@@ -0,0 +1,31 @@
+package services
+
+import (
+	"math/rand"
+	"strings"
+)
+
+const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+const (
+	letterIdxBits = 6
+	letterIdxMask = 1<<letterIdxBits - 1
+	letterIdxMax  = 63 / letterIdxBits
+)
+
+func RandString(length int) string {
+	builder := strings.Builder{}
+	builder.Grow(length)
+	for idx, cache, remain := length-1, rand.Int63(), letterIdxMax; idx >= 0; {
+		if remain == 0 {
+			cache, remain = rand.Int63(), letterIdxMax
+		}
+		if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
+			builder.WriteByte(letterBytes[idx])
+			idx--
+		}
+		cache >>= letterIdxBits
+		remain--
+	}
+
+	return builder.String()
+}
diff --git a/settings.toml b/settings.toml
index 084764e..afdd5fe 100644
--- a/settings.toml
+++ b/settings.toml
@@ -5,11 +5,11 @@ grpc_bind = "0.0.0.0:7443"
 domain = "usercontent.solsynth.dev"
 
 accepts_usage = [
-  "p.avatar",
-  "p.banner",
-  "i.attachment",
-  "m.attachment",
-  "sticker",
+    "p.avatar",
+    "p.banner",
+    "i.attachment",
+    "m.attachment",
+    "sticker",
 ]
 
 [workers]
@@ -33,6 +33,13 @@ refresh_token_duration = 2592000
 dsn = "host=localhost user=postgres password=password dbname=hy_paperclip port=5432 sslmode=disable"
 prefix = "paperclip_"
 
+[pools.aliases]
+"p.avatar" = "avatar"
+"p.banner" = "avatar"
+"i.attachment" = "interactive"
+"m.attachment" = "messaging"
+"sticker" = "sticker"
+
 [destinations.temporary]
 type = "local"
 path = "uploads"