Philipp Heckel 4 лет назад
Родитель
Сommit
5eca20469f
5 измененных файлов с 76 добавлено и 7 удалено
  1. 13 0
      cmd/serve.go
  2. 1 1
      server/config.go
  3. 4 4
      server/message.go
  4. 27 2
      util/util.go
  5. 31 0
      util/util_test.go

+ 13 - 0
cmd/serve.go

@@ -21,6 +21,7 @@ var flagsServe = []cli.Flag{
 	altsrc.NewStringFlag(&cli.StringFlag{Name: "cache-file", Aliases: []string{"C"}, EnvVars: []string{"NTFY_CACHE_FILE"}, Usage: "cache file used for message caching"}),
 	altsrc.NewDurationFlag(&cli.DurationFlag{Name: "cache-duration", Aliases: []string{"b"}, EnvVars: []string{"NTFY_CACHE_DURATION"}, Value: server.DefaultCacheDuration, Usage: "buffer messages for this time to allow `since` requests"}),
 	altsrc.NewStringFlag(&cli.StringFlag{Name: "attachment-cache-dir", EnvVars: []string{"NTFY_ATTACHMENT_CACHE_DIR"}, Usage: "cache directory for attached files"}),
+	altsrc.NewStringFlag(&cli.StringFlag{Name: "attachment-size-limit", Aliases: []string{"A"}, EnvVars: []string{"NTFY_ATTACHMENT_SIZE_LIMIT"}, DefaultText: "15M", Usage: "attachment size limit (e.g. 10k, 2M)"}),
 	altsrc.NewDurationFlag(&cli.DurationFlag{Name: "keepalive-interval", Aliases: []string{"k"}, EnvVars: []string{"NTFY_KEEPALIVE_INTERVAL"}, Value: server.DefaultKeepaliveInterval, Usage: "interval of keepalive messages"}),
 	altsrc.NewDurationFlag(&cli.DurationFlag{Name: "manager-interval", Aliases: []string{"m"}, EnvVars: []string{"NTFY_MANAGER_INTERVAL"}, Value: server.DefaultManagerInterval, Usage: "interval of for message pruning and stats printing"}),
 	altsrc.NewStringFlag(&cli.StringFlag{Name: "smtp-sender-addr", EnvVars: []string{"NTFY_SMTP_SENDER_ADDR"}, Usage: "SMTP server address (host:port) for outgoing emails"}),
@@ -71,6 +72,7 @@ func execServe(c *cli.Context) error {
 	cacheFile := c.String("cache-file")
 	cacheDuration := c.Duration("cache-duration")
 	attachmentCacheDir := c.String("attachment-cache-dir")
+	attachmentSizeLimitStr := c.String("attachment-size-limit")
 	keepaliveInterval := c.Duration("keepalive-interval")
 	managerInterval := c.Duration("manager-interval")
 	smtpSenderAddr := c.String("smtp-sender-addr")
@@ -109,6 +111,16 @@ func execServe(c *cli.Context) error {
 		return errors.New("if smtp-server-listen is set, smtp-server-domain must also be set")
 	}
 
+	// Convert
+	attachmentSizeLimit := server.DefaultAttachmentSizeLimit
+	if attachmentSizeLimitStr != "" {
+		var err error
+		attachmentSizeLimit, err = util.ParseSize(attachmentSizeLimitStr)
+		if err != nil {
+			return err
+		}
+	}
+
 	// Run server
 	conf := server.NewConfig()
 	conf.BaseURL = baseURL
@@ -120,6 +132,7 @@ func execServe(c *cli.Context) error {
 	conf.CacheFile = cacheFile
 	conf.CacheDuration = cacheDuration
 	conf.AttachmentCacheDir = attachmentCacheDir
+	conf.AttachmentSizeLimit = attachmentSizeLimit
 	conf.KeepaliveInterval = keepaliveInterval
 	conf.ManagerInterval = managerInterval
 	conf.SMTPSenderAddr = smtpSenderAddr

+ 1 - 1
server/config.go

@@ -14,7 +14,7 @@ const (
 	DefaultMinDelay                  = 10 * time.Second
 	DefaultMaxDelay                  = 3 * 24 * time.Hour
 	DefaultMessageLimit              = 4096 // Bytes
-	DefaultAttachmentSizeLimit       = 15 * 1024 * 1024
+	DefaultAttachmentSizeLimit       = int64(15 * 1024 * 1024)
 	DefaultAttachmentSizePreviewMax  = 20 * 1024 * 1024 // Bytes
 	DefaultAttachmentExpiryDuration  = 3 * time.Hour
 	DefaultFirebaseKeepaliveInterval = 3 * time.Hour // Not too frequently to save battery

+ 4 - 4
server/message.go

@@ -32,10 +32,10 @@ type message struct {
 
 type attachment struct {
 	Name       string `json:"name"`
-	Type       string `json:"type"`
-	Size       int64  `json:"size"`
-	Expires    int64  `json:"expires"`
-	PreviewURL string `json:"preview_url"`
+	Type       string `json:"type,omitempty"`
+	Size       int64  `json:"size,omitempty"`
+	Expires    int64  `json:"expires,omitempty"`
+	PreviewURL string `json:"preview_url,omitempty"`
 	URL        string `json:"url"`
 }
 

+ 27 - 2
util/util.go

@@ -6,6 +6,8 @@ import (
 	"math/rand"
 	"mime"
 	"os"
+	"regexp"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -16,8 +18,9 @@ const (
 )
 
 var (
-	random      = rand.New(rand.NewSource(time.Now().UnixNano()))
-	randomMutex = sync.Mutex{}
+	random       = rand.New(rand.NewSource(time.Now().UnixNano()))
+	randomMutex  = sync.Mutex{}
+	sizeStrRegex = regexp.MustCompile(`(?i)^(\d+)([gmkb])?$`)
 
 	errInvalidPriority = errors.New("invalid priority")
 )
@@ -178,3 +181,25 @@ func ExtensionByType(contentType string) string {
 		return ".bin"
 	}
 }
+
+// ParseSize parses a size string like 2K or 2M into bytes. If no unit is found, e.g. 123, bytes is assumed.
+func ParseSize(s string) (int64, error) {
+	matches := sizeStrRegex.FindStringSubmatch(s)
+	if matches == nil {
+		return -1, fmt.Errorf("invalid size %s", s)
+	}
+	value, err := strconv.Atoi(matches[1])
+	if err != nil {
+		return -1, fmt.Errorf("cannot convert number %s", matches[1])
+	}
+	switch strings.ToUpper(matches[2]) {
+	case "G":
+		return int64(value) * 1024 * 1024 * 1024, nil
+	case "M":
+		return int64(value) * 1024 * 1024, nil
+	case "K":
+		return int64(value) * 1024, nil
+	default:
+		return int64(value), nil
+	}
+}

+ 31 - 0
util/util_test.go

@@ -121,3 +121,34 @@ func TestShortTopicURL(t *testing.T) {
 	require.Equal(t, "ntfy.sh/mytopic", ShortTopicURL("http://ntfy.sh/mytopic"))
 	require.Equal(t, "lalala", ShortTopicURL("lalala"))
 }
+
+func TestParseSize_10GSuccess(t *testing.T) {
+	s, err := ParseSize("10G")
+	if err != nil {
+		t.Fatal(err)
+	}
+	require.Equal(t, 10*1024*1024*1024, s)
+}
+
+func TestParseSize_10MUpperCaseSuccess(t *testing.T) {
+	s, err := ParseSize("10M")
+	if err != nil {
+		t.Fatal(err)
+	}
+	require.Equal(t, 10*1024*1024, s)
+}
+
+func TestParseSize_10kLowerCaseSuccess(t *testing.T) {
+	s, err := ParseSize("10k")
+	if err != nil {
+		t.Fatal(err)
+	}
+	require.Equal(t, 10*1024, s)
+}
+
+func TestParseSize_FailureInvalid(t *testing.T) {
+	_, err := ParseSize("not a size")
+	if err == nil {
+		t.Fatalf("expected error, but got none")
+	}
+}