Skip to content

Commit 315dabe

Browse files
bemotybemoty
committed
feat: added server side ttl feature
1 parent 764ce33 commit 315dabe

File tree

8 files changed

+281
-46
lines changed

8 files changed

+281
-46
lines changed

‎cmd/client/cmd/root.go‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -180,9 +180,6 @@ func initConfig() {
180180
if err := viper.BindPFlag("key", RootCmd.PersistentFlags().Lookup("key")); err != nil {
181181
panic(err)
182182
}
183-
if err := viper.BindPFlag("ttl", RootCmd.Flags().Lookup("ttl")); err != nil {
184-
panic(err)
185-
}
186183

187184
if err := viper.ReadInConfig(); err != nil {
188185
var notFound viper.ConfigFileNotFoundError

‎cmd/server/config.go‎

Lines changed: 32 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,30 +5,49 @@ import (
55
"os"
66
"strconv"
77
"strings"
8+
"time"
89
)
910

1011
type Config struct {
11-
Port string
12-
StoragePath string
13-
AuthKey string
14-
BaseURL string
15-
MaxFileMB int64
16-
PasteStyle string
12+
Port string
13+
StoragePath string
14+
AuthKey string
15+
BaseURL string
16+
MaxFileMB int64
17+
PasteStyle string
18+
DefaultTTL time.Duration
19+
SweepInterval time.Duration
1720
}
1821

19-
func LoadConfig() (Config, error) {
22+
func NewConfig() (Config, error) {
2023
maxFileMBStr := getEnv("MAX_FILE_MB", "100")
2124
maxFileMB, err := strconv.ParseInt(maxFileMBStr, 10, 64)
2225
if err != nil || maxFileMB <= 0 {
2326
return Config{}, fmt.Errorf("MAX_FILE_MB must be a positive integer, got %q", maxFileMBStr)
2427
}
28+
29+
sweepInterval, err := time.ParseDuration(getEnv("SWEEP_INTERVAL", "1h"))
30+
if err != nil {
31+
return Config{}, fmt.Errorf("SWEEP_INTERVAL must be a valid duration, got %q", getEnv("SWEEP_INTERVAL", "1h"))
32+
}
33+
34+
var defaultTTL time.Duration
35+
if s := getEnv("DEFAULT_TTL", ""); s != "" {
36+
defaultTTL, err = parseTTL(s)
37+
if err != nil {
38+
return Config{}, fmt.Errorf("DEFAULT_TTL is invalid: %w", err)
39+
}
40+
}
41+
2542
return Config{
26-
Port: normalizeAddress(getEnv("PORT", ":8080")),
27-
StoragePath: getEnv("STORAGE_PATH", "./data"),
28-
AuthKey: getEnv("AUTH_KEY", "no-auth"),
29-
BaseURL: getEnv("BASE_URL", "https://i.bemoty.dev"),
30-
MaxFileMB: maxFileMB,
31-
PasteStyle: getEnv("PASTE_STYLE", "dracula"),
43+
Port: normalizeAddress(getEnv("PORT", ":8080")),
44+
StoragePath: getEnv("STORAGE_PATH", "./data"),
45+
AuthKey: getEnv("AUTH_KEY", "no-auth"),
46+
BaseURL: getEnv("BASE_URL", "https://i.bemoty.dev"),
47+
MaxFileMB: maxFileMB,
48+
PasteStyle: getEnv("PASTE_STYLE", "dracula"),
49+
DefaultTTL: defaultTTL,
50+
SweepInterval: sweepInterval,
3251
}, nil
3352
}
3453

‎cmd/server/handlers.go‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,17 @@ import (
44
"bytes"
55
"crypto/subtle"
66
"errors"
7+
"fmt"
78
"io"
89
"log/slog"
910
"mime"
1011
"net/http"
1112
"os"
1213
"path/filepath"
1314
"regexp"
15+
"strconv"
1416
"strings"
17+
"time"
1518

1619
"github.com/alecthomas/chroma"
1720
"github.com/alecthomas/chroma/formatters/html"
@@ -26,6 +29,8 @@ type Server struct {
2629

2730
var langRegex = regexp.MustCompile(`^[a-z0-9]{1,20}$`)
2831

32+
const maxTTL = 365 * 24 * time.Hour
33+
2934
func (s *Server) HandleUpload(w http.ResponseWriter, r *http.Request) {
3035
authHeader := r.Header.Get("Authorization")
3136
expectedAuthHeader := "Bearer " + s.config.AuthKey
@@ -34,6 +39,16 @@ func (s *Server) HandleUpload(w http.ResponseWriter, r *http.Request) {
3439
return
3540
}
3641

42+
ttl := s.config.DefaultTTL
43+
if ttlStr := r.URL.Query().Get("ttl"); ttlStr != "" {
44+
var err error
45+
ttl, err = parseTTL(ttlStr)
46+
if err != nil {
47+
http.Error(w, "Invalid TTL", http.StatusBadRequest)
48+
return
49+
}
50+
}
51+
3752
r.Body = http.MaxBytesReader(w, r.Body, s.config.MaxFileMB<<20)
3853
defer func(Body io.ReadCloser) {
3954
err := Body.Close()
@@ -64,6 +79,12 @@ func (s *Server) HandleUpload(w http.ResponseWriter, r *http.Request) {
6479
return
6580
}
6681

82+
if ttl > 0 {
83+
if err := s.store.WriteMeta(id, ttl); err != nil {
84+
slog.Warn("failed to write ttl metadata", "id", id, "error", err)
85+
}
86+
}
87+
6788
fullURL := s.config.BaseURL + "/" + id
6889
if _, err := w.Write([]byte(fullURL)); err != nil {
6990
slog.Warn("failed to write response to client", "error", err)
@@ -99,6 +120,23 @@ func (s *Server) HandleServe(w http.ResponseWriter, r *http.Request) {
99120
http.NotFound(w, r)
100121
return
101122
}
123+
124+
expiry, metaErr := s.store.GetMeta(id)
125+
if metaErr != nil && !errors.Is(metaErr, os.ErrNotExist) {
126+
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
127+
return
128+
}
129+
if metaErr == nil && time.Now().After(expiry) {
130+
if err := s.store.DeleteFile(id); err != nil {
131+
slog.Warn("failed to delete expired file", "id", id, "error", err)
132+
}
133+
if err := s.store.DeleteMeta(id); err != nil {
134+
slog.Warn("failed to delete expired file metadata", "id", id, "error", err)
135+
}
136+
http.Error(w, "Gone", http.StatusGone)
137+
return
138+
}
139+
102140
w.Header().Set("X-Content-Type-Options", "nosniff")
103141

104142
fileType := mime.TypeByExtension(filepath.Ext(path))
@@ -119,6 +157,31 @@ func (s *Server) HandleServe(w http.ResponseWriter, r *http.Request) {
119157
}
120158
}
121159

160+
func parseTTL(s string) (time.Duration, error) {
161+
if strings.HasSuffix(s, "d") {
162+
n, err := strconv.Atoi(strings.TrimSuffix(s, "d"))
163+
if err != nil {
164+
return 0, fmt.Errorf("invalid ttl %q", s)
165+
}
166+
return validateTTL(time.Duration(n) * 24 * time.Hour)
167+
}
168+
d, err := time.ParseDuration(s)
169+
if err != nil {
170+
return 0, err
171+
}
172+
return validateTTL(d)
173+
}
174+
175+
func validateTTL(d time.Duration) (time.Duration, error) {
176+
if d > maxTTL {
177+
return 0, errors.New("ttl too large")
178+
}
179+
if d <= 0 {
180+
return 0, errors.New("ttl cannot be negative or zero")
181+
}
182+
return d, nil
183+
}
184+
122185
func renderText(w http.ResponseWriter, path string, content []byte, style string) error {
123186
lexer := lexers.Match(path)
124187
if lexer == nil {

‎cmd/server/main.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package main
22

33
import (
4+
"context"
45
"log/slog"
56
"net/http"
67
"os"
@@ -10,12 +11,14 @@ func main() {
1011
handler := slog.NewTextHandler(os.Stdout, nil)
1112
slog.SetDefault(slog.New(handler))
1213

13-
config, err := LoadConfig()
14+
config, err := NewConfig()
1415
if err != nil {
1516
slog.Error("invalid config", "error", err)
1617
os.Exit(1)
1718
}
18-
store := &DiskStore{config.StoragePath}
19+
ctx, cancel := context.WithCancel(context.Background())
20+
defer cancel()
21+
store := NewDiskStore(ctx, config)
1922
server := Server{config, store}
2023

2124
mux := http.NewServeMux()

‎cmd/server/storage.go‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,17 @@
11
package main
22

33
import (
4+
"context"
45
"crypto/rand"
56
"encoding/base64"
7+
"errors"
68
"io"
9+
"io/fs"
10+
"log/slog"
711
"os"
812
"path/filepath"
13+
"strings"
14+
"time"
915
)
1016

1117
type DiskStore struct {
@@ -15,6 +21,70 @@ type DiskStore struct {
1521
// IdByteLength must not be less than 3, or the sharding logic will panic (minimum 5 for sensible file names)
1622
const IdByteLength = 6
1723

24+
func NewDiskStore(ctx context.Context, config Config) *DiskStore {
25+
store := &DiskStore{BaseDir: config.StoragePath}
26+
go store.sweep(ctx, config.SweepInterval)
27+
return store
28+
}
29+
30+
func (s *DiskStore) sweep(ctx context.Context, sweepInterval time.Duration) {
31+
ticker := time.NewTicker(sweepInterval)
32+
defer ticker.Stop()
33+
for {
34+
select {
35+
case <-ctx.Done():
36+
return
37+
case <-ticker.C:
38+
s.performCleanup()
39+
}
40+
}
41+
}
42+
43+
func (s *DiskStore) performCleanup() {
44+
metaDir := filepath.Join(s.BaseDir, ".meta")
45+
err := filepath.WalkDir(metaDir, func(path string, d fs.DirEntry, err error) error {
46+
if err != nil {
47+
slog.Warn("sweep: walk error", "path", path, "error", err)
48+
return nil
49+
}
50+
if d.IsDir() {
51+
return nil
52+
}
53+
data, err := os.ReadFile(path)
54+
if err != nil {
55+
slog.Warn("sweep: failed to read meta file", "path", path, "error", err)
56+
return nil
57+
}
58+
expiry, err := time.Parse(time.RFC3339, string(data))
59+
if err != nil {
60+
slog.Warn("sweep: failed to parse meta file", "path", path, "error", err)
61+
return nil
62+
}
63+
if !time.Now().After(expiry) {
64+
return nil
65+
}
66+
rel, err := filepath.Rel(metaDir, path)
67+
if err != nil {
68+
return nil
69+
}
70+
parts := strings.Split(rel, string(filepath.Separator))
71+
if len(parts) != 3 {
72+
return nil
73+
}
74+
id := parts[0] + parts[1] + parts[2]
75+
if err := s.DeleteFile(id); err != nil && !errors.Is(err, os.ErrNotExist) {
76+
slog.Warn("sweep: failed to delete expired file", "id", id, "error", err)
77+
}
78+
if err := s.DeleteMeta(id); err != nil && !errors.Is(err, os.ErrNotExist) {
79+
slog.Warn("sweep: failed to delete expired meta", "id", id, "error", err)
80+
}
81+
return nil
82+
})
83+
if err != nil && !errors.Is(err, os.ErrNotExist) {
84+
slog.Warn("sweep failed", "error", err)
85+
}
86+
}
87+
1888
func (s *DiskStore) SaveFile(r io.Reader, ext string) (string, error) {
1989
for {
2090
id, err := generateId(IdByteLength)
@@ -68,6 +138,51 @@ func (s *DiskStore) GetFile(id string) (string, bool) {
68138
return matches[0], true
69139
}
70140

141+
func (s *DiskStore) DeleteFile(id string) error {
142+
path, ok := s.GetFile(id)
143+
if !ok {
144+
return os.ErrNotExist
145+
}
146+
if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) {
147+
return err
148+
}
149+
dir := filepath.Dir(path)
150+
_ = os.Remove(dir)
151+
_ = os.Remove(filepath.Dir(dir))
152+
return nil
153+
}
154+
155+
func (s *DiskStore) GetMeta(id string) (time.Time, error) {
156+
contents, err := os.ReadFile(s.metaPath(id))
157+
if err != nil {
158+
return time.Time{}, err
159+
}
160+
return time.Parse(time.RFC3339, string(contents))
161+
}
162+
163+
func (s *DiskStore) WriteMeta(id string, expiry time.Duration) error {
164+
path := s.metaPath(id)
165+
if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
166+
return err
167+
}
168+
return os.WriteFile(path, []byte(time.Now().Add(expiry).Format(time.RFC3339)), 0644)
169+
}
170+
171+
func (s *DiskStore) DeleteMeta(id string) error {
172+
path := s.metaPath(id)
173+
if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) {
174+
return err
175+
}
176+
dir := filepath.Dir(path)
177+
_ = os.Remove(dir)
178+
_ = os.Remove(filepath.Dir(dir))
179+
return nil
180+
}
181+
182+
func (s *DiskStore) metaPath(id string) string {
183+
return filepath.Join(s.BaseDir, ".meta", id[:2], id[2:4], id[4:])
184+
}
185+
71186
func generateId(length int) (string, error) {
72187
buffer := make([]byte, length)
73188
if _, err := rand.Read(buffer); err != nil {

0 commit comments

Comments
 (0)