package controller import ( "context" "fmt" "io" "os" "path/filepath" "strings" "sync" "time" gonanoid "github.com/matoous/go-nanoid/v2" "github.com/loveuer/nf/nft/log" "github.com/loveuer/ushare/internal/model" "github.com/loveuer/ushare/internal/opt" "github.com/pkg/errors" "github.com/spf13/viper" ) type metaInfo struct { f *os.File name string create time.Time last time.Time size int64 cursor int64 user string maxDownloads int expiresAt int64 } func (m *metaInfo) generateMeta(code string) error { content := fmt.Sprintf( "filename=%s\ncreated_at=%d\nsize=%d\nuploader=%s\nmax_downloads=%d\nexpires_at=%d\ndownloads=0", m.name, m.create.UnixMilli(), m.size, m.user, m.maxDownloads, m.expiresAt, ) return os.WriteFile(opt.MetaPath(code), []byte(content), 0644) } type meta struct { sync.Mutex ctx context.Context m map[string]*metaInfo } var ( MetaManager = &meta{m: make(map[string]*metaInfo)} ) const letters = "1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ" // New creates a new upload session. // maxDownloads: 0 = unlimited; expiresIn: seconds from now (minimum opt.MinExpiresIn). func (m *meta) New(size int64, filename, ip string, maxDownloads int, expiresIn int64) (string, error) { now := time.Now() if expiresIn < opt.MinExpiresIn { expiresIn = opt.DefaultExpiresIn } if maxDownloads < 0 { maxDownloads = 0 } code, err := gonanoid.Generate(letters, opt.CodeLength) if err != nil { return "", err } f, err := os.Create(opt.FilePath(code)) if err != nil { return "", err } if err = f.Truncate(size); err != nil { f.Close() return "", err } m.Lock() defer m.Unlock() m.m[code] = &metaInfo{ f: f, name: filename, last: now, size: size, cursor: 0, create: now, user: ip, maxDownloads: maxDownloads, expiresAt: now.Unix() + expiresIn, } return code, nil } func (m *meta) Write(code string, start, end int64, reader io.Reader) (total, cursor int64, err error) { m.Lock() defer m.Unlock() if _, ok := m.m[code]; !ok { return 0, 0, fmt.Errorf("code not exist") } w, err := io.CopyN(m.m[code].f, reader, end-start+1) if err != nil { return 0, 0, err } m.m[code].cursor += w m.m[code].last = time.Now() total = m.m[code].size cursor = m.m[code].cursor if m.m[code].cursor == m.m[code].size { defer delete(m.m, code) if err = m.m[code].generateMeta(code); err != nil { return 0, 0, err } } return total, cursor, nil } // CheckAndIncrDownload reads the meta file, validates expiry and download limit, // increments the download counter, and writes the meta file back. // Returns the meta on success, or an error if the file is unavailable. func (m *meta) CheckAndIncrDownload(code string) (*model.Meta, error) { m.Lock() defer m.Unlock() metaPath := opt.MetaPath(code) v := viper.New() v.SetConfigFile(metaPath) v.SetConfigType("env") if err := v.ReadInConfig(); err != nil { return nil, errors.New("文件不存在或已过期") } info := new(model.Meta) if err := v.Unmarshal(info); err != nil { return nil, errors.New("文件元数据损坏") } now := time.Now().Unix() // Check expiry if info.ExpiresAt > 0 && now > info.ExpiresAt { // Clean up expired files go func() { _ = os.RemoveAll(opt.FilePath(code)) _ = os.RemoveAll(metaPath) }() return nil, errors.New("文件已过期") } // Check download limit if info.MaxDownloads > 0 && info.Downloads >= info.MaxDownloads { return nil, errors.New("文件下载次数已达上限") } // Increment downloads and write back info.Downloads++ content := fmt.Sprintf( "filename=%s\ncreated_at=%d\nsize=%d\nuploader=%s\nmax_downloads=%d\nexpires_at=%d\ndownloads=%d", info.Filename, info.CreatedAt, info.Size, info.Uploader, info.MaxDownloads, info.ExpiresAt, info.Downloads, ) if err := os.WriteFile(metaPath, []byte(content), 0644); err != nil { log.Warn("meta.CheckAndIncrDownload: write back failed: %s", err.Error()) } // If this was the last allowed download, clean up after serving if info.MaxDownloads > 0 && info.Downloads >= info.MaxDownloads { go func() { time.Sleep(5 * time.Second) _ = os.RemoveAll(opt.FilePath(code)) _ = os.RemoveAll(metaPath) }() } return info, nil } func (m *meta) Start(ctx context.Context) { ticker := time.NewTicker(time.Minute) m.ctx = ctx if err := os.MkdirAll(opt.Cfg.DataPath, 0644); err != nil { log.Fatal("controller.MetaManager.Start: mkdir datapath failed, path = %s, err = %s", opt.Cfg.DataPath, err.Error()) } // Clean uploads with no activity for 2 minutes go func() { for { select { case <-ctx.Done(): return case now := <-ticker.C: for code, info := range m.m { if now.Sub(info.last) > 2*time.Minute { m.Lock() if err := info.f.Close(); err != nil { log.Warn("handler.Meta: [timer] close file failed, file = %s, err = %s", opt.FilePath(code), err.Error()) } if err := os.RemoveAll(opt.FilePath(code)); err != nil { log.Warn("handler.Meta: [timer] remove file failed, file = %s, err = %s", opt.FilePath(code), err.Error()) } delete(m.m, code) m.Unlock() log.Warn("MetaController: code timeout removed, code = %s", code) } } } } }() // Clean expired files by walking the data directory go func() { if opt.Cfg.CleanInterval <= 0 { log.Warn("meta.Clean: no clean interval set, plz clean manual!!!") return } ticker := time.NewTicker(5 * time.Minute) duration := time.Duration(opt.Cfg.CleanInterval) * time.Hour for { select { case <-ctx.Done(): return case now := <-ticker.C: _ = filepath.Walk(opt.Cfg.DataPath, func(path string, info os.FileInfo, err error) error { if info == nil { return nil } if info.IsDir() { return nil } name := filepath.Base(info.Name()) if !strings.HasPrefix(name, ".meta.") { return nil } v := viper.New() v.SetConfigFile(path) v.SetConfigType("env") if err = v.ReadInConfig(); err != nil { return nil } mi := new(model.Meta) if err = v.Unmarshal(mi); err != nil { return nil } code := strings.TrimPrefix(name, ".meta.") // Remove if past explicit expiry if mi.ExpiresAt > 0 && now.Unix() > mi.ExpiresAt { log.Debug("controller.meta: file expired, code = %s", code) _ = os.RemoveAll(opt.FilePath(code)) _ = os.RemoveAll(path) return nil } // Remove if past global clean interval if now.Sub(time.UnixMilli(mi.CreatedAt)) > duration { log.Debug("controller.meta: file out of date, code = %s", code) _ = os.RemoveAll(opt.FilePath(code)) _ = os.RemoveAll(path) } return nil }) } } }() }