120 lines
2.2 KiB
Go
120 lines
2.2 KiB
Go
package store
|
|
|
|
import (
|
|
"log"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
// StoredImage represents an image with its data and expiration time
|
|
type Image struct {
|
|
Data []byte
|
|
ContentType string
|
|
}
|
|
|
|
// ImageStore holds uploaded images in memory
|
|
type Store struct {
|
|
mu sync.RWMutex
|
|
images map[string]*Image
|
|
}
|
|
|
|
// NewStore creates a new image store
|
|
func newStore() *Store {
|
|
s := &Store{
|
|
images: make(map[string]*Image),
|
|
}
|
|
// Start cleanup goroutine
|
|
return s
|
|
}
|
|
|
|
var store *Store
|
|
|
|
// Get retrieves an image by ID
|
|
func (s *Store) Get(id string) (*Image, bool) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
image, exists := s.images[id]
|
|
if !exists {
|
|
return nil, false
|
|
} else {
|
|
|
|
// drop this guy 30 seconds after the first access
|
|
go func(id string) {
|
|
log.Println("delete", id, "in ~30s (after access)")
|
|
time.Sleep(time.Second * 30)
|
|
|
|
log.Println("lock map...")
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
log.Println("delete", id)
|
|
delete(s.images, id)
|
|
}(id)
|
|
}
|
|
|
|
return image, true
|
|
}
|
|
|
|
// isImageContentType checks if the content type is a valid image
|
|
func extFromContentType(contentType string) (string, bool) {
|
|
val, ok := map[string]string{
|
|
"image/jpeg": "jpg",
|
|
"image/jpg": "jpg",
|
|
"image/png": "png",
|
|
"image/gif": "gif",
|
|
"image/webp": "webp",
|
|
"image/bmp": "bmp",
|
|
}[contentType]
|
|
return val, ok
|
|
}
|
|
|
|
// Save stores an image and returns its ID
|
|
func (s *Store) Save(data []byte, contentType string) string {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
ext, ok := extFromContentType(contentType)
|
|
if !ok {
|
|
ext = ""
|
|
} else {
|
|
ext = "." + ext
|
|
}
|
|
|
|
id := uuid.New().String() + ext
|
|
s.images[id] = &Image{
|
|
Data: data,
|
|
ContentType: contentType,
|
|
}
|
|
|
|
// start a goroutine to drop this guy in 1 hour
|
|
go func(id string) {
|
|
log.Println("delete", id, "in ~1h (after creation)")
|
|
time.Sleep(time.Hour * 1)
|
|
|
|
log.Println("lock map...")
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
log.Println("delete", id)
|
|
delete(s.images, id)
|
|
}(id)
|
|
|
|
return id
|
|
}
|
|
|
|
func Get(id string) (*Image, bool) {
|
|
if store == nil {
|
|
store = newStore()
|
|
}
|
|
|
|
return store.Get(id)
|
|
}
|
|
|
|
func Save(data []byte, contentType string) string {
|
|
if store == nil {
|
|
store = newStore()
|
|
}
|
|
return store.Save(data, contentType)
|
|
}
|