ultisuite-backend/internal/api/mail/handlers.go

package mail

import (
	"errors"
	"log/slog"
	"net/http"

	"github.com/go-chi/chi/v5"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/ultisuite/ulti-backend/internal/api/apiresponse"
	"github.com/ultisuite/ulti-backend/internal/api/apivalidate"
	"github.com/ultisuite/ulti-backend/internal/api/middleware"
	"github.com/ultisuite/ulti-backend/internal/api/query"
	"github.com/ultisuite/ulti-backend/internal/mail/credentials"
	"github.com/ultisuite/ulti-backend/internal/mail/limits"
	"github.com/ultisuite/ulti-backend/internal/api/mail/sendguard"
	"github.com/ultisuite/ulti-backend/internal/mail/storage"
	"github.com/ultisuite/ulti-backend/internal/securityaudit"
)

type Handler struct {
	svc         ServiceAPI
	logger      *slog.Logger
	sendLimiter *sendguard.RateLimiter
}

func NewHandlerWithService(svc ServiceAPI) *Handler {
	return &Handler{
		svc:    svc,
		logger: slog.Default().With("component", "mail-api"),
	}
}

func NewHandler(
	db *pgxpool.Pool,
	audit *securityaudit.Logger,
	credentialManager *credentials.Manager,
	objectStorage *storage.Client,
	attachmentsBucket string,
	sendLimiter *sendguard.RateLimiter,
) *Handler {
	h := NewHandlerWithService(NewService(db, audit, credentialManager, objectStorage, attachmentsBucket))
	h.sendLimiter = sendLimiter
	return h
}

func (h *Handler) Routes() chi.Router {
	r := chi.NewRouter()

	r.Get("/accounts", h.ListAccounts)
	r.Post("/accounts", h.CreateAccount)
	r.Get("/accounts/{accountID}", h.GetAccount)
	r.Delete("/accounts/{accountID}", h.DeleteAccount)
	r.Get("/accounts/{accountID}/identities", h.ListIdentities)
	r.Post("/accounts/{accountID}/identities", h.CreateIdentity)

	r.Get("/identities/{identityID}", h.GetIdentity)
	r.Put("/identities/{identityID}", h.UpdateIdentity)
	r.Delete("/identities/{identityID}", h.DeleteIdentity)

	r.Mount("/", h.FolderLabelRoutes())

	r.Get("/search", h.SearchMessages)

	r.Get("/drafts", h.ListDrafts)
	r.Post("/drafts", h.CreateDraft)
	r.Get("/drafts/{draftID}", h.GetDraft)
	r.Put("/drafts/{draftID}", h.UpdateDraft)
	r.Delete("/drafts/{draftID}", h.DeleteDraft)
	r.Post("/drafts/{draftID}/attachments", h.UploadDraftAttachment)
	r.Get("/drafts/{draftID}/attachments/{attachmentID}", h.DownloadDraftAttachment)
	r.Get("/drafts/{draftID}/attachments/{attachmentID}/inline", h.DownloadDraftAttachment)

	r.Get("/messages", h.ListMessages)
	r.Get("/messages/{messageID}/attachments", h.ListMessageAttachments)
	r.Get("/messages/{messageID}/attachments/cid-map", h.MessageAttachmentCIDMap)
	r.Post("/messages/{messageID}/attachments", h.UploadMessageAttachment)
	r.Get("/messages/{messageID}", h.GetMessage)
	r.Put("/messages/{messageID}/labels", h.UpdateLabels)
	r.Put("/messages/{messageID}/flags", h.UpdateFlags)
	r.Delete("/messages/{messageID}", h.DeleteMessage)

	r.Get("/attachments/{attachmentID}", h.DownloadAttachment)
	r.Get("/attachments/{attachmentID}/inline", h.DownloadAttachment)

	r.Get("/threads/{threadID}", h.GetThread)

	r.Post("/send", h.SendMessage)

	r.Get("/rules", h.ListRules)
	r.Post("/rules", h.CreateRule)
	r.Put("/rules/{ruleID}", h.UpdateRule)
	r.Delete("/rules/{ruleID}", h.DeleteRule)

	r.Get("/webhooks", h.ListWebhooks)
	r.Post("/webhooks", h.CreateWebhook)
	r.Delete("/webhooks/{webhookID}", h.DeleteWebhook)

	return r
}

func (h *Handler) ListAccounts(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())
	params, err := query.ParseListRequest(r)
	if err != nil {
		apivalidate.WriteQueryError(w, r, err)
		return
	}

	result, err := h.svc.ListAccounts(r.Context(), claims.Sub, params)
	if err != nil {
		h.logger.Error("list accounts", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusOK, result)
}

func (h *Handler) CreateAccount(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())

	var req createAccountRequest
	if err := apivalidate.DecodeJSON(w, r, maxAccountRequestBody, &req); err != nil {
		return
	}
	if verr := validateCreateAccount(&req); verr != nil {
		apivalidate.WriteValidationError(w, r, verr)
		return
	}

	id, err := h.svc.CreateAccount(r.Context(), claims.Sub, &req)
	if err != nil {
		if errors.Is(err, ErrCredentialsUnavailable) {
			apiresponse.WriteError(w, r, http.StatusInternalServerError, apiresponse.CodeInternal, "credentials encryption unavailable", nil)
			return
		}
		h.logger.Error("create account", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusCreated, map[string]string{"id": id})
}

func (h *Handler) GetAccount(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())
	account, err := h.svc.GetAccount(r.Context(), claims.Sub, chi.URLParam(r, "accountID"))
	if err != nil {
		if errors.Is(err, ErrNotFound) {
			apivalidate.WriteNotFound(w, r, "not found")
			return
		}
		h.logger.Error("get account", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusOK, account)
}

func (h *Handler) DeleteAccount(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())
	if err := h.svc.DeleteAccount(r.Context(), claims.Sub, chi.URLParam(r, "accountID")); err != nil {
		if errors.Is(err, ErrUserNotProvisioned) {
			apiresponse.WriteError(w, r, http.StatusUnauthorized, apiresponse.CodeAuthUnauthorized, "user not provisioned", nil)
			return
		}
		if errors.Is(err, ErrNotFound) {
			apivalidate.WriteNotFound(w, r, "not found")
			return
		}
		h.logger.Error("delete account", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (h *Handler) ListMessages(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())
	params, err := query.ParseListRequest(r)
	if err != nil {
		apivalidate.WriteQueryError(w, r, err)
		return
	}

	filter := MessageListFilter{
		Folder:    r.URL.Query().Get("folder"),
		AccountID: r.URL.Query().Get("account_id"),
	}

	result, err := h.svc.ListMessages(r.Context(), claims.Sub, filter, params)
	if err != nil {
		h.logger.Error("list messages", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusOK, result)
}

func (h *Handler) GetMessage(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())
	msg, err := h.svc.GetMessage(r.Context(), claims.Sub, chi.URLParam(r, "messageID"))
	if err != nil {
		if errors.Is(err, ErrNotFound) {
			apivalidate.WriteNotFound(w, r, "not found")
			return
		}
		h.logger.Error("get message", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusOK, msg)
}

func (h *Handler) UpdateLabels(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())

	var req updateLabelsRequest
	if err := apivalidate.DecodeJSON(w, r, maxFlagsLabelsBody, &req); err != nil {
		return
	}
	if verr := validateUpdateLabels(&req); verr != nil {
		apivalidate.WriteValidationError(w, r, verr)
		return
	}

	if err := h.svc.UpdateLabels(r.Context(), claims.Sub, chi.URLParam(r, "messageID"), req.Labels); err != nil {
		if errors.Is(err, ErrNotFound) {
			apivalidate.WriteNotFound(w, r, "not found")
			return
		}
		h.logger.Error("update labels", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (h *Handler) UpdateFlags(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())

	var req updateFlagsRequest
	if err := apivalidate.DecodeJSON(w, r, maxFlagsLabelsBody, &req); err != nil {
		return
	}
	if verr := validateUpdateFlags(&req); verr != nil {
		apivalidate.WriteValidationError(w, r, verr)
		return
	}

	if err := h.svc.UpdateFlags(r.Context(), claims.Sub, chi.URLParam(r, "messageID"), req.Flags); err != nil {
		if errors.Is(err, ErrNotFound) {
			apivalidate.WriteNotFound(w, r, "not found")
			return
		}
		h.logger.Error("update flags", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (h *Handler) DeleteMessage(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())

	if err := h.svc.DeleteMessage(r.Context(), claims.Sub, chi.URLParam(r, "messageID")); err != nil {
		if errors.Is(err, ErrNotFound) {
			apivalidate.WriteNotFound(w, r, "not found")
			return
		}
		h.logger.Error("delete message", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (h *Handler) GetThread(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())
	result, err := h.svc.GetThread(r.Context(), claims.Sub, chi.URLParam(r, "threadID"))
	if err != nil {
		h.logger.Error("get thread", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusOK, result)
}

func (h *Handler) SendMessage(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())
	userID, err := h.svc.ResolveUserID(r.Context(), claims.Sub)
	if err != nil {
		h.writeUserResolveError(w, r, err)
		return
	}

	if h.sendLimiter != nil {
		if err := h.sendLimiter.Allow(userID); err != nil {
			apiresponse.WriteError(w, r, http.StatusTooManyRequests, apiresponse.CodeRateLimited, "send rate limit exceeded", nil)
			return
		}
	}

	idempotencyKey, ok := normalizeIdempotencyKey(r.Header.Get("Idempotency-Key"))
	if !ok {
		apivalidate.WriteValidationError(w, r, apivalidate.NewValidationError(apivalidate.FieldDetail{
			Field: "Idempotency-Key", Message: "invalid",
		}))
		return
	}

	var req sendMessageRequest
	if err := apivalidate.DecodeJSON(w, r, limits.MaxSendRequestBodyBytes, &req); err != nil {
		return
	}
	req.IdempotencyKey = idempotencyKey
	if verr := validateSendMessage(&req); verr != nil {
		apivalidate.WriteValidationError(w, r, verr)
		return
	}

	id, status, err := h.svc.SendMessage(r.Context(), userID, &req)
	if err != nil {
		if errors.Is(err, ErrAccountNotFound) {
			apivalidate.WriteNotFound(w, r, "account not found")
			return
		}
		h.logger.Error("send message", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusAccepted, map[string]string{"id": id, "status": status})
}

func (h *Handler) ListRules(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())
	params, err := query.ParseListRequest(r)
	if err != nil {
		apivalidate.WriteQueryError(w, r, err)
		return
	}

	result, err := h.svc.ListRules(r.Context(), claims.Sub, params)
	if err != nil {
		h.logger.Error("list rules", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusOK, result)
}

func (h *Handler) CreateRule(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())
	userID, err := h.svc.ResolveUserID(r.Context(), claims.Sub)
	if err != nil {
		h.writeUserResolveError(w, r, err)
		return
	}

	var req createRuleRequest
	if err := apivalidate.DecodeJSON(w, r, maxRulesRequestBody, &req); err != nil {
		return
	}
	if verr := validateCreateRule(&req); verr != nil {
		apivalidate.WriteValidationError(w, r, verr)
		return
	}

	id, err := h.svc.CreateRule(r.Context(), userID, &req)
	if err != nil {
		if errors.Is(err, ErrAccountNotFound) {
			apivalidate.WriteNotFound(w, r, "account not found")
			return
		}
		h.logger.Error("create rule", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusCreated, map[string]string{"id": id})
}

func (h *Handler) UpdateRule(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())

	var req updateRuleRequest
	if err := apivalidate.DecodeJSON(w, r, maxRulesRequestBody, &req); err != nil {
		return
	}
	if verr := validateUpdateRule(&req); verr != nil {
		apivalidate.WriteValidationError(w, r, verr)
		return
	}

	if err := h.svc.UpdateRule(r.Context(), claims.Sub, chi.URLParam(r, "ruleID"), &req); err != nil {
		if errors.Is(err, ErrNotFound) {
			apivalidate.WriteNotFound(w, r, "not found")
			return
		}
		h.logger.Error("update rule", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (h *Handler) DeleteRule(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())

	if err := h.svc.DeleteRule(r.Context(), claims.Sub, chi.URLParam(r, "ruleID")); err != nil {
		if errors.Is(err, ErrNotFound) {
			apivalidate.WriteNotFound(w, r, "not found")
			return
		}
		h.logger.Error("delete rule", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (h *Handler) ListWebhooks(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())
	params, err := query.ParseListRequest(r)
	if err != nil {
		apivalidate.WriteQueryError(w, r, err)
		return
	}

	result, err := h.svc.ListWebhooks(r.Context(), claims.Sub, params)
	if err != nil {
		h.logger.Error("list webhooks", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusOK, result)
}

func (h *Handler) CreateWebhook(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())

	var req createWebhookRequest
	if err := apivalidate.DecodeJSON(w, r, maxWebhookRequestBody, &req); err != nil {
		return
	}
	method, verr := validateCreateWebhook(&req)
	if verr != nil {
		apivalidate.WriteValidationError(w, r, verr)
		return
	}

	id, err := h.svc.CreateWebhook(r.Context(), claims.Sub, &req, method)
	if err != nil {
		h.logger.Error("create webhook", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	apiresponse.WriteJSON(w, http.StatusCreated, map[string]string{"id": id})
}

func (h *Handler) DeleteWebhook(w http.ResponseWriter, r *http.Request) {
	claims := middleware.ClaimsFromContext(r.Context())

	if err := h.svc.DeleteWebhook(r.Context(), claims.Sub, chi.URLParam(r, "webhookID")); err != nil {
		if errors.Is(err, ErrNotFound) {
			apivalidate.WriteNotFound(w, r, "not found")
			return
		}
		h.logger.Error("delete webhook", "error", err)
		apivalidate.WriteInternal(w, r)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (h *Handler) writeUserResolveError(w http.ResponseWriter, r *http.Request, err error) {
	if errors.Is(err, ErrUserNotProvisioned) {
		apiresponse.WriteError(w, r, http.StatusUnauthorized, apiresponse.CodeAuthUnauthorized, "user not provisioned", nil)
		return
	}
	h.logger.Error("resolve user id", "error", err)
	apivalidate.WriteInternal(w, r)
}