From 297afc5ce699823cca171aa0d44280e8b5e75b26 Mon Sep 17 00:00:00 2001
From: Drew DeVault <sir@cmpwn.com>
Date: Thu, 19 Nov 2020 09:27:59 -0500
Subject: Limit total size of unsent attachments

---
 plugins/base/routes.go |  8 +++++++-
 server.go              |  2 +-
 session.go             | 24 +++++++++++++++---------
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/plugins/base/routes.go b/plugins/base/routes.go
index c8d112e..3f379dd 100644
--- a/plugins/base/routes.go
+++ b/plugins/base/routes.go
@@ -692,7 +692,13 @@ func handleComposeAttachment(ctx *alps.Context) error {
 	var uuids []string
 	for _, fh := range form.File["attachments"] {
 		uuid, err := ctx.Session.PutAttachment(fh, form)
-		if err != nil {
+		if err == alps.ErrAttachmentCacheSize {
+			form.RemoveAll()
+			return ctx.JSON(http.StatusBadRequest, map[string]string{
+				"error": "The total size of unset attachments on your session exceeds the maximum file size. Remove some attachments and try again.",
+			})
+		} else if err != nil {
+			form.RemoveAll()
 			ctx.Logger().Printf("PutAttachment: %v\n", err)
 			return ctx.JSON(http.StatusBadRequest, map[string]string{
 				"error": "failed to store attachment",
diff --git a/server.go b/server.go
index 44f0956..7159697 100644
--- a/server.go
+++ b/server.go
@@ -426,7 +426,7 @@ func New(e *echo.Echo, options *Options) (*Server, error) {
 			}
 
 			ctx.Session, err = ctx.Server.Sessions.get(cookie.Value)
-			if err == errSessionExpired {
+			if err == ErrSessionExpired {
 				ctx.SetSession(nil)
 				return handleUnauthenticated(next, ctx)
 			} else if err != nil {
diff --git a/session.go b/session.go
index 8e69565..950b324 100644
--- a/session.go
+++ b/session.go
@@ -20,6 +20,7 @@ import (
 
 // TODO: make this configurable
 const sessionDuration = 30 * time.Minute
+const maxAttachmentSize = 32 << 20 // 32 MiB
 
 func generateToken() (string, error) {
 	b := make([]byte, 32)
@@ -30,7 +31,10 @@ func generateToken() (string, error) {
 	return base64.URLEncoding.EncodeToString(b), nil
 }
 
-var errSessionExpired = errors.New("session expired")
+var (
+	ErrSessionExpired = errors.New("session expired")
+	ErrAttachmentCacheSize = errors.New("Attachments on session exceed maximum file size")
+)
 
 // AuthError wraps an authentication error.
 type AuthError struct {
@@ -145,15 +149,17 @@ func (s *Session) Close() {
 // Puts an attachment and returns a generated UUID
 func (s *Session) PutAttachment(in *multipart.FileHeader,
 	form *multipart.Form) (string, error) {
-	// TODO: Prevent users from uploading too many attachments, or too large
-	//
-	// Probably just set a cap on the maximum combined size of all files in the
-	// user's session
-	//
-	// TODO: Figure out what to do if the user abandons the compose window
-	// after adding some attachments
 	id := uuid.New()
 	s.attachmentsLocker.Lock()
+
+	var size int64
+	for _, a := range s.attachments {
+		size += a.File.Size
+	}
+	if size + in.Size > maxAttachmentSize {
+		return "", ErrAttachmentCacheSize
+	}
+
 	s.attachments[id.String()] = &Attachment{
 		File:     in,
 		Form:     form,
@@ -241,7 +247,7 @@ func (sm *SessionManager) get(token string) (*Session, error) {
 
 	session, ok := sm.sessions[token]
 	if !ok {
-		return nil, errSessionExpired
+		return nil, ErrSessionExpired
 	}
 	return session, nil
 }
-- 
cgit v1.2.3-59-g8ed1b