latexgrambot: inline LaTeX bot for Telegram in Go
Renders LaTeX expressions into PNG pictures, PDF documents and native Telegram rich text messages, entirely locally with TeX Live and poppler. - Chat mode with /image, /pdf, /rich plus per-user preamble and DPI - Inline mode delivering picture, PDF and rich text results - Broad default package set preloaded into a TeX format at image build - MongoDB persistence: settings, users, groups, render statistics - Private S3-compatible bucket with short-lived presigned URLs - OpenTelemetry traces over OTLP - Kubernetes manifests and Argo CD application, Docker Compose, GitHub Actions CI, BSD 3-Clause license
This commit is contained in:
corrado.mulas
882
internal/bot/bot.go
Normal file
882
internal/bot/bot.go
Normal file
@@ -0,0 +1,882 @@
|
||||
// Package bot implements the Telegram-facing behavior: messages render LaTeX
|
||||
// into a PNG picture, a PDF document or a Telegram rich message, inline
|
||||
// queries offer the same three formats, and commands manage per-user
|
||||
// settings.
|
||||
package bot
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"latexgrambot/internal/config"
|
||||
"latexgrambot/internal/render"
|
||||
"latexgrambot/internal/rich"
|
||||
"latexgrambot/internal/store"
|
||||
|
||||
tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
|
||||
"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
|
||||
"go.opentelemetry.io/otel"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
const (
|
||||
pollTimeout = 30 // seconds, Telegram long-poll
|
||||
pollStale = 3 * time.Minute
|
||||
minDPI = 72
|
||||
maxDPI = 1200
|
||||
displayLimit = 3800 // long texts are truncated before being echoed back
|
||||
|
||||
// Included in content hashes so renderer/template changes never reuse
|
||||
// stale Telegram file_ids or S3 objects from an older deployment.
|
||||
renderCacheVersion = "presigned-v1"
|
||||
)
|
||||
|
||||
// outputFormat is one of the three generation formats.
|
||||
type outputFormat int
|
||||
|
||||
const (
|
||||
formatImage outputFormat = iota
|
||||
formatPDF
|
||||
formatRich
|
||||
)
|
||||
|
||||
// chatAction returns the Telegram chat action shown while the format is
|
||||
// being produced.
|
||||
func (f outputFormat) chatAction() string {
|
||||
switch f {
|
||||
case formatPDF:
|
||||
return tgbotapi.ChatUploadDocument
|
||||
case formatRich:
|
||||
return tgbotapi.ChatTyping
|
||||
default:
|
||||
return tgbotapi.ChatUploadPhoto
|
||||
}
|
||||
}
|
||||
|
||||
// renderer renders LaTeX into PNG and PDF bytes.
|
||||
type renderer interface {
|
||||
RenderFormats(ctx context.Context, preamble, expression string, dpi int, formats render.Formats) (*render.Result, error)
|
||||
}
|
||||
|
||||
// settingsProvider holds per-chat settings.
|
||||
type settingsProvider interface {
|
||||
Get(chatID int64) store.Settings
|
||||
Set(chatID int64, settings store.Settings) error
|
||||
}
|
||||
|
||||
// recorder persists users, group memberships and render statistics, and
|
||||
// reports the aggregates behind /stats.
|
||||
type recorder interface {
|
||||
RecordUser(user store.User)
|
||||
RecordGroup(chat store.Chat, member bool)
|
||||
RecordRender(stat store.RenderStat)
|
||||
Stats() store.Stats
|
||||
}
|
||||
|
||||
// fileCache stores correctly typed Telegram file_ids captured from ordinary
|
||||
// chat sends. Inline answers can reuse those IDs without rendering or
|
||||
// uploading anything.
|
||||
type fileCache interface {
|
||||
FileRefs(hash string) store.FileRefs
|
||||
SavePhotoID(hash, fileID string)
|
||||
SaveDocID(hash, fileID string)
|
||||
}
|
||||
|
||||
// FileStore stores rendered outputs and returns presigned URLs that Telegram
|
||||
// can fetch without the bot sending a temporary message to any chat.
|
||||
type FileStore interface {
|
||||
Put(ctx context.Context, ext, contentType string, data []byte) (string, error)
|
||||
}
|
||||
|
||||
// Bot bridges Telegram and the local LaTeX renderer.
|
||||
type Bot struct {
|
||||
api *tgbotapi.BotAPI
|
||||
cfg *config.Config
|
||||
ren renderer
|
||||
user settingsProvider
|
||||
rec recorder
|
||||
cache fileCache
|
||||
files FileStore
|
||||
tracer trace.Tracer
|
||||
|
||||
lastPoll atomic.Int64 // unix time of the last successful getUpdates
|
||||
|
||||
settingsMu sync.Mutex
|
||||
settingsCache map[int64]store.Settings
|
||||
refsMu sync.Mutex
|
||||
refsCache map[string]refsEntry
|
||||
urlsMu sync.Mutex
|
||||
urlsCache map[string]urlEntry
|
||||
}
|
||||
|
||||
const refsCacheTTL = 15 * time.Minute
|
||||
|
||||
type refsEntry struct {
|
||||
refs store.FileRefs
|
||||
at time.Time
|
||||
}
|
||||
|
||||
// urlEntry caches the presigned URLs of a recent render so repeated
|
||||
// identical inline queries skip rendering and uploading again.
|
||||
type urlEntry struct {
|
||||
photoURL string
|
||||
pdfURL string
|
||||
expires time.Time
|
||||
}
|
||||
|
||||
// New authenticates the bot and instruments its HTTP client.
|
||||
func New(cfg *config.Config, renderer renderer, users settingsProvider, rec recorder, cache fileCache, files FileStore) (*Bot, error) {
|
||||
api, err := tgbotapi.NewBotAPI(cfg.TelegramBotToken)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("telegram bot init: %w", err)
|
||||
}
|
||||
api.Client = &http.Client{Transport: otelhttp.NewTransport(http.DefaultTransport)}
|
||||
log.Printf("Authorised as @%s", api.Self.UserName)
|
||||
return &Bot{
|
||||
api: api,
|
||||
cfg: cfg,
|
||||
ren: renderer,
|
||||
user: users,
|
||||
rec: rec,
|
||||
cache: cache,
|
||||
files: files,
|
||||
tracer: otel.Tracer("latexgrambot/bot"),
|
||||
settingsCache: map[int64]store.Settings{},
|
||||
refsCache: map[string]refsEntry{},
|
||||
urlsCache: map[string]urlEntry{},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Healthy reports whether Telegram polling has been failing for a while
|
||||
// (e.g. a 409 conflict because another process polls the same token).
|
||||
func (b *Bot) Healthy() error {
|
||||
last := b.lastPoll.Load()
|
||||
if last == 0 {
|
||||
return nil // not started yet
|
||||
}
|
||||
if age := time.Since(time.Unix(last, 0)); age > pollStale {
|
||||
return fmt.Errorf("no successful Telegram poll for %s", age.Round(time.Second))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Run polls Telegram until ctx is cancelled. Updates are processed by a small
|
||||
// worker pool so slow chat renders do not stall inline queries.
|
||||
func (b *Bot) Run(ctx context.Context) error {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
workers := b.cfg.MaxConcurrency + 2
|
||||
updatesCh := make(chan tgbotapi.Update, workers*2)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for upd := range updatesCh {
|
||||
b.handleUpdate(upd)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
u := tgbotapi.NewUpdate(0)
|
||||
u.Timeout = pollTimeout
|
||||
u.AllowedUpdates = []string{"message", "inline_query", "my_chat_member"}
|
||||
b.lastPoll.Store(time.Now().Unix())
|
||||
log.Println("Listening for Telegram updates")
|
||||
for ctx.Err() == nil {
|
||||
updates, err := b.api.GetUpdates(u)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
log.Printf("telegram getUpdates failed: %v", err)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-time.After(time.Second):
|
||||
}
|
||||
continue
|
||||
}
|
||||
// A successful response (even an empty one) proves the polling
|
||||
// loop is healthy; long-poll timeouts with no updates are normal.
|
||||
b.lastPoll.Store(time.Now().Unix())
|
||||
if len(updates) > 0 {
|
||||
u.Offset = updates[len(updates)-1].UpdateID + 1
|
||||
for _, upd := range updates {
|
||||
updatesCh <- upd
|
||||
}
|
||||
}
|
||||
}
|
||||
close(updatesCh)
|
||||
wg.Wait()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *Bot) handleUpdate(upd tgbotapi.Update) {
|
||||
switch {
|
||||
case upd.Message != nil:
|
||||
b.handleMessage(upd.Message)
|
||||
case upd.InlineQuery != nil:
|
||||
b.handleInlineQuery(upd.InlineQuery)
|
||||
case upd.MyChatMember != nil:
|
||||
b.handleMyChatMember(upd.MyChatMember)
|
||||
}
|
||||
}
|
||||
|
||||
// handleMyChatMember records groups the bot is added to or removed from.
|
||||
func (b *Bot) handleMyChatMember(update *tgbotapi.ChatMemberUpdated) {
|
||||
status := update.NewChatMember.Status
|
||||
member := status == "member" || status == "administrator" || status == "creator" || status == "restricted"
|
||||
go b.rec.RecordGroup(store.Chat{
|
||||
ID: update.Chat.ID,
|
||||
Type: update.Chat.Type,
|
||||
Title: update.Chat.Title,
|
||||
Username: update.Chat.UserName,
|
||||
}, member)
|
||||
}
|
||||
|
||||
func (b *Bot) handleMessage(msg *tgbotapi.Message) {
|
||||
ctx, span := b.tracer.Start(context.Background(), "handle_message",
|
||||
trace.WithAttributes(attribute.Int64("chat.id", msg.Chat.ID)))
|
||||
defer span.End()
|
||||
|
||||
if msg.From != nil {
|
||||
go b.rec.RecordUser(store.User{
|
||||
ID: msg.From.ID,
|
||||
Username: msg.From.UserName,
|
||||
FirstName: msg.From.FirstName,
|
||||
LastName: msg.From.LastName,
|
||||
})
|
||||
}
|
||||
var userID int64
|
||||
if msg.From != nil {
|
||||
userID = msg.From.ID
|
||||
}
|
||||
if msg.IsCommand() {
|
||||
b.handleCommand(msg, userID)
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(msg.Text) == "" {
|
||||
return
|
||||
}
|
||||
b.renderExpression(ctx, msg.Chat.ID, userID, msg.Text, formatRich)
|
||||
}
|
||||
|
||||
// handleCommand routes the known commands. Unknown commands (e.g. \frac) are
|
||||
// rendered as ordinary LaTeX and answered with the default rich text format.
|
||||
func (b *Bot) handleCommand(msg *tgbotapi.Message, userID int64) {
|
||||
switch msg.Command() {
|
||||
case "start":
|
||||
b.sendText(msg.Chat.ID, welcomeText(b.api.Self.UserName))
|
||||
case "help":
|
||||
b.sendText(msg.Chat.ID, helpText(b.api.Self.UserName))
|
||||
case "preamble":
|
||||
b.showPreamble(msg)
|
||||
case "setpreamble":
|
||||
b.setPreamble(msg)
|
||||
case "resetpreamble":
|
||||
b.resetPreamble(msg)
|
||||
case "setdpi":
|
||||
b.setDPI(msg)
|
||||
case "settings":
|
||||
b.showSettings(msg)
|
||||
case "stats":
|
||||
b.sendText(msg.Chat.ID, statsText(b.rec.Stats()))
|
||||
case "pdf", "image", "rich":
|
||||
b.renderCommand(msg, userID)
|
||||
default:
|
||||
b.renderExpression(context.Background(), msg.Chat.ID, userID, msg.Text, formatRich)
|
||||
}
|
||||
}
|
||||
|
||||
// renderCommand renders the arguments of a /pdf, /image or /rich command and
|
||||
// replies with exactly that format.
|
||||
func (b *Bot) renderCommand(msg *tgbotapi.Message, userID int64) {
|
||||
format := formatImage
|
||||
switch msg.Command() {
|
||||
case "pdf":
|
||||
format = formatPDF
|
||||
case "rich":
|
||||
format = formatRich
|
||||
}
|
||||
expression := strings.TrimSpace(msg.CommandArguments())
|
||||
if expression == "" {
|
||||
b.sendText(msg.Chat.ID, fmt.Sprintf("Usage: /%s <LaTeX code>", msg.Command()))
|
||||
return
|
||||
}
|
||||
b.renderExpression(context.Background(), msg.Chat.ID, userID, expression, format)
|
||||
}
|
||||
|
||||
// renderExpression renders the expression and sends only the requested
|
||||
// format back to the chat, recording the outcome in the statistics.
|
||||
func (b *Bot) renderExpression(ctx context.Context, chatID, userID int64, expression string, format outputFormat) {
|
||||
started := time.Now()
|
||||
if len(expression) > b.cfg.MaxExpressionLen {
|
||||
b.sendText(chatID, fmt.Sprintf("Expression is too long (max %d characters).", b.cfg.MaxExpressionLen))
|
||||
return
|
||||
}
|
||||
if format == formatRich {
|
||||
b.sendRichText(chatID, expression)
|
||||
go b.rec.RecordRender(store.RenderStat{
|
||||
Source: "chat", Format: "rich", ChatID: chatID, UserID: userID,
|
||||
ExpressionLen: len(expression), OK: true, DurationMs: time.Since(started).Milliseconds(),
|
||||
})
|
||||
return
|
||||
}
|
||||
if _, err := b.api.Send(tgbotapi.NewChatAction(chatID, format.chatAction())); err != nil {
|
||||
log.Printf("send chat action to %d: %v", chatID, err)
|
||||
}
|
||||
|
||||
settings := b.getSettings(chatID)
|
||||
dpi := b.dpiFor(settings)
|
||||
hash := contentHash(settings.Preamble, expression, dpi)
|
||||
formats := render.FormatPNG
|
||||
if format == formatPDF {
|
||||
formats = render.FormatPDF
|
||||
}
|
||||
renderCtx, cancel := context.WithTimeout(ctx, b.cfg.RenderTimeout)
|
||||
defer cancel()
|
||||
result, err := b.ren.RenderFormats(renderCtx, settings.Preamble, expression, dpi, formats)
|
||||
|
||||
formatName := "png"
|
||||
if format == formatPDF {
|
||||
formatName = "pdf"
|
||||
}
|
||||
stat := store.RenderStat{
|
||||
Source: "chat", Format: formatName, ChatID: chatID, UserID: userID,
|
||||
DPI: dpi, ExpressionLen: len(expression),
|
||||
OK: err == nil, DurationMs: time.Since(started).Milliseconds(),
|
||||
}
|
||||
if err != nil {
|
||||
stat.Error = err.Error()
|
||||
b.sendRenderError(chatID, err)
|
||||
go b.rec.RecordRender(stat)
|
||||
return
|
||||
}
|
||||
go b.rec.RecordRender(stat)
|
||||
|
||||
switch format {
|
||||
case formatPDF:
|
||||
if fileID := b.sendDocument(chatID, result.PDF); fileID != "" && b.cache != nil {
|
||||
b.saveDocID(hash, fileID)
|
||||
}
|
||||
default:
|
||||
if fileID := b.sendPhoto(chatID, result.PNG); fileID != "" && b.cache != nil {
|
||||
b.savePhotoID(hash, fileID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// sendRichText replies with the rich-message rendering of the expression: a
|
||||
// single native math block, no source echo. Failures are non-fatal. The
|
||||
// pinned library predates the rich message API, so the request is built as
|
||||
// raw params.
|
||||
func (b *Bot) sendRichText(chatID int64, expression string) {
|
||||
if !b.cfg.RichTextEnabled {
|
||||
return
|
||||
}
|
||||
if _, err := b.api.Send(tgbotapi.NewChatAction(chatID, tgbotapi.ChatTyping)); err != nil {
|
||||
log.Printf("send chat action to %d: %v", chatID, err)
|
||||
}
|
||||
params := tgbotapi.Params{}
|
||||
params.AddNonZero64("chat_id", chatID)
|
||||
if err := params.AddInterface("rich_message", rich.Document(expression, b.cfg.RichDefaultMath)); err != nil {
|
||||
log.Printf("encode rich message: %v", err)
|
||||
return
|
||||
}
|
||||
if _, err := b.api.MakeRequest("sendRichMessage", params); err != nil {
|
||||
log.Printf("sendRichMessage to %d: %v", chatID, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bot) handleInlineQuery(query *tgbotapi.InlineQuery) {
|
||||
ctx, span := b.tracer.Start(context.Background(), "answer_inline_query",
|
||||
trace.WithAttributes(attribute.String("chat.type", query.ChatType)))
|
||||
defer span.End()
|
||||
|
||||
go b.rec.RecordUser(store.User{
|
||||
ID: query.From.ID,
|
||||
Username: query.From.UserName,
|
||||
FirstName: query.From.FirstName,
|
||||
LastName: query.From.LastName,
|
||||
})
|
||||
|
||||
answer := tgbotapi.InlineConfig{
|
||||
InlineQueryID: query.ID,
|
||||
Results: b.inlineResults(ctx, query),
|
||||
IsPersonal: true,
|
||||
}
|
||||
if _, err := b.api.Request(answer); err != nil {
|
||||
log.Printf("answer inline query %s: %v", query.ID, err)
|
||||
}
|
||||
}
|
||||
|
||||
// getSettings serves settings from the write-through cache; every write
|
||||
// goes through saveSettings, so the cache stays consistent.
|
||||
func (b *Bot) getSettings(chatID int64) store.Settings {
|
||||
b.settingsMu.Lock()
|
||||
settings, ok := b.settingsCache[chatID]
|
||||
b.settingsMu.Unlock()
|
||||
if ok {
|
||||
return settings
|
||||
}
|
||||
settings = b.user.Get(chatID)
|
||||
b.settingsMu.Lock()
|
||||
if b.settingsCache == nil {
|
||||
b.settingsCache = map[int64]store.Settings{}
|
||||
}
|
||||
b.settingsCache[chatID] = settings
|
||||
b.settingsMu.Unlock()
|
||||
return settings
|
||||
}
|
||||
|
||||
func (b *Bot) saveSettings(chatID int64, settings store.Settings) error {
|
||||
if err := b.user.Set(chatID, settings); err != nil {
|
||||
return err
|
||||
}
|
||||
b.settingsMu.Lock()
|
||||
if b.settingsCache == nil {
|
||||
b.settingsCache = map[int64]store.Settings{}
|
||||
}
|
||||
b.settingsCache[chatID] = settings
|
||||
b.settingsMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// fileRefs serves Telegram file_ids from a short-lived cache to keep the
|
||||
// inline hot path free of database round-trips.
|
||||
func (b *Bot) fileRefs(hash string) store.FileRefs {
|
||||
b.refsMu.Lock()
|
||||
entry, ok := b.refsCache[hash]
|
||||
b.refsMu.Unlock()
|
||||
if ok && time.Since(entry.at) < refsCacheTTL {
|
||||
return entry.refs
|
||||
}
|
||||
refs := store.FileRefs{}
|
||||
if b.cache != nil {
|
||||
refs = b.cache.FileRefs(hash)
|
||||
}
|
||||
b.refsMu.Lock()
|
||||
if b.refsCache == nil {
|
||||
b.refsCache = map[string]refsEntry{}
|
||||
}
|
||||
b.refsCache[hash] = refsEntry{refs: refs, at: time.Now()}
|
||||
b.refsMu.Unlock()
|
||||
return refs
|
||||
}
|
||||
|
||||
func (b *Bot) savePhotoID(hash, fileID string) {
|
||||
if b.cache == nil {
|
||||
return
|
||||
}
|
||||
b.savePhotoID(hash, fileID)
|
||||
b.refsMu.Lock()
|
||||
entry := b.refsCache[hash]
|
||||
entry.refs.PhotoID = fileID
|
||||
entry.at = time.Now()
|
||||
b.refsCache[hash] = entry
|
||||
b.refsMu.Unlock()
|
||||
}
|
||||
|
||||
func (b *Bot) saveDocID(hash, fileID string) {
|
||||
if b.cache == nil {
|
||||
return
|
||||
}
|
||||
b.saveDocID(hash, fileID)
|
||||
b.refsMu.Lock()
|
||||
entry := b.refsCache[hash]
|
||||
entry.refs.DocID = fileID
|
||||
entry.at = time.Now()
|
||||
b.refsCache[hash] = entry
|
||||
b.refsMu.Unlock()
|
||||
}
|
||||
|
||||
// cachedURLs returns the presigned URLs of a recent render of the same
|
||||
// content, so repeated identical queries skip render and upload entirely.
|
||||
func (b *Bot) cachedURLs(hash string) (photoURL, pdfURL string, ok bool) {
|
||||
b.urlsMu.Lock()
|
||||
defer b.urlsMu.Unlock()
|
||||
entry, found := b.urlsCache[hash]
|
||||
if !found || time.Now().After(entry.expires) {
|
||||
return "", "", false
|
||||
}
|
||||
return entry.photoURL, entry.pdfURL, true
|
||||
}
|
||||
|
||||
func (b *Bot) cacheURLs(hash, photoURL, pdfURL string) {
|
||||
ttl := b.cfg.S3PresignTTL - 10*time.Minute
|
||||
if ttl < 5*time.Minute {
|
||||
ttl = 5 * time.Minute
|
||||
}
|
||||
b.urlsMu.Lock()
|
||||
if b.urlsCache == nil {
|
||||
b.urlsCache = map[string]urlEntry{}
|
||||
}
|
||||
b.urlsCache[hash] = urlEntry{photoURL: photoURL, pdfURL: pdfURL, expires: time.Now().Add(ttl)}
|
||||
b.urlsMu.Unlock()
|
||||
}
|
||||
|
||||
// inlineResults stores rendered files in private object storage and returns
|
||||
// URL-based results with short-lived signatures. Telegram fetches the selected
|
||||
// file directly; generating an inline answer never sends or deletes a message
|
||||
// in the user's chat.
|
||||
func (b *Bot) inlineResults(ctx context.Context, query *tgbotapi.InlineQuery) []interface{} {
|
||||
raw := strings.TrimSpace(query.Query)
|
||||
if raw == "" {
|
||||
const example = `\int_0^\infty e^{-x^2}\,dx`
|
||||
return []interface{}{b.article(
|
||||
"empty", "Type a LaTeX expression", example,
|
||||
"e.g. "+example,
|
||||
)}
|
||||
}
|
||||
if len(raw) > b.cfg.MaxExpressionLen {
|
||||
return []interface{}{b.article(
|
||||
"long", "Expression too long",
|
||||
fmt.Sprintf("The expression is limited to %d characters.", b.cfg.MaxExpressionLen),
|
||||
fmt.Sprintf("The expression is limited to %d characters.", b.cfg.MaxExpressionLen),
|
||||
)}
|
||||
}
|
||||
expr := raw
|
||||
|
||||
settings := b.getSettings(query.From.ID)
|
||||
dpi := b.dpiFor(settings)
|
||||
hash := contentHash(settings.Preamble, expr, dpi)
|
||||
id := hash[:40]
|
||||
|
||||
refs := b.fileRefs(hash)
|
||||
needPhoto := refs.PhotoID == ""
|
||||
needDoc := refs.DocID == ""
|
||||
var photoURL, pdfURL string
|
||||
if needPhoto || needDoc {
|
||||
if cachedPhoto, cachedPDF, ok := b.cachedURLs(hash); ok {
|
||||
photoURL, pdfURL = cachedPhoto, cachedPDF
|
||||
needPhoto, needDoc = false, false
|
||||
}
|
||||
}
|
||||
if b.files != nil && (needPhoto || needDoc) {
|
||||
renderCtx, cancel := context.WithTimeout(ctx, b.cfg.RenderTimeout)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
result, err := b.ren.RenderFormats(renderCtx, settings.Preamble, expr, dpi, render.FormatJPEG|render.FormatPDF)
|
||||
stat := store.RenderStat{
|
||||
Source: "inline", Format: "inline", UserID: query.From.ID,
|
||||
DPI: dpi, ExpressionLen: len(expr), OK: err == nil,
|
||||
DurationMs: time.Since(started).Milliseconds(),
|
||||
}
|
||||
if err != nil {
|
||||
stat.Error = err.Error()
|
||||
go b.rec.RecordRender(stat)
|
||||
return []interface{}{b.errorArticle(err)}
|
||||
}
|
||||
go b.rec.RecordRender(stat)
|
||||
|
||||
// Upload both formats concurrently: each Put is an object store
|
||||
// round-trip.
|
||||
var wg sync.WaitGroup
|
||||
var photoPutURL, pdfPutURL string
|
||||
if needPhoto {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
storedURL, putErr := b.files.Put(ctx, "jpg", "image/jpeg", result.JPEG)
|
||||
if putErr != nil {
|
||||
log.Printf("store inline picture: %v", putErr)
|
||||
return
|
||||
}
|
||||
photoPutURL = storedURL
|
||||
}()
|
||||
}
|
||||
if needDoc {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
storedURL, putErr := b.files.Put(ctx, "pdf", "application/pdf", result.PDF)
|
||||
if putErr != nil {
|
||||
log.Printf("store inline document: %v", putErr)
|
||||
return
|
||||
}
|
||||
pdfPutURL = storedURL
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if photoPutURL != "" {
|
||||
photoURL = photoPutURL
|
||||
}
|
||||
if pdfPutURL != "" {
|
||||
pdfURL = pdfPutURL
|
||||
}
|
||||
if photoPutURL != "" || pdfPutURL != "" {
|
||||
b.cacheURLs(hash, photoURL, pdfURL)
|
||||
}
|
||||
}
|
||||
|
||||
results := []interface{}{}
|
||||
if refs.PhotoID != "" {
|
||||
photo := tgbotapi.NewInlineQueryResultCachedPhoto("photo-"+id, refs.PhotoID)
|
||||
photo.Title = "PNG"
|
||||
results = append(results, photo)
|
||||
} else if photoURL != "" {
|
||||
photo := tgbotapi.NewInlineQueryResultPhoto("photo-"+id, photoURL)
|
||||
photo.ThumbURL = photoURL
|
||||
photo.Title = "PNG"
|
||||
results = append(results, photo)
|
||||
}
|
||||
if refs.DocID != "" {
|
||||
document := tgbotapi.NewInlineQueryResultCachedDocument("doc-"+id, refs.DocID, "PDF")
|
||||
document.Title = "PDF"
|
||||
results = append(results, document)
|
||||
} else if pdfURL != "" {
|
||||
document := tgbotapi.NewInlineQueryResultDocument("doc-"+id, pdfURL, "formula.pdf", "application/pdf")
|
||||
document.Title = "PDF"
|
||||
results = append(results, document)
|
||||
}
|
||||
if b.cfg.RichTextEnabled {
|
||||
results = append(results, b.richTextResult("formula-rich-"+id, raw))
|
||||
}
|
||||
if len(results) == 0 {
|
||||
return []interface{}{b.article(
|
||||
"nofiles", "Rendering unavailable",
|
||||
"No output format is available for inline answers.",
|
||||
"No output format is available for inline answers.",
|
||||
)}
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
// The formula renders inline with a one-space gap from the bot mention.
|
||||
func (b *Bot) richTextResult(id, expression string) interface{} {
|
||||
result := tgbotapi.NewInlineQueryResultArticle(id, "Rich text", "")
|
||||
result.Description = truncate(expression, 64)
|
||||
result.InputMessageContent = inputRichMessageContent{RichMessage: rich.InlineDocument(expression, b.cfg.RichDefaultMath)}
|
||||
return result
|
||||
}
|
||||
|
||||
// inputRichMessageContent implements the Bot API 10.x InputRichMessageContent
|
||||
// used as inline input_message_content.
|
||||
type inputRichMessageContent struct {
|
||||
RichMessage *rich.Message `json:"rich_message"`
|
||||
}
|
||||
|
||||
// sendPhoto sends a rendered picture to the chat and returns its reusable
|
||||
// Telegram file_id.
|
||||
func (b *Bot) sendPhoto(chatID int64, png []byte) string {
|
||||
photo := tgbotapi.NewPhoto(chatID, tgbotapi.FileBytes{Name: "formula.png", Bytes: png})
|
||||
msg, err := b.api.Send(photo)
|
||||
if err != nil {
|
||||
log.Printf("send png to %d: %v", chatID, err)
|
||||
return ""
|
||||
}
|
||||
var best tgbotapi.PhotoSize
|
||||
for _, size := range msg.Photo {
|
||||
if size.FileSize > best.FileSize {
|
||||
best = size
|
||||
}
|
||||
}
|
||||
return best.FileID
|
||||
}
|
||||
|
||||
// sendDocument sends the rendered PDF to the chat and returns its reusable
|
||||
// Telegram file_id.
|
||||
func (b *Bot) sendDocument(chatID int64, pdf []byte) string {
|
||||
document := tgbotapi.NewDocument(chatID, tgbotapi.FileBytes{Name: "formula.pdf", Bytes: pdf})
|
||||
msg, err := b.api.Send(document)
|
||||
if err != nil {
|
||||
log.Printf("send pdf to %d: %v", chatID, err)
|
||||
return ""
|
||||
}
|
||||
if msg.Document == nil {
|
||||
return ""
|
||||
}
|
||||
return msg.Document.FileID
|
||||
}
|
||||
|
||||
// article builds a plain article result; tapping it sends text back to the
|
||||
// chat. description, when non-empty, is shown under the title in the query
|
||||
// menu.
|
||||
func (b *Bot) article(id, title, text, description string) interface{} {
|
||||
result := tgbotapi.NewInlineQueryResultArticle(id, title, text)
|
||||
if description != "" {
|
||||
result.Description = truncate(description, 200)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (b *Bot) errorArticle(err error) interface{} {
|
||||
var latexErr *render.LatexError
|
||||
switch {
|
||||
case errors.As(err, &latexErr):
|
||||
return b.article("latex-error", "LaTeX error", latexErr.Excerpt, latexErr.Excerpt)
|
||||
case errors.Is(err, context.DeadlineExceeded):
|
||||
return b.article("timeout", "Rendering timed out",
|
||||
"Rendering took too long. Simplify the expression or check the preamble.",
|
||||
"Rendering took too long. Simplify the expression or check the preamble.")
|
||||
default:
|
||||
return b.article("failed", "Rendering failed", "Something went wrong while rendering.",
|
||||
"Something went wrong while rendering.")
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bot) sendRenderError(chatID int64, err error) {
|
||||
var latexErr *render.LatexError
|
||||
switch {
|
||||
case errors.As(err, &latexErr):
|
||||
b.sendText(chatID, "LaTeX error:\n\n"+latexErr.Excerpt)
|
||||
case errors.Is(err, context.DeadlineExceeded):
|
||||
b.sendText(chatID, "Rendering timed out. Simplify the expression or check the preamble.")
|
||||
default:
|
||||
b.sendText(chatID, "Rendering failed: "+err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bot) sendText(chatID int64, text string) {
|
||||
msg := tgbotapi.NewMessage(chatID, text)
|
||||
if _, err := b.api.Send(msg); err != nil {
|
||||
log.Printf("reply to %d: %v", chatID, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bot) showPreamble(msg *tgbotapi.Message) {
|
||||
settings := b.getSettings(msg.Chat.ID)
|
||||
if settings.Preamble == "" {
|
||||
b.sendText(msg.Chat.ID, "You are using the default preamble. Set your own with /setpreamble.")
|
||||
return
|
||||
}
|
||||
b.sendText(msg.Chat.ID, "Your preamble additions:\n\n"+truncate(settings.Preamble, displayLimit))
|
||||
}
|
||||
|
||||
func (b *Bot) setPreamble(msg *tgbotapi.Message) {
|
||||
code := strings.TrimSpace(msg.CommandArguments())
|
||||
if code == "" {
|
||||
b.sendText(msg.Chat.ID, "Usage: /setpreamble <LaTeX packages and definitions>\n"+
|
||||
"Your additions are loaded after the default package set.")
|
||||
return
|
||||
}
|
||||
if len(code) > b.cfg.MaxPreambleLen {
|
||||
b.sendText(msg.Chat.ID, fmt.Sprintf("Preamble is too long (max %d characters).", b.cfg.MaxPreambleLen))
|
||||
return
|
||||
}
|
||||
settings := b.getSettings(msg.Chat.ID)
|
||||
settings.Preamble = code
|
||||
if err := b.saveSettings(msg.Chat.ID, settings); err != nil {
|
||||
log.Printf("store preamble for %d: %v", msg.Chat.ID, err)
|
||||
b.sendText(msg.Chat.ID, "Could not save the preamble, sorry.")
|
||||
return
|
||||
}
|
||||
b.sendText(msg.Chat.ID, "Preamble additions updated. Use /resetpreamble to remove them.")
|
||||
}
|
||||
|
||||
func (b *Bot) resetPreamble(msg *tgbotapi.Message) {
|
||||
settings := b.getSettings(msg.Chat.ID)
|
||||
settings.Preamble = ""
|
||||
if err := b.saveSettings(msg.Chat.ID, settings); err != nil {
|
||||
log.Printf("reset preamble for %d: %v", msg.Chat.ID, err)
|
||||
b.sendText(msg.Chat.ID, "Could not reset the preamble, sorry.")
|
||||
return
|
||||
}
|
||||
b.sendText(msg.Chat.ID, "Preamble reset to the default.")
|
||||
}
|
||||
|
||||
func (b *Bot) setDPI(msg *tgbotapi.Message) {
|
||||
value, err := strconv.Atoi(strings.TrimSpace(msg.CommandArguments()))
|
||||
if err != nil || value < minDPI || value > maxDPI {
|
||||
b.sendText(msg.Chat.ID, fmt.Sprintf("Usage: /setdpi <%d-%d> (default %d).", minDPI, maxDPI, b.cfg.RenderDPI))
|
||||
return
|
||||
}
|
||||
settings := b.getSettings(msg.Chat.ID)
|
||||
settings.DPI = value
|
||||
if err := b.saveSettings(msg.Chat.ID, settings); err != nil {
|
||||
log.Printf("store dpi for %d: %v", msg.Chat.ID, err)
|
||||
b.sendText(msg.Chat.ID, "Could not save the resolution, sorry.")
|
||||
return
|
||||
}
|
||||
b.sendText(msg.Chat.ID, fmt.Sprintf("Resolution set to %d DPI. Use /setdpi %d to restore the default.", value, b.cfg.RenderDPI))
|
||||
}
|
||||
|
||||
func (b *Bot) showSettings(msg *tgbotapi.Message) {
|
||||
settings := b.getSettings(msg.Chat.ID)
|
||||
dpi := b.dpiFor(settings)
|
||||
var text strings.Builder
|
||||
fmt.Fprintf(&text, "Resolution: %d DPI\n", dpi)
|
||||
if settings.Preamble == "" {
|
||||
text.WriteString("Preamble: default")
|
||||
} else {
|
||||
text.WriteString("Preamble additions:\n\n")
|
||||
text.WriteString(truncate(settings.Preamble, displayLimit))
|
||||
}
|
||||
b.sendText(msg.Chat.ID, text.String())
|
||||
}
|
||||
|
||||
func (b *Bot) dpiFor(settings store.Settings) int {
|
||||
if settings.DPI > 0 {
|
||||
return settings.DPI
|
||||
}
|
||||
return b.cfg.RenderDPI
|
||||
}
|
||||
|
||||
func truncate(text string, limit int) string {
|
||||
if len(text) <= limit {
|
||||
return text
|
||||
}
|
||||
return text[:limit] + "..."
|
||||
}
|
||||
|
||||
// contentHash identifies the rendered output of an expression with a given
|
||||
// preamble and DPI. It also names the stored files.
|
||||
func contentHash(preamble, expression string, dpi int) string {
|
||||
sum := sha256.Sum256([]byte(fmt.Sprintf("%s\x00%s\x00%s\x00%d", renderCacheVersion, preamble, expression, dpi)))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
// statsText renders the public /stats summary.
|
||||
func statsText(stats store.Stats) string {
|
||||
return fmt.Sprintf(
|
||||
"latexgrambot stats\n\n"+
|
||||
"Renders: %d (chat %d, inline %d)\n"+
|
||||
"Failed renders: %d\n"+
|
||||
"Renders in the last 24h: %d\n"+
|
||||
"Users: %d\n"+
|
||||
"Groups: %d",
|
||||
stats.Renders, stats.ChatRenders, stats.InlineRenders,
|
||||
stats.FailedRenders, stats.Renders24h, stats.Users, stats.Groups)
|
||||
}
|
||||
|
||||
func welcomeText(username string) string {
|
||||
return fmt.Sprintf(
|
||||
"Hi! I render LaTeX into rich text messages.\n\n"+
|
||||
"Send me any LaTeX expression in this chat and I will reply with a rich text message. "+
|
||||
"Use /image or /pdf to get files instead. "+
|
||||
"You can also use me inline: in any chat, type @%s followed by an expression.\n\n"+
|
||||
"See /help for the available commands.", username)
|
||||
}
|
||||
|
||||
func helpText(username string) string {
|
||||
return fmt.Sprintf(
|
||||
"Usage\n"+
|
||||
"- Chat mode: send a message containing LaTeX code (e.g. \\int_0^\\infty e^{-x^2}\\,dx) "+
|
||||
"and I reply with a rich text message.\n"+
|
||||
"- /rich <code> — reply with the rich text message (default)\n"+
|
||||
"- /image <code> — reply with the PNG picture\n"+
|
||||
"- /pdf <code> — reply with the PDF document\n"+
|
||||
"- Inline mode: in any chat, type @%s <expression> and pick the rich text, picture or PDF result.\n"+
|
||||
"- The document class is fixed to the standalone preview layout; "+
|
||||
"your code goes between \\begin{document} and \\end{document}.\n\n"+
|
||||
"Commands\n"+
|
||||
"/preamble — show your current preamble\n"+
|
||||
"/setpreamble <code> — replace the packages block with your own preamble\n"+
|
||||
"/resetpreamble — restore the default preamble\n"+
|
||||
"/setdpi <72-1200> — set the picture resolution (default 600)\n"+
|
||||
"/settings — show your current settings\n"+
|
||||
"/stats — show bot usage statistics\n"+
|
||||
"/help — this message", username)
|
||||
}
|
||||
269
internal/bot/bot_test.go
Normal file
269
internal/bot/bot_test.go
Normal file
@@ -0,0 +1,269 @@
|
||||
package bot
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"latexgrambot/internal/config"
|
||||
"latexgrambot/internal/render"
|
||||
"latexgrambot/internal/rich"
|
||||
"latexgrambot/internal/store"
|
||||
|
||||
tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
|
||||
"github.com/google/uuid"
|
||||
"go.opentelemetry.io/otel"
|
||||
)
|
||||
|
||||
type stubRenderer struct {
|
||||
results *render.Result
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *stubRenderer) RenderFormats(_ context.Context, _, _ string, _ int, _ render.Formats) (*render.Result, error) {
|
||||
return s.results, s.err
|
||||
}
|
||||
|
||||
type fakeSettings struct {
|
||||
settings store.Settings
|
||||
}
|
||||
|
||||
func (f *fakeSettings) Get(int64) store.Settings { return f.settings }
|
||||
func (f *fakeSettings) Set(int64, store.Settings) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type fakeRecorder struct{}
|
||||
|
||||
func (f *fakeRecorder) RecordUser(store.User) {}
|
||||
func (f *fakeRecorder) RecordGroup(store.Chat, bool) {}
|
||||
func (f *fakeRecorder) RecordRender(store.RenderStat) {}
|
||||
|
||||
func (f *fakeRecorder) Stats() store.Stats {
|
||||
return store.Stats{
|
||||
Renders: 42, ChatRenders: 30, InlineRenders: 12,
|
||||
FailedRenders: 3, Renders24h: 7, Users: 9, Groups: 2,
|
||||
}
|
||||
}
|
||||
|
||||
type fakeCache struct {
|
||||
refs store.FileRefs
|
||||
}
|
||||
|
||||
func (f *fakeCache) FileRefs(string) store.FileRefs { return f.refs }
|
||||
func (f *fakeCache) SavePhotoID(_, fileID string) { f.refs.PhotoID = fileID }
|
||||
func (f *fakeCache) SaveDocID(_, fileID string) { f.refs.DocID = fileID }
|
||||
|
||||
type fakeFiles struct {
|
||||
mu sync.Mutex
|
||||
puts map[string][]byte
|
||||
putCalls int
|
||||
}
|
||||
|
||||
func (f *fakeFiles) Put(_ context.Context, ext, _ string, data []byte) (string, error) {
|
||||
f.mu.Lock()
|
||||
if f.puts == nil {
|
||||
f.puts = map[string][]byte{}
|
||||
}
|
||||
f.puts[ext] = data
|
||||
f.putCalls++
|
||||
f.mu.Unlock()
|
||||
return "https://s3.example/latexgrambot/content/" + uuid.NewString() + "." + ext + "?X-Amz-Signature=test", nil
|
||||
}
|
||||
|
||||
func testBot() *Bot {
|
||||
return &Bot{
|
||||
api: &tgbotapi.BotAPI{Self: tgbotapi.User{UserName: "latexgrambot"}},
|
||||
cfg: &config.Config{MaxExpressionLen: 4000, RichTextEnabled: true, RichDefaultMath: true, RenderDPI: 600},
|
||||
ren: &stubRenderer{results: &render.Result{PNG: []byte("png"), JPEG: []byte("jpg"), PDF: []byte("pdf")}},
|
||||
user: &fakeSettings{},
|
||||
rec: &fakeRecorder{},
|
||||
files: &fakeFiles{},
|
||||
tracer: otel.GetTracerProvider().Tracer("test"),
|
||||
settingsCache: map[int64]store.Settings{},
|
||||
refsCache: map[string]refsEntry{},
|
||||
urlsCache: map[string]urlEntry{},
|
||||
}
|
||||
}
|
||||
|
||||
func article(t *testing.T, result interface{}) tgbotapi.InlineQueryResultArticle {
|
||||
t.Helper()
|
||||
a, ok := result.(tgbotapi.InlineQueryResultArticle)
|
||||
if !ok {
|
||||
t.Fatalf("result is %T, want InlineQueryResultArticle", result)
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
func TestInlineResultsUsesPresignedFiles(t *testing.T) {
|
||||
b := testBot()
|
||||
results := b.inlineResults(context.Background(), &tgbotapi.InlineQuery{Query: "$F=ma$", ChatType: "supergroup", From: &tgbotapi.User{ID: 1}})
|
||||
if len(results) != 3 {
|
||||
t.Fatalf("got %d results, want 3", len(results))
|
||||
}
|
||||
|
||||
photo, ok := results[0].(tgbotapi.InlineQueryResultPhoto)
|
||||
if !ok {
|
||||
t.Fatalf("result 0 is %T, want InlineQueryResultPhoto", results[0])
|
||||
}
|
||||
if !strings.Contains(photo.URL, ".jpg?") || photo.ThumbURL != photo.URL {
|
||||
t.Errorf("photo URL/thumb = %q/%q", photo.URL, photo.ThumbURL)
|
||||
}
|
||||
if photo.Title != "PNG" {
|
||||
t.Errorf("photo title = %q", photo.Title)
|
||||
}
|
||||
|
||||
document, ok := results[1].(tgbotapi.InlineQueryResultDocument)
|
||||
if !ok {
|
||||
t.Fatalf("result 1 is %T, want InlineQueryResultDocument", results[1])
|
||||
}
|
||||
if !strings.Contains(document.URL, ".pdf?") || document.MimeType != "application/pdf" {
|
||||
t.Errorf("document URL/type = %q/%q", document.URL, document.MimeType)
|
||||
}
|
||||
if document.Title != "PDF" {
|
||||
t.Errorf("document title = %q", document.Title)
|
||||
}
|
||||
|
||||
richResult := article(t, results[2])
|
||||
if richResult.Title != "Rich text" {
|
||||
t.Errorf("rich title = %q", richResult.Title)
|
||||
}
|
||||
if richResult.ReplyMarkup != nil {
|
||||
t.Error("rich result must not carry a button")
|
||||
}
|
||||
content, ok := richResult.InputMessageContent.(inputRichMessageContent)
|
||||
if !ok || content.RichMessage == nil || len(content.RichMessage.Blocks) != 1 {
|
||||
t.Fatalf("unexpected rich content: %+v", richResult.InputMessageContent)
|
||||
}
|
||||
parts, ok := content.RichMessage.Blocks[0].Text.([]any)
|
||||
if !ok || len(parts) != 2 {
|
||||
t.Fatalf("unexpected paragraph content: %+v", content.RichMessage.Blocks[0].Text)
|
||||
}
|
||||
math, ok := parts[1].(rich.Inline)
|
||||
if !ok || math.Type != "mathematical_expression" || math.Expression != "F=ma" {
|
||||
t.Fatalf("unexpected inline math node: %+v", parts[1])
|
||||
}
|
||||
files := b.files.(*fakeFiles)
|
||||
if string(files.puts["jpg"]) != "jpg" || string(files.puts["pdf"]) != "pdf" {
|
||||
t.Fatalf("stored files = %#v", files.puts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInlineResultsCachesPresignedURLs(t *testing.T) {
|
||||
b := testBot()
|
||||
query := &tgbotapi.InlineQuery{Query: "$F=ma$", ChatType: "private", From: &tgbotapi.User{ID: 1}}
|
||||
first := b.inlineResults(context.Background(), query)
|
||||
second := b.inlineResults(context.Background(), query)
|
||||
files := b.files.(*fakeFiles)
|
||||
if files.putCalls != 2 {
|
||||
t.Fatalf("put calls = %d, want 2 (one per format, cached afterwards)", files.putCalls)
|
||||
}
|
||||
firstPhoto := first[0].(tgbotapi.InlineQueryResultPhoto)
|
||||
secondPhoto := second[0].(tgbotapi.InlineQueryResultPhoto)
|
||||
if firstPhoto.URL != secondPhoto.URL {
|
||||
t.Fatalf("second query re-rendered: %q vs %q", firstPhoto.URL, secondPhoto.URL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInlineResultsReusesTelegramFileIDs(t *testing.T) {
|
||||
b := testBot()
|
||||
b.cache = &fakeCache{refs: store.FileRefs{PhotoID: "photo-id", DocID: "doc-id"}}
|
||||
b.ren = &stubRenderer{err: &render.LatexError{Excerpt: "must not render"}}
|
||||
results := b.inlineResults(context.Background(), &tgbotapi.InlineQuery{Query: "$x$", ChatType: "private", From: &tgbotapi.User{ID: 1}})
|
||||
if len(results) != 3 {
|
||||
t.Fatalf("got %d results, want 3", len(results))
|
||||
}
|
||||
photo, ok := results[0].(tgbotapi.InlineQueryResultCachedPhoto)
|
||||
if !ok || photo.PhotoID != "photo-id" {
|
||||
t.Fatalf("photo = %#v", results[0])
|
||||
}
|
||||
document, ok := results[1].(tgbotapi.InlineQueryResultCachedDocument)
|
||||
if !ok || document.DocumentID != "doc-id" {
|
||||
t.Fatalf("document = %#v", results[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestInlineResultsNoFileStore(t *testing.T) {
|
||||
b := testBot()
|
||||
b.files = nil
|
||||
b.ren = &stubRenderer{err: &render.LatexError{Excerpt: "must not render"}}
|
||||
results := b.inlineResults(context.Background(), &tgbotapi.InlineQuery{Query: "$x$", ChatType: "private", From: &tgbotapi.User{ID: 1}})
|
||||
if len(results) != 1 || article(t, results[0]).Title != "Rich text" {
|
||||
t.Fatalf("results = %#v", results)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInlineResultsRichDisabled(t *testing.T) {
|
||||
b := testBot()
|
||||
b.cfg.RichTextEnabled = false
|
||||
results := b.inlineResults(context.Background(), &tgbotapi.InlineQuery{Query: "$x$", ChatType: "private", From: &tgbotapi.User{ID: 1}})
|
||||
if len(results) != 2 {
|
||||
t.Fatalf("got %d results, want 2", len(results))
|
||||
}
|
||||
if _, ok := results[0].(tgbotapi.InlineQueryResultPhoto); !ok {
|
||||
t.Fatalf("result 0 is %T, want InlineQueryResultPhoto", results[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestInlineResultsRenderError(t *testing.T) {
|
||||
b := testBot()
|
||||
b.ren = &stubRenderer{err: &render.LatexError{Excerpt: "! Undefined control sequence."}}
|
||||
results := b.inlineResults(context.Background(), &tgbotapi.InlineQuery{Query: "\\bad", ChatType: "private", From: &tgbotapi.User{ID: 1}})
|
||||
if len(results) != 1 {
|
||||
t.Fatalf("got %d results, want 1", len(results))
|
||||
}
|
||||
if got := article(t, results[0]).Title; got != "LaTeX error" {
|
||||
t.Errorf("title = %q, want LaTeX error", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInlineResultsEmptyQuery(t *testing.T) {
|
||||
b := testBot()
|
||||
results := b.inlineResults(context.Background(), &tgbotapi.InlineQuery{Query: " ", ChatType: "private", From: &tgbotapi.User{ID: 1}})
|
||||
if len(results) != 1 {
|
||||
t.Fatalf("got %d results, want 1", len(results))
|
||||
}
|
||||
if got := article(t, results[0]).Title; got != "Type a LaTeX expression" {
|
||||
t.Errorf("title = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestContentHashStability(t *testing.T) {
|
||||
a := contentHash("pre", "x^2", 600)
|
||||
b := contentHash("pre", "x^2", 600)
|
||||
c := contentHash("pre", "x^2", 300)
|
||||
if a != b || len(a) != 64 {
|
||||
t.Fatalf("hash must be stable: %q %q", a, b)
|
||||
}
|
||||
if a == c {
|
||||
t.Fatal("different DPI must change the hash")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatsText(t *testing.T) {
|
||||
text := statsText(store.Stats{
|
||||
Renders: 42, ChatRenders: 30, InlineRenders: 12,
|
||||
FailedRenders: 3, Renders24h: 7, Users: 9, Groups: 2,
|
||||
})
|
||||
for _, fragment := range []string{
|
||||
"Renders: 42 (chat 30, inline 12)",
|
||||
"Failed renders: 3",
|
||||
"Renders in the last 24h: 7",
|
||||
"Users: 9",
|
||||
"Groups: 2",
|
||||
} {
|
||||
if !strings.Contains(text, fragment) {
|
||||
t.Errorf("stats text missing %q:\n%s", fragment, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTruncate(t *testing.T) {
|
||||
if got := truncate("hello", 5); got != "hello" {
|
||||
t.Errorf("truncate = %q", got)
|
||||
}
|
||||
if got := truncate("hello", 3); got != "hel..." {
|
||||
t.Errorf("truncate = %q", got)
|
||||
}
|
||||
}
|
||||
192
internal/config/config.go
Normal file
192
internal/config/config.go
Normal file
@@ -0,0 +1,192 @@
|
||||
// Package config reads the latexgrambot configuration from environment variables.
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Config holds every runtime setting.
|
||||
type Config struct {
|
||||
TelegramBotToken string
|
||||
|
||||
// MongoURI is the shared cluster connection string. MongoDatabase is
|
||||
// the bot's database inside it.
|
||||
MongoURI string
|
||||
MongoDatabase string
|
||||
|
||||
// S3 hosts rendered inline files. All required values must be set
|
||||
// together; otherwise inline mode offers only native rich text.
|
||||
S3Endpoint string
|
||||
S3Region string
|
||||
S3Bucket string
|
||||
S3AccessKeyID string
|
||||
S3SecretAccessKey string
|
||||
S3Prefix string
|
||||
S3PublicEndpoint string
|
||||
S3PresignTTL time.Duration
|
||||
|
||||
HealthAddr string
|
||||
|
||||
// RichTextEnabled controls the rich-message replies (native math block).
|
||||
// The format needs a recent Bot API server; when disabled the bot falls
|
||||
// back to PNG and PDF only.
|
||||
RichTextEnabled bool
|
||||
|
||||
// RichDefaultMath controls how the rich format treats a bare source
|
||||
// (no math delimiters, no text-mode command): true renders it as a math
|
||||
// expression, false parses it as text so formulas need $...$ delimiters.
|
||||
RichDefaultMath bool
|
||||
|
||||
RenderDPI int
|
||||
RenderTimeout time.Duration
|
||||
MaxConcurrency int
|
||||
|
||||
// MaxExpressionLen caps chat-mode expressions. MaxPreambleLen caps the
|
||||
// per-user preamble (Telegram messages are at most 4096 characters).
|
||||
MaxExpressionLen int
|
||||
MaxPreambleLen int
|
||||
|
||||
PdflatexBin string
|
||||
PdftoppmBin string
|
||||
|
||||
// LatexFormat is the preloaded TeX format name (built from the default
|
||||
// preamble at image build time). Empty disables it.
|
||||
LatexFormat string
|
||||
}
|
||||
|
||||
// Load builds the Config and validates the mandatory keys.
|
||||
func Load() (*Config, error) {
|
||||
cfg := &Config{
|
||||
TelegramBotToken: strings.TrimSpace(os.Getenv("TELEGRAM_BOT_TOKEN")),
|
||||
MongoURI: strings.TrimSpace(os.Getenv("MONGODB_URI")),
|
||||
MongoDatabase: getenv("MONGODB_DATABASE", "latexgram"),
|
||||
S3Endpoint: strings.TrimRight(strings.TrimSpace(os.Getenv("S3_ENDPOINT_URL")), "/"),
|
||||
S3Region: getenv("S3_REGION", "us-east-1"),
|
||||
S3Bucket: strings.TrimSpace(os.Getenv("S3_BUCKET")),
|
||||
S3AccessKeyID: strings.TrimSpace(os.Getenv("AWS_ACCESS_KEY_ID")),
|
||||
S3SecretAccessKey: strings.TrimSpace(os.Getenv("AWS_SECRET_ACCESS_KEY")),
|
||||
S3Prefix: strings.Trim(getenv("S3_PREFIX", "content"), "/"),
|
||||
S3PublicEndpoint: strings.TrimRight(strings.TrimSpace(os.Getenv("S3_PUBLIC_ENDPOINT_URL")), "/"),
|
||||
HealthAddr: getenv("HEALTH_ADDR", ":8080"),
|
||||
PdflatexBin: getenv("PDFLATEX_BIN", "pdflatex"),
|
||||
PdftoppmBin: getenv("PDFTOPPM_BIN", "pdftoppm"),
|
||||
LatexFormat: getenv("LATEX_FORMAT", ""),
|
||||
}
|
||||
|
||||
var err error
|
||||
if cfg.RichTextEnabled, err = getenvBool("RICH_TEXT_ENABLED", true); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cfg.RichDefaultMath, err = getenvBool("RICH_DEFAULT_MATH", true); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cfg.RenderDPI, err = getenvInt("RENDER_DPI", 600); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cfg.MaxConcurrency, err = getenvInt("RENDER_MAX_CONCURRENCY", 2); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cfg.MaxExpressionLen, err = getenvInt("MAX_EXPRESSION_LEN", 4000); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cfg.MaxPreambleLen, err = getenvInt("MAX_PREAMBLE_LEN", 4000); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cfg.RenderTimeout, err = getenvDuration("RENDER_TIMEOUT", 30*time.Second); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cfg.S3PresignTTL, err = getenvDuration("S3_PRESIGN_TTL", time.Hour); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if cfg.TelegramBotToken == "" {
|
||||
return nil, fmt.Errorf("missing required environment variable: TELEGRAM_BOT_TOKEN")
|
||||
}
|
||||
if cfg.MongoURI == "" {
|
||||
return nil, fmt.Errorf("missing required environment variable: MONGODB_URI")
|
||||
}
|
||||
if cfg.RenderDPI < 72 || cfg.RenderDPI > 1200 {
|
||||
return nil, fmt.Errorf("RENDER_DPI must be between 72 and 1200")
|
||||
}
|
||||
if cfg.RenderTimeout <= 0 {
|
||||
return nil, fmt.Errorf("RENDER_TIMEOUT must be positive")
|
||||
}
|
||||
if cfg.MaxConcurrency < 1 {
|
||||
return nil, fmt.Errorf("RENDER_MAX_CONCURRENCY must be at least 1")
|
||||
}
|
||||
if cfg.MaxExpressionLen < 1 || cfg.MaxPreambleLen < 1 {
|
||||
return nil, fmt.Errorf("MAX_EXPRESSION_LEN and MAX_PREAMBLE_LEN must be positive")
|
||||
}
|
||||
s3Set := 0
|
||||
for _, value := range []string{
|
||||
cfg.S3Endpoint, cfg.S3Bucket, cfg.S3AccessKeyID,
|
||||
cfg.S3SecretAccessKey, cfg.S3PublicEndpoint,
|
||||
} {
|
||||
if value != "" {
|
||||
s3Set++
|
||||
}
|
||||
}
|
||||
if s3Set != 0 && s3Set != 5 {
|
||||
return nil, fmt.Errorf(
|
||||
"incomplete S3 configuration: set S3_ENDPOINT_URL, S3_BUCKET, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY and S3_PUBLIC_ENDPOINT_URL together")
|
||||
}
|
||||
if cfg.S3PresignTTL <= 0 {
|
||||
return nil, fmt.Errorf("S3_PRESIGN_TTL must be positive")
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
// S3Enabled reports whether private inline PNG/PDF assets can be stored and
|
||||
// exposed through short-lived presigned URLs.
|
||||
func (c *Config) S3Enabled() bool {
|
||||
return c.S3Endpoint != "" && c.S3Bucket != "" &&
|
||||
c.S3AccessKeyID != "" && c.S3SecretAccessKey != "" &&
|
||||
c.S3PublicEndpoint != ""
|
||||
}
|
||||
|
||||
func getenv(key, fallback string) string {
|
||||
if value := strings.TrimSpace(os.Getenv(key)); value != "" {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func getenvInt(key string, fallback int) (int, error) {
|
||||
value := strings.TrimSpace(os.Getenv(key))
|
||||
if value == "" {
|
||||
return fallback, nil
|
||||
}
|
||||
parsed, err := strconv.Atoi(value)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("invalid %s: %w", key, err)
|
||||
}
|
||||
return parsed, nil
|
||||
}
|
||||
|
||||
func getenvBool(key string, fallback bool) (bool, error) {
|
||||
value := strings.TrimSpace(os.Getenv(key))
|
||||
if value == "" {
|
||||
return fallback, nil
|
||||
}
|
||||
parsed, err := strconv.ParseBool(value)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("invalid %s: %w", key, err)
|
||||
}
|
||||
return parsed, nil
|
||||
}
|
||||
|
||||
func getenvDuration(key string, fallback time.Duration) (time.Duration, error) {
|
||||
value := strings.TrimSpace(os.Getenv(key))
|
||||
if value == "" {
|
||||
return fallback, nil
|
||||
}
|
||||
parsed, err := time.ParseDuration(value)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("invalid %s: %w", key, err)
|
||||
}
|
||||
return parsed, nil
|
||||
}
|
||||
143
internal/config/config_test.go
Normal file
143
internal/config/config_test.go
Normal file
@@ -0,0 +1,143 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func setRequired(t *testing.T) {
|
||||
t.Helper()
|
||||
t.Setenv("TELEGRAM_BOT_TOKEN", "123:abc")
|
||||
t.Setenv("MONGODB_URI", "mongodb://example:27017")
|
||||
}
|
||||
|
||||
func TestLoadDefaults(t *testing.T) {
|
||||
setRequired(t)
|
||||
cfg, err := Load()
|
||||
if err != nil {
|
||||
t.Fatalf("Load: %v", err)
|
||||
}
|
||||
if cfg.RenderDPI != 600 {
|
||||
t.Errorf("RenderDPI = %d, want 600", cfg.RenderDPI)
|
||||
}
|
||||
if cfg.RenderTimeout != 30*time.Second {
|
||||
t.Errorf("RenderTimeout = %s, want 30s", cfg.RenderTimeout)
|
||||
}
|
||||
if cfg.MaxConcurrency != 2 {
|
||||
t.Errorf("MaxConcurrency = %d, want 2", cfg.MaxConcurrency)
|
||||
}
|
||||
if cfg.MaxExpressionLen != 4000 || cfg.MaxPreambleLen != 4000 {
|
||||
t.Errorf("unexpected length caps: %d/%d", cfg.MaxExpressionLen, cfg.MaxPreambleLen)
|
||||
}
|
||||
if cfg.HealthAddr != ":8080" {
|
||||
t.Errorf("HealthAddr = %q, want :8080", cfg.HealthAddr)
|
||||
}
|
||||
if cfg.MongoDatabase != "latexgram" {
|
||||
t.Errorf("MongoDatabase = %q, want latexgram", cfg.MongoDatabase)
|
||||
}
|
||||
if !cfg.RichTextEnabled {
|
||||
t.Error("RichTextEnabled = false, want true by default")
|
||||
}
|
||||
if !cfg.RichDefaultMath {
|
||||
t.Error("RichDefaultMath = false, want true by default")
|
||||
}
|
||||
if cfg.S3PresignTTL != time.Hour {
|
||||
t.Errorf("S3PresignTTL = %s, want 1h", cfg.S3PresignTTL)
|
||||
}
|
||||
if cfg.S3Enabled() {
|
||||
t.Error("S3Enabled = true without S3 configuration")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadS3(t *testing.T) {
|
||||
setRequired(t)
|
||||
t.Setenv("S3_ENDPOINT_URL", "http://minio:9000/")
|
||||
t.Setenv("S3_BUCKET", "latexgrambot")
|
||||
t.Setenv("AWS_ACCESS_KEY_ID", "access")
|
||||
t.Setenv("AWS_SECRET_ACCESS_KEY", "secret")
|
||||
t.Setenv("S3_PUBLIC_ENDPOINT_URL", "https://s3.example/")
|
||||
t.Setenv("S3_PRESIGN_TTL", "45m")
|
||||
cfg, err := Load()
|
||||
if err != nil {
|
||||
t.Fatalf("Load: %v", err)
|
||||
}
|
||||
if !cfg.S3Enabled() || cfg.S3Endpoint != "http://minio:9000" || cfg.S3PublicEndpoint != "https://s3.example" {
|
||||
t.Fatalf("unexpected S3 config: %+v", cfg)
|
||||
}
|
||||
if cfg.S3PresignTTL != 45*time.Minute {
|
||||
t.Errorf("S3PresignTTL = %s", cfg.S3PresignTTL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadRejectsIncompleteS3(t *testing.T) {
|
||||
setRequired(t)
|
||||
t.Setenv("S3_BUCKET", "latexgrambot")
|
||||
if _, err := Load(); err == nil || !strings.Contains(err.Error(), "incomplete S3 configuration") {
|
||||
t.Fatalf("expected incomplete S3 error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadRichDefaultMathToggle(t *testing.T) {
|
||||
setRequired(t)
|
||||
t.Setenv("RICH_DEFAULT_MATH", "false")
|
||||
cfg, err := Load()
|
||||
if err != nil {
|
||||
t.Fatalf("Load: %v", err)
|
||||
}
|
||||
if cfg.RichDefaultMath {
|
||||
t.Error("RichDefaultMath = true, want false")
|
||||
}
|
||||
|
||||
t.Setenv("RICH_DEFAULT_MATH", "banana")
|
||||
if _, err := Load(); err == nil || !strings.Contains(err.Error(), "RICH_DEFAULT_MATH") {
|
||||
t.Fatalf("expected RICH_DEFAULT_MATH validation error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadRichTextToggle(t *testing.T) {
|
||||
setRequired(t)
|
||||
t.Setenv("RICH_TEXT_ENABLED", "false")
|
||||
cfg, err := Load()
|
||||
if err != nil {
|
||||
t.Fatalf("Load: %v", err)
|
||||
}
|
||||
if cfg.RichTextEnabled {
|
||||
t.Error("RichTextEnabled = true, want false")
|
||||
}
|
||||
|
||||
t.Setenv("RICH_TEXT_ENABLED", "banana")
|
||||
if _, err := Load(); err == nil || !strings.Contains(err.Error(), "RICH_TEXT_ENABLED") {
|
||||
t.Fatalf("expected RICH_TEXT_ENABLED validation error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadMissingToken(t *testing.T) {
|
||||
t.Setenv("MONGODB_URI", "mongodb://example:27017")
|
||||
_, err := Load()
|
||||
if err == nil || !strings.Contains(err.Error(), "TELEGRAM_BOT_TOKEN") {
|
||||
t.Fatalf("expected missing token error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadMissingMongoURI(t *testing.T) {
|
||||
t.Setenv("TELEGRAM_BOT_TOKEN", "123:abc")
|
||||
_, err := Load()
|
||||
if err == nil || !strings.Contains(err.Error(), "MONGODB_URI") {
|
||||
t.Fatalf("expected missing MONGODB_URI error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadInvalidValues(t *testing.T) {
|
||||
setRequired(t)
|
||||
t.Setenv("RENDER_DPI", "10")
|
||||
if _, err := Load(); err == nil || !strings.Contains(err.Error(), "RENDER_DPI") {
|
||||
t.Fatalf("expected RENDER_DPI validation error, got %v", err)
|
||||
}
|
||||
|
||||
t.Setenv("RENDER_DPI", "300")
|
||||
t.Setenv("RENDER_TIMEOUT", "banana")
|
||||
if _, err := Load(); err == nil || !strings.Contains(err.Error(), "RENDER_TIMEOUT") {
|
||||
t.Fatalf("expected RENDER_TIMEOUT validation error, got %v", err)
|
||||
}
|
||||
}
|
||||
145
internal/files/files.go
Normal file
145
internal/files/files.go
Normal file
@@ -0,0 +1,145 @@
|
||||
// Package files stores rendered outputs in a private S3-compatible bucket
|
||||
// and issues short-lived presigned URLs for Telegram inline-query results.
|
||||
package files
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
awsconfig "github.com/aws/aws-sdk-go-v2/config"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// Config points the store at the internal S3-compatible API and the public
|
||||
// API endpoint used only when signing download URLs.
|
||||
type Config struct {
|
||||
Endpoint string
|
||||
Region string
|
||||
Bucket string
|
||||
AccessKeyID string
|
||||
SecretAccessKey string
|
||||
Prefix string
|
||||
PublicEndpoint string
|
||||
PresignTTL time.Duration
|
||||
HTTPClient aws.HTTPClient
|
||||
}
|
||||
|
||||
// Store keeps content-addressed rendered assets in S3.
|
||||
type Store struct {
|
||||
client *s3.Client
|
||||
presigner *s3.PresignClient
|
||||
bucket string
|
||||
prefix string
|
||||
presignTTL time.Duration
|
||||
}
|
||||
|
||||
// New builds an S3-compatible object store.
|
||||
func New(ctx context.Context, cfg Config) (*Store, error) {
|
||||
awsCfg, err := awsconfig.LoadDefaultConfig(ctx,
|
||||
awsconfig.WithRegion(cfg.Region),
|
||||
awsconfig.WithCredentialsProvider(
|
||||
credentials.NewStaticCredentialsProvider(cfg.AccessKeyID, cfg.SecretAccessKey, ""),
|
||||
),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("load aws config: %w", err)
|
||||
}
|
||||
if cfg.HTTPClient != nil {
|
||||
awsCfg.HTTPClient = cfg.HTTPClient
|
||||
}
|
||||
client := s3.NewFromConfig(awsCfg, func(options *s3.Options) {
|
||||
options.BaseEndpoint = aws.String(cfg.Endpoint)
|
||||
options.UsePathStyle = true
|
||||
})
|
||||
publicClient := s3.NewFromConfig(awsCfg, func(options *s3.Options) {
|
||||
options.BaseEndpoint = aws.String(cfg.PublicEndpoint)
|
||||
options.UsePathStyle = true
|
||||
})
|
||||
return &Store{
|
||||
client: client,
|
||||
presigner: s3.NewPresignClient(publicClient),
|
||||
bucket: cfg.Bucket,
|
||||
prefix: strings.Trim(cfg.Prefix, "/"),
|
||||
presignTTL: cfg.PresignTTL,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Put uploads the data under a fresh UUIDv4 object name and returns a
|
||||
// short-lived presigned GET URL.
|
||||
func (s *Store) Put(ctx context.Context, ext, contentType string, data []byte) (string, error) {
|
||||
key := s.key(uuid.NewString(), ext)
|
||||
_, err := s.client.PutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: aws.String(s.bucket),
|
||||
Key: aws.String(key),
|
||||
Body: bytes.NewReader(data),
|
||||
ContentType: aws.String(contentType),
|
||||
CacheControl: aws.String("private, max-age=86400"),
|
||||
})
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("put %s object: %w", ext, err)
|
||||
}
|
||||
|
||||
request, err := s.presigner.PresignGetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: aws.String(s.bucket),
|
||||
Key: aws.String(key),
|
||||
}, func(options *s3.PresignOptions) {
|
||||
options.Expires = s.presignTTL
|
||||
})
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("presign %s object: %w", ext, err)
|
||||
}
|
||||
return request.URL, nil
|
||||
}
|
||||
|
||||
// DeleteExpired removes objects older than age and returns how many were
|
||||
// deleted. Telegram downloads the file when an inline result is picked and
|
||||
// the presigned link expires after the presign TTL; objects past their
|
||||
// lifetime can no longer be fetched, so they are garbage collected.
|
||||
func (s *Store) DeleteExpired(ctx context.Context, age time.Duration) (int, error) {
|
||||
cutoff := time.Now().Add(-age)
|
||||
prefix := ""
|
||||
if s.prefix != "" {
|
||||
prefix = s.prefix + "/"
|
||||
}
|
||||
deleted := 0
|
||||
var token *string
|
||||
for {
|
||||
list, err := s.client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
|
||||
Bucket: aws.String(s.bucket),
|
||||
Prefix: aws.String(prefix),
|
||||
ContinuationToken: token,
|
||||
})
|
||||
if err != nil {
|
||||
return deleted, fmt.Errorf("list objects: %w", err)
|
||||
}
|
||||
for _, object := range list.Contents {
|
||||
if object.LastModified == nil || !object.LastModified.Before(cutoff) {
|
||||
continue
|
||||
}
|
||||
if _, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
|
||||
Bucket: aws.String(s.bucket),
|
||||
Key: object.Key,
|
||||
}); err != nil {
|
||||
return deleted, fmt.Errorf("delete %s: %w", aws.ToString(object.Key), err)
|
||||
}
|
||||
deleted++
|
||||
}
|
||||
if list.IsTruncated == nil || !*list.IsTruncated {
|
||||
return deleted, nil
|
||||
}
|
||||
token = list.NextContinuationToken
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Store) key(name, ext string) string {
|
||||
if s.prefix == "" {
|
||||
return name + "." + ext
|
||||
}
|
||||
return s.prefix + "/" + name + "." + ext
|
||||
}
|
||||
176
internal/files/files_test.go
Normal file
176
internal/files/files_test.go
Normal file
@@ -0,0 +1,176 @@
|
||||
package files
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func TestPutUploadsUniqueUUIDObjectsWithPresignedURLs(t *testing.T) {
|
||||
var puts atomic.Int32
|
||||
keys := make(chan string, 4)
|
||||
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
||||
status := http.StatusOK
|
||||
switch request.Method {
|
||||
case http.MethodPut:
|
||||
puts.Add(1)
|
||||
keys <- request.URL.Path
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: status,
|
||||
Status: http.StatusText(status),
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader("")),
|
||||
Request: request,
|
||||
}, nil
|
||||
})}
|
||||
|
||||
store, err := New(context.Background(), Config{
|
||||
Endpoint: "https://s3.internal",
|
||||
PublicEndpoint: "https://s3.example",
|
||||
Region: "us-east-1",
|
||||
Bucket: "latexgrambot",
|
||||
AccessKeyID: "access",
|
||||
SecretAccessKey: "secret",
|
||||
Prefix: "content",
|
||||
PresignTTL: time.Hour,
|
||||
HTTPClient: client,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
seen := map[string]bool{}
|
||||
for range 2 {
|
||||
signed, err := store.Put(context.Background(), "jpg", "image/jpeg", []byte("image"))
|
||||
if err != nil {
|
||||
t.Fatalf("Put: %v", err)
|
||||
}
|
||||
parsed, err := url.Parse(signed)
|
||||
if err != nil {
|
||||
t.Fatalf("parse signed URL: %v", err)
|
||||
}
|
||||
base := strings.TrimSuffix(strings.TrimPrefix(parsed.Path, "/latexgrambot/content/"), ".jpg")
|
||||
if _, err := uuid.Parse(base); err != nil || uuid.MustParse(base).Version() != 4 {
|
||||
t.Errorf("object name %q is not a UUIDv4", base)
|
||||
}
|
||||
if seen[parsed.Path] {
|
||||
t.Errorf("object name %q was reused", parsed.Path)
|
||||
}
|
||||
seen[parsed.Path] = true
|
||||
if parsed.Query().Get("X-Amz-Signature") == "" || parsed.Query().Get("X-Amz-Expires") != "3600" {
|
||||
t.Errorf("URL is not signed for one hour: %s", signed)
|
||||
}
|
||||
}
|
||||
if got := puts.Load(); got != 2 {
|
||||
t.Fatalf("PUT requests = %d, want 2 (one object per request)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteExpired(t *testing.T) {
|
||||
var deleted []string
|
||||
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
||||
header := make(http.Header)
|
||||
body := ""
|
||||
status := http.StatusOK
|
||||
switch request.Method {
|
||||
case http.MethodGet: // ListObjectsV2
|
||||
header.Set("Content-Type", "application/xml")
|
||||
body = `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
|
||||
<Name>latexgrambot</Name><Prefix>content/</Prefix><KeyCount>2</KeyCount><MaxKeys>1000</MaxKeys><IsTruncated>false</IsTruncated>
|
||||
<Contents><Key>content/old.jpg</Key><LastModified>2000-01-01T00:00:00.000Z</LastModified><ETag>"a"</ETag><Size>1</Size><StorageClass>STANDARD</StorageClass></Contents>
|
||||
<Contents><Key>content/fresh.pdf</Key><LastModified>2099-01-01T00:00:00.000Z</LastModified><ETag>"b"</ETag><Size>1</Size><StorageClass>STANDARD</StorageClass></Contents>
|
||||
</ListBucketResult>`
|
||||
case http.MethodDelete:
|
||||
deleted = append(deleted, request.URL.Path)
|
||||
status = http.StatusNoContent
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: status,
|
||||
Status: http.StatusText(status),
|
||||
Header: header,
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
Request: request,
|
||||
}, nil
|
||||
})}
|
||||
|
||||
store, err := New(context.Background(), Config{
|
||||
Endpoint: "https://s3.internal",
|
||||
PublicEndpoint: "https://s3.example",
|
||||
Region: "us-east-1",
|
||||
Bucket: "latexgrambot",
|
||||
AccessKeyID: "access",
|
||||
SecretAccessKey: "secret",
|
||||
Prefix: "content",
|
||||
PresignTTL: time.Hour,
|
||||
HTTPClient: client,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
removed, err := store.DeleteExpired(context.Background(), time.Hour)
|
||||
if err != nil {
|
||||
t.Fatalf("DeleteExpired: %v", err)
|
||||
}
|
||||
if removed != 1 {
|
||||
t.Fatalf("removed = %d, want 1", removed)
|
||||
}
|
||||
if len(deleted) != 1 || !strings.HasSuffix(deleted[0], "/content/old.jpg") {
|
||||
t.Fatalf("deleted = %v, want the expired object only", deleted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestS3Integration(t *testing.T) {
|
||||
endpoint := os.Getenv("LATEXGRAMBOT_S3_TEST_ENDPOINT")
|
||||
accessKey := os.Getenv("LATEXGRAMBOT_S3_TEST_ACCESS_KEY")
|
||||
secretKey := os.Getenv("LATEXGRAMBOT_S3_TEST_SECRET_KEY")
|
||||
if endpoint == "" || accessKey == "" || secretKey == "" {
|
||||
t.Skip("S3 integration environment is not configured")
|
||||
}
|
||||
store, err := New(context.Background(), Config{
|
||||
Endpoint: endpoint,
|
||||
PublicEndpoint: endpoint,
|
||||
Region: "us-east-1",
|
||||
Bucket: "latexgrambot",
|
||||
AccessKeyID: accessKey,
|
||||
SecretAccessKey: secretKey,
|
||||
Prefix: "integration",
|
||||
PresignTTL: 5 * time.Minute,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
const payload = "latexgrambot S3 integration probe\n"
|
||||
signed, err := store.Put(context.Background(), "txt", "text/plain", []byte(payload))
|
||||
if err != nil {
|
||||
t.Fatalf("Put: %v", err)
|
||||
}
|
||||
response, err := http.Get(signed)
|
||||
if err != nil {
|
||||
t.Fatalf("GET presigned URL: %v", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
body, err := io.ReadAll(response.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("read presigned response: %v", err)
|
||||
}
|
||||
if response.StatusCode != http.StatusOK || string(body) != payload {
|
||||
t.Fatalf("presigned GET = %s, body %q", response.Status, body)
|
||||
}
|
||||
}
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
return f(request)
|
||||
}
|
||||
446
internal/render/render.go
Normal file
446
internal/render/render.go
Normal file
@@ -0,0 +1,446 @@
|
||||
// Package render turns LaTeX expressions into PNG and PDF files using a
|
||||
// locally installed TeX Live (pdflatex) and poppler (pdftoppm). No network
|
||||
// service is involved.
|
||||
package render
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"image"
|
||||
_ "image/jpeg"
|
||||
_ "image/png"
|
||||
"math"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"latexgrambot/internal/rich"
|
||||
|
||||
"go.opentelemetry.io/otel"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/codes"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
var renderTracer = otel.Tracer("latexgrambot/render")
|
||||
|
||||
const (
|
||||
// minRenderedSide is the minimum size of the smallest PNG side. Small
|
||||
// formulas rasterize to a few hundred pixels at any DPI; Telegram then
|
||||
// upscales and JPEG-recompresses them, blurring thin glyphs. Upscaling
|
||||
// the raster beforehand keeps the final picture crisp.
|
||||
minRenderedSide = 2560
|
||||
|
||||
// maxRenderedSide guards against pathologically tall formulas.
|
||||
maxRenderedSide = 4096
|
||||
)
|
||||
|
||||
// documentClass is the fixed preview layout: the standalone class crops
|
||||
// tightly around the content.
|
||||
const documentClass = `\documentclass[preview,varwidth,border=4pt]{standalone}`
|
||||
|
||||
// defaultPackages is inserted into the document when the user has no custom
|
||||
// preamble: the broad standard package set so documents compile without a
|
||||
// custom preamble.
|
||||
const defaultPackages = `\usepackage[utf8]{inputenc}
|
||||
\usepackage[T1]{fontenc}
|
||||
\usepackage{lmodern}
|
||||
\usepackage{textcomp}
|
||||
\usepackage{amsmath}
|
||||
\usepackage{amssymb}
|
||||
\usepackage{amsfonts}
|
||||
\usepackage{amsthm}
|
||||
\usepackage{amscd}
|
||||
\usepackage{mathtools}
|
||||
\usepackage{mathrsfs}
|
||||
\usepackage{bm}
|
||||
\usepackage{cancel}
|
||||
\usepackage{braket}
|
||||
\usepackage{siunitx}
|
||||
\usepackage{mhchem}
|
||||
\usepackage{chemfig}
|
||||
\usepackage{systeme}
|
||||
\usepackage{polynom}
|
||||
\usepackage{wasysym}
|
||||
\usepackage{stmaryrd}
|
||||
\usepackage{pifont}
|
||||
\usepackage{eurosym}
|
||||
\usepackage{graphicx}
|
||||
\usepackage{xcolor}
|
||||
\usepackage{colortbl}
|
||||
\usepackage{tikz}
|
||||
\usepackage{pgfplots}
|
||||
\pgfplotsset{compat=1.18}
|
||||
\usepackage{tikz-cd}
|
||||
\usepackage{tcolorbox}
|
||||
\usepackage{qrcode}
|
||||
\usepackage{float}
|
||||
\usepackage{array}
|
||||
\usepackage{multirow}
|
||||
\usepackage{makecell}
|
||||
\usepackage{booktabs}
|
||||
\usepackage{tabularx}
|
||||
\usepackage{longtable}
|
||||
\usepackage{multicol}
|
||||
\usepackage{enumitem}
|
||||
\usepackage{caption}
|
||||
\usepackage{subcaption}
|
||||
\usepackage{wrapfig}
|
||||
\usepackage{rotating}
|
||||
\usepackage[normalem]{ulem}
|
||||
\usepackage{soul}
|
||||
\usepackage{listings}
|
||||
\usepackage{algorithm}
|
||||
\usepackage{algpseudocode}
|
||||
\usepackage{microtype}
|
||||
\usepackage{xspace}
|
||||
\usepackage{etoolbox}
|
||||
\usepackage{hyperref}`
|
||||
|
||||
// DefaultPreamble returns the preamble file used at image build time to
|
||||
// preload the default class and packages into the "latexgrambot" TeX
|
||||
// format, so renders no longer pay the package-load cost.
|
||||
func DefaultPreamble() string {
|
||||
return documentClass + "\n" + defaultPackages + "\n\\endofdump\n"
|
||||
}
|
||||
|
||||
// Result carries the compiled output. JPEG is used for URL-based Telegram
|
||||
// inline photos, while PNG is used for direct chat replies.
|
||||
type Result struct {
|
||||
PNG []byte
|
||||
JPEG []byte
|
||||
PDF []byte
|
||||
}
|
||||
|
||||
// LatexError wraps a LaTeX compile failure and carries the log excerpt shown
|
||||
// to the user.
|
||||
type LatexError struct {
|
||||
Excerpt string
|
||||
}
|
||||
|
||||
func (e *LatexError) Error() string { return "latex compilation failed" }
|
||||
|
||||
// Renderer compiles LaTeX documents locally.
|
||||
type Renderer struct {
|
||||
Pdflatex string
|
||||
Pdftoppm string
|
||||
Timeout time.Duration
|
||||
WorkDir string // base directory for temp dirs; os.TempDir() when empty
|
||||
|
||||
// Format is the preloaded TeX format (built from DefaultPreamble) used
|
||||
// when the default preamble is in effect. Empty disables it.
|
||||
Format string
|
||||
|
||||
sem chan struct{}
|
||||
}
|
||||
|
||||
// New creates a Renderer that allows at most maxConcurrent simultaneous
|
||||
// pdflatex runs.
|
||||
func New(pdflatex, pdftoppm string, timeout time.Duration, maxConcurrent int, workDir string) *Renderer {
|
||||
if maxConcurrent < 1 {
|
||||
maxConcurrent = 1
|
||||
}
|
||||
return &Renderer{
|
||||
Pdflatex: pdflatex,
|
||||
Pdftoppm: pdftoppm,
|
||||
Timeout: timeout,
|
||||
WorkDir: workDir,
|
||||
sem: make(chan struct{}, maxConcurrent),
|
||||
}
|
||||
}
|
||||
|
||||
// dangerousCommands are TeX primitives that read or write files; user input
|
||||
// must never reach them.
|
||||
var dangerousCommands = regexp.MustCompile(`\\(input|include|openin|openout|write|read|catcode|newwrite|immediate)\b`)
|
||||
|
||||
// normalizeExpression validates and trims user LaTeX without changing its
|
||||
// mode. The bot has already applied the RICH_DEFAULT_MATH policy before it
|
||||
// calls the renderer; wrapping everything here would incorrectly place text
|
||||
// commands such as \hyperlink and \section inside display math.
|
||||
func normalizeExpression(expression string) (string, error) {
|
||||
expr := strings.TrimSpace(expression)
|
||||
if expr == "" {
|
||||
return "", &LatexError{Excerpt: "the expression is empty"}
|
||||
}
|
||||
if dangerousCommands.MatchString(expr) {
|
||||
return "", &LatexError{Excerpt: "unsupported TeX command in the expression"}
|
||||
}
|
||||
return expr, nil
|
||||
}
|
||||
|
||||
// stripOuterMathDelimiters removes matching outer math delimiters, so a
|
||||
// forced math retry never nests delimiters.
|
||||
func stripOuterMathDelimiters(expr string) string {
|
||||
for _, pair := range [][2]string{{`$$`, `$$`}, {`$`, `$`}, {`\[`, `\]`}, {`\(`, `\)`}} {
|
||||
if len(expr) >= len(pair[0])+len(pair[1]) &&
|
||||
strings.HasPrefix(expr, pair[0]) && strings.HasSuffix(expr, pair[1]) {
|
||||
return strings.TrimSpace(expr[len(pair[0]) : len(expr)-len(pair[1])])
|
||||
}
|
||||
}
|
||||
return expr
|
||||
}
|
||||
|
||||
// Formats selects which outputs a render produces.
|
||||
type Formats uint8
|
||||
|
||||
const (
|
||||
FormatPNG Formats = 1 << iota
|
||||
FormatJPEG
|
||||
FormatPDF
|
||||
)
|
||||
|
||||
// Render compiles the expression and produces every format.
|
||||
func (r *Renderer) Render(ctx context.Context, preamble, expression string, dpi int) (*Result, error) {
|
||||
return r.RenderFormats(ctx, preamble, expression, dpi, FormatPNG|FormatJPEG|FormatPDF)
|
||||
}
|
||||
|
||||
// RenderFormats compiles the expression inside the given preamble and
|
||||
// produces only the requested formats, so callers never pay for
|
||||
// rasterization they do not use. The expression is normalised (math
|
||||
// delimiters stripped, plain expressions wrapped in math mode) and inserted
|
||||
// between \begin{document} and \end{document}.
|
||||
func (r *Renderer) RenderFormats(ctx context.Context, preamble, expression string, dpi int, formats Formats) (result *Result, err error) {
|
||||
normalized, err := normalizeExpression(expression)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ctx, span := renderTracer.Start(ctx, "render",
|
||||
trace.WithAttributes(
|
||||
attribute.Int("dpi", dpi),
|
||||
attribute.Int("expression.length", len(expression)),
|
||||
attribute.Int("expression.normalized_length", len(normalized)),
|
||||
attribute.Int("preamble.length", len(preamble)),
|
||||
),
|
||||
)
|
||||
defer span.End()
|
||||
defer func() {
|
||||
if err != nil {
|
||||
span.RecordError(err)
|
||||
span.SetStatus(codes.Error, "render failed")
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case r.sem <- struct{}{}:
|
||||
defer func() { <-r.sem }()
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
workdir, err := os.MkdirTemp(r.WorkDir, "latexgrambot-")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create workdir: %w", err)
|
||||
}
|
||||
defer os.RemoveAll(workdir)
|
||||
|
||||
texPath := filepath.Join(workdir, "document.tex")
|
||||
// Preloaded format: with the default preamble the class and packages
|
||||
// are already in the format, so only the body is compiled.
|
||||
useFormat := r.Format != "" && strings.TrimSpace(preamble) == ""
|
||||
compile := func(body string) error {
|
||||
document := buildDocument(preamble, body)
|
||||
if useFormat {
|
||||
document = "\\begin{document}\n" + body + "\n\\end{document}\n"
|
||||
}
|
||||
if err := os.WriteFile(texPath, []byte(document), 0o600); err != nil {
|
||||
return fmt.Errorf("write document: %w", err)
|
||||
}
|
||||
// -no-shell-escape is the default, spelled out here on purpose: the
|
||||
// \write18 escapes must never be enabled for user-supplied LaTeX.
|
||||
args := []string{
|
||||
"-interaction=nonstopmode", "-halt-on-error", "-no-shell-escape",
|
||||
"-output-directory=" + workdir,
|
||||
}
|
||||
if useFormat {
|
||||
args = append(args, "-fmt="+r.Format)
|
||||
}
|
||||
return r.run(ctx, workdir, r.Pdflatex, append(args, "document.tex")...)
|
||||
}
|
||||
|
||||
// Entirely-math expressions are put into math mode up front: waiting
|
||||
// for the "Missing $ inserted" error would compile every formula twice.
|
||||
firstAttempt := normalized
|
||||
if rich.IsMathExpression(expression) && !rich.ContainsMathEnvironment(expression) {
|
||||
firstAttempt = `\[` + stripOuterMathDelimiters(strings.TrimSpace(normalized)) + `\]`
|
||||
}
|
||||
|
||||
if err := compile(firstAttempt); err != nil {
|
||||
logContent, readErr := os.ReadFile(filepath.Join(workdir, "document.log"))
|
||||
if readErr != nil {
|
||||
return nil, err
|
||||
}
|
||||
// One retry for cases the up-front check missed: when the compiler
|
||||
// only complained about missing math delimiters and the whole
|
||||
// expression is math. If the retry fails too, the original error is
|
||||
// reported so the user decides what to do.
|
||||
if strings.Contains(string(logContent), "Missing $ inserted") &&
|
||||
rich.IsMathExpression(expression) && firstAttempt == normalized {
|
||||
forced := `\[` + stripOuterMathDelimiters(strings.TrimSpace(expression)) + `\]`
|
||||
for _, name := range []string{"document.aux", "document.log", "document.pdf"} {
|
||||
_ = os.Remove(filepath.Join(workdir, name))
|
||||
}
|
||||
if retryErr := compile(forced); retryErr != nil {
|
||||
return nil, &LatexError{Excerpt: extractError(string(logContent))}
|
||||
}
|
||||
} else {
|
||||
return nil, &LatexError{Excerpt: extractError(string(logContent))}
|
||||
}
|
||||
}
|
||||
|
||||
pdfPath := filepath.Join(workdir, "document.pdf")
|
||||
result = &Result{}
|
||||
if formats&FormatPDF != 0 {
|
||||
result.PDF, err = os.ReadFile(pdfPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read compiled pdf: %w", err)
|
||||
}
|
||||
}
|
||||
if formats&FormatPNG != 0 {
|
||||
pngPath := filepath.Join(workdir, "document.png")
|
||||
if err := r.renderRaster(ctx, workdir, dpi, pdfPath, pngPath, "png"); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.PNG, err = os.ReadFile(pngPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read rendered png: %w", err)
|
||||
}
|
||||
}
|
||||
if formats&FormatJPEG != 0 {
|
||||
jpegPath := filepath.Join(workdir, "document.jpg")
|
||||
if err := r.renderRaster(ctx, workdir, dpi, pdfPath, jpegPath, "jpeg"); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.JPEG, err = os.ReadFile(jpegPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read rendered jpeg: %w", err)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// renderRaster rasterizes the first PDF page to a PNG or JPEG, re-rendering
|
||||
// it upscaled when the natural size is too small for a crisp Telegram
|
||||
// display. Rasterizing the JPEG directly avoids decoding and re-encoding a
|
||||
// full-size PNG in Go.
|
||||
func (r *Renderer) renderRaster(ctx context.Context, workdir string, dpi int, pdfPath, outPath, format string) error {
|
||||
base := filepath.Join(workdir, "document")
|
||||
args := []string{"-" + format, "-r", strconv.Itoa(dpi), "-singlefile", "-f", "1", "-l", "1"}
|
||||
if err := r.run(ctx, workdir, r.Pdftoppm, append(args, pdfPath, base)...); err != nil {
|
||||
return err
|
||||
}
|
||||
width, height, err := imageSize(outPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("inspect rendered %s: %w", format, err)
|
||||
}
|
||||
small := min(width, height)
|
||||
large := max(width, height)
|
||||
scale := 1.0
|
||||
if small > 0 && small < minRenderedSide {
|
||||
scale = float64(minRenderedSide) / float64(small)
|
||||
if scale*float64(large) > maxRenderedSide {
|
||||
scale = float64(maxRenderedSide) / float64(large)
|
||||
}
|
||||
}
|
||||
if scale <= 1 {
|
||||
return nil
|
||||
}
|
||||
scaledWidth := int(math.Round(float64(width) * scale))
|
||||
args = append(args,
|
||||
"-scale-to-x", strconv.Itoa(scaledWidth), "-scale-to-y", "-1")
|
||||
return r.run(ctx, workdir, r.Pdftoppm, append(args, pdfPath, base)...)
|
||||
}
|
||||
|
||||
func imageSize(path string) (int, int, error) {
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
defer file.Close()
|
||||
config, _, err := image.DecodeConfig(file)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
return config.Width, config.Height, nil
|
||||
}
|
||||
|
||||
// buildDocument assembles the full .tex source. User preamble additions are
|
||||
// loaded after the broad default package set, so setting custom definitions
|
||||
// cannot accidentally remove core packages such as hyperref.
|
||||
func buildDocument(preamble, expression string) string {
|
||||
pkgs := defaultPackages
|
||||
if strings.TrimSpace(preamble) != "" {
|
||||
pkgs += "\n" + strings.TrimSpace(preamble)
|
||||
}
|
||||
return documentClass + "\n" +
|
||||
pkgs + "\n" +
|
||||
"\\begin{document}\n" +
|
||||
expression + "\n" +
|
||||
"\\end{document}\n"
|
||||
}
|
||||
|
||||
// extractError pulls the first "! ..." error line and its "l.NN ..." context
|
||||
// line from the pdflatex log.
|
||||
func extractError(log string) string {
|
||||
var first, context string
|
||||
for _, line := range strings.Split(log, "\n") {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if first == "" && strings.HasPrefix(trimmed, "!") {
|
||||
first = trimmed
|
||||
continue
|
||||
}
|
||||
if first != "" && strings.HasPrefix(trimmed, "l.") {
|
||||
context = trimmed
|
||||
break
|
||||
}
|
||||
}
|
||||
if first == "" {
|
||||
return "compilation failed without an explicit error message"
|
||||
}
|
||||
if context == "" {
|
||||
return first
|
||||
}
|
||||
return first + "\n" + context
|
||||
}
|
||||
|
||||
// run executes a command in dir, inheriting the timeout from ctx. Home,
|
||||
// TMPDIR and TEXMFOUTPUT point into the workdir so TeX never needs to write
|
||||
// outside it.
|
||||
func (r *Renderer) run(ctx context.Context, dir, bin string, args ...string) error {
|
||||
cmd := exec.CommandContext(ctx, bin, args...)
|
||||
cmd.Dir = dir
|
||||
cmd.Env = append(os.Environ(),
|
||||
"HOME="+dir,
|
||||
"TMPDIR="+dir,
|
||||
"TEXMFOUTPUT="+dir,
|
||||
)
|
||||
var output bytes.Buffer
|
||||
cmd.Stdout = &output
|
||||
cmd.Stderr = &output
|
||||
if err := cmd.Run(); err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return fmt.Errorf("%s: %w", filepath.Base(bin), ctx.Err())
|
||||
}
|
||||
var execErr *exec.Error
|
||||
if errors.As(err, &execErr) {
|
||||
return fmt.Errorf("%s is not available: %w", bin, execErr)
|
||||
}
|
||||
return fmt.Errorf("%s failed: %w (%s)", filepath.Base(bin), err, firstLine(output.String()))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// firstLine returns the first non-empty line of s, or an empty string.
|
||||
func firstLine(s string) string {
|
||||
for _, line := range strings.Split(s, "\n") {
|
||||
if strings.TrimSpace(line) != "" {
|
||||
return strings.TrimSpace(line)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
243
internal/render/render_test.go
Normal file
243
internal/render/render_test.go
Normal file
@@ -0,0 +1,243 @@
|
||||
package render
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"image"
|
||||
_ "image/png"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestBuildDocument(t *testing.T) {
|
||||
t.Run("default preamble", func(t *testing.T) {
|
||||
doc := buildDocument("", `$x^2$`)
|
||||
if !strings.Contains(doc, "\\documentclass[preview,varwidth,border=4pt]{standalone}") {
|
||||
t.Errorf("missing document class:\n%s", doc)
|
||||
}
|
||||
if !strings.Contains(doc, defaultPackages) {
|
||||
t.Errorf("missing default packages:\n%s", doc)
|
||||
}
|
||||
if !strings.Contains(doc, "\\begin{document}\n$x^2$\n\\end{document}") {
|
||||
t.Errorf("expression not placed in the document body:\n%s", doc)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("custom preamble extends packages", func(t *testing.T) {
|
||||
doc := buildDocument(`\usepackage{chemist}`, `H_2O`)
|
||||
if !strings.Contains(doc, defaultPackages) {
|
||||
t.Errorf("default packages missing:\n%s", doc)
|
||||
}
|
||||
if !strings.Contains(doc, "\\usepackage{chemist}") {
|
||||
t.Errorf("custom preamble missing:\n%s", doc)
|
||||
}
|
||||
if !strings.Contains(doc, "\\begin{document}\nH_2O\n\\end{document}") {
|
||||
t.Errorf("expression not placed in the document body:\n%s", doc)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestExtractError(t *testing.T) {
|
||||
log := strings.Join([]string{
|
||||
"This is pdfTeX, Version 3.141592653-2.6-1.40.26",
|
||||
"! Undefined control sequence.",
|
||||
"l.4 \\begin{document}",
|
||||
" \\badcommand",
|
||||
"",
|
||||
"! LaTeX Error: File `missing.sty' not found.",
|
||||
"l.5 \\usepackage{missing}",
|
||||
}, "\n")
|
||||
|
||||
excerpt := extractError(log)
|
||||
if !strings.Contains(excerpt, "! Undefined control sequence.") {
|
||||
t.Errorf("missing first error line in %q", excerpt)
|
||||
}
|
||||
if !strings.Contains(excerpt, "l.4 \\begin{document}") {
|
||||
t.Errorf("missing context line in %q", excerpt)
|
||||
}
|
||||
if strings.Contains(excerpt, "missing.sty") {
|
||||
t.Errorf("second error must not be reported with -halt-on-error behavior: %q", excerpt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractErrorWithoutErrorLine(t *testing.T) {
|
||||
excerpt := extractError("some unrelated log content")
|
||||
if !strings.Contains(excerpt, "without an explicit error") {
|
||||
t.Errorf("unexpected excerpt: %q", excerpt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderIntegration(t *testing.T) {
|
||||
if _, err := exec.LookPath("pdflatex"); err != nil {
|
||||
t.Skip("pdflatex not installed")
|
||||
}
|
||||
if _, err := exec.LookPath("pdftoppm"); err != nil {
|
||||
t.Skip("pdftoppm not installed")
|
||||
}
|
||||
|
||||
r := New("pdflatex", "pdftoppm", 30*time.Second, 2, "")
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
|
||||
result, err := r.Render(ctx, "", `$\int_0^\infty e^{-x^2}\,dx = \frac{\sqrt{\pi}}{2}$`, 150)
|
||||
if err != nil {
|
||||
t.Fatalf("render failed: %v", err)
|
||||
}
|
||||
if len(result.PNG) == 0 {
|
||||
t.Fatal("empty PNG")
|
||||
}
|
||||
if !strings.HasPrefix(string(result.PNG), "\x89PNG") {
|
||||
t.Fatal("output is not a PNG file")
|
||||
}
|
||||
width, height, err := pngSizeBytes(result.PNG)
|
||||
if err != nil {
|
||||
t.Fatalf("decode png size: %v", err)
|
||||
}
|
||||
if width < minRenderedSide && height < minRenderedSide {
|
||||
t.Errorf("png is %dx%d, want the smallest side at least %d", width, height, minRenderedSide)
|
||||
}
|
||||
if width > maxRenderedSide || height > maxRenderedSide {
|
||||
t.Errorf("png is %dx%d, want no side beyond %d", width, height, maxRenderedSide)
|
||||
}
|
||||
if len(result.PDF) == 0 || !strings.HasPrefix(string(result.PDF), "%PDF") {
|
||||
t.Fatal("output is not a PDF file")
|
||||
}
|
||||
if len(result.JPEG) == 0 || !strings.HasPrefix(string(result.JPEG), "\xff\xd8") {
|
||||
t.Fatal("output is not a JPEG file")
|
||||
}
|
||||
|
||||
linked, err := r.Render(ctx, "", `\hypertarget{details}{Details}\par\hyperlink{details}{Go to details}`, 150)
|
||||
if err != nil {
|
||||
t.Fatalf("hyperlink render failed: %v", err)
|
||||
}
|
||||
if len(linked.PDF) == 0 || len(linked.PNG) == 0 || len(linked.JPEG) == 0 {
|
||||
t.Fatal("hyperlink render returned an empty format")
|
||||
}
|
||||
}
|
||||
|
||||
func pngSizeBytes(data []byte) (int, int, error) {
|
||||
config, _, err := image.DecodeConfig(bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
return config.Width, config.Height, nil
|
||||
}
|
||||
|
||||
func TestRenderErrorIntegration(t *testing.T) {
|
||||
if _, err := exec.LookPath("pdflatex"); err != nil {
|
||||
t.Skip("pdflatex not installed")
|
||||
}
|
||||
|
||||
r := New("pdflatex", "pdftoppm", 30*time.Second, 2, "")
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
|
||||
_, err := r.Render(ctx, "", `\undefinedcommandxyz`, 150)
|
||||
latexErr, ok := err.(*LatexError)
|
||||
if !ok {
|
||||
t.Fatalf("expected *LatexError, got %T: %v", err, err)
|
||||
}
|
||||
if !strings.Contains(latexErr.Excerpt, "! Undefined control sequence") {
|
||||
t.Errorf("unexpected excerpt: %q", latexErr.Excerpt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderPreWrapsBareMath(t *testing.T) {
|
||||
if _, err := exec.LookPath("pdflatex"); err != nil {
|
||||
t.Skip("pdflatex not installed")
|
||||
}
|
||||
r := New("pdflatex", "pdftoppm", 30*time.Second, 2, "")
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// Bare math is wrapped up front, so a single compile succeeds.
|
||||
result, err := r.Render(ctx, "", `\int_0^1 x\,dx`, 150)
|
||||
if err != nil {
|
||||
t.Fatalf("render failed: %v", err)
|
||||
}
|
||||
if len(result.PNG) == 0 || len(result.PDF) == 0 {
|
||||
t.Fatal("empty output")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderRetriesMissingDollarForMath(t *testing.T) {
|
||||
if _, err := exec.LookPath("pdflatex"); err != nil {
|
||||
t.Skip("pdflatex not installed")
|
||||
}
|
||||
r := New("pdflatex", "pdftoppm", 30*time.Second, 2, "")
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// A math environment is not pre-wrapped; the compile fails once with
|
||||
// "Missing $ inserted" and the retry wraps the whole expression.
|
||||
result, err := r.Render(ctx, "", "\\begin{matrix}a&b\\\\c&d\\end{matrix} \\int_0^1 x\\,dx", 150)
|
||||
if err != nil {
|
||||
t.Fatalf("render failed: %v", err)
|
||||
}
|
||||
if len(result.PNG) == 0 || len(result.PDF) == 0 {
|
||||
t.Fatal("empty output")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderNoRetryForMixedInput(t *testing.T) {
|
||||
if _, err := exec.LookPath("pdflatex"); err != nil {
|
||||
t.Skip("pdflatex not installed")
|
||||
}
|
||||
r := New("pdflatex", "pdftoppm", 30*time.Second, 2, "")
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// Mixed text with a math-only command inside is not retried: the error
|
||||
// is reported to the user as-is.
|
||||
_, err := r.Render(ctx, "", "\\begin{itemize}\\item \\int\\end{itemize}", 150)
|
||||
latexErr, ok := err.(*LatexError)
|
||||
if !ok {
|
||||
t.Fatalf("expected *LatexError, got %T: %v", err, err)
|
||||
}
|
||||
if !strings.Contains(latexErr.Excerpt, "Missing $ inserted") {
|
||||
t.Errorf("unexpected excerpt: %q", latexErr.Excerpt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultPreamble(t *testing.T) {
|
||||
preamble := DefaultPreamble()
|
||||
for _, fragment := range []string{
|
||||
"\\documentclass[preview,varwidth,border=4pt]{standalone}",
|
||||
"\\usepackage{amsmath}",
|
||||
"\\usepackage{tikz}",
|
||||
"\\endofdump",
|
||||
} {
|
||||
if !strings.Contains(preamble, fragment) {
|
||||
t.Errorf("preamble missing %q", fragment)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeExpression(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
`\mathbb{ R }`: `\mathbb{ R }`,
|
||||
`$\mathbb{R}$`: `$\mathbb{R}$`,
|
||||
`$$x^2 + y^2$$`: `$$x^2 + y^2$$`,
|
||||
`\[ \frac{1}{2} \]`: `\[ \frac{1}{2} \]`,
|
||||
`\(a+b\)`: `\(a+b\)`,
|
||||
" x ": `x`,
|
||||
`\begin{align} a &= b \\ c &= d \end{align}`: `\begin{align} a &= b \\ c &= d \end{align}`,
|
||||
}
|
||||
for input, want := range cases {
|
||||
got, err := normalizeExpression(input)
|
||||
if err != nil {
|
||||
t.Fatalf("normalizeExpression(%q): %v", input, err)
|
||||
}
|
||||
if got != want {
|
||||
t.Fatalf("normalizeExpression(%q) = %q, want %q", input, got, want)
|
||||
}
|
||||
}
|
||||
if _, err := normalizeExpression(`\input{/etc/passwd}`); err == nil {
|
||||
t.Fatal("dangerous command was accepted")
|
||||
}
|
||||
if _, err := normalizeExpression(" "); err == nil {
|
||||
t.Fatal("empty expression was accepted")
|
||||
}
|
||||
}
|
||||
742
internal/rich/latex.go
Normal file
742
internal/rich/latex.go
Normal file
@@ -0,0 +1,742 @@
|
||||
package rich
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"unicode"
|
||||
)
|
||||
|
||||
// textCommands marks a document as text-mode LaTeX (as opposed to a bare
|
||||
// math expression) when any of these is present.
|
||||
var textCommands = map[string]bool{
|
||||
"textbf": true, "textit": true, "emph": true, "underline": true,
|
||||
"uline": true, "texttt": true, "textsuperscript": true,
|
||||
"textsubscript": true, "textrm": true, "textsf": true, "textnormal": true,
|
||||
"textup": true, "textmd": true, "textcolor": true,
|
||||
"href": true, "url": true, "hyperlink": true, "hypertarget": true,
|
||||
"section": true, "subsection": true,
|
||||
"subsubsection": true, "item": true, "maketitle": true, "title": true,
|
||||
"author": true, "documentclass": true, "par": true,
|
||||
}
|
||||
|
||||
// wrapperEnvironments are environments whose markers are dropped from text
|
||||
// mode; their content is rendered inline.
|
||||
var wrapperEnvironments = map[string]bool{
|
||||
"itemize": true, "enumerate": true, "center": true, "flushleft": true,
|
||||
"flushright": true, "quote": true, "quotation": true, "description": true,
|
||||
"small": true, "footnotesize": true, "large": true, "Large": true,
|
||||
"document": true,
|
||||
}
|
||||
|
||||
// documentClassEnvironments are math environments that keep a \begin{...}
|
||||
// source in math mode.
|
||||
var mathEnvironments = map[string]bool{
|
||||
"equation": true, "equation*": true, "align": true, "align*": true,
|
||||
"aligned": true, "gather": true, "gather*": true, "multline": true,
|
||||
"multline*": true, "eqnarray": true, "eqnarray*": true, "array": true,
|
||||
"matrix": true, "pmatrix": true, "bmatrix": true, "vmatrix": true,
|
||||
"Vmatrix": true, "Bmatrix": true, "cases": true, "split": true,
|
||||
"displaymath": true, "math": true,
|
||||
}
|
||||
|
||||
// symbolCommands maps exotic symbol commands that client math engines
|
||||
// commonly lack to Unicode. Standard LaTeX/AMS math commands (\int, \sum,
|
||||
// Greek letters, arrows, relations, ...) are deliberately NOT translated:
|
||||
// the client engine renders them natively, while Unicode replacements are
|
||||
// not drawn by every math font.
|
||||
var symbolCommands = map[string]string{
|
||||
"heartsuit": "♡", "diamondsuit": "♢", "spadesuit": "♠", "clubsuit": "♣",
|
||||
"checkmark": "✓", "textbullet": "•", "degree": "°", "celsius": "°C",
|
||||
"copyright": "©", "registered": "®", "circledR": "®", "pounds": "£",
|
||||
"euro": "€", "yen": "¥", "male": "♂", "female": "♀",
|
||||
"flat": "♭", "natural": "♮", "sharp": "♯",
|
||||
"Box": "□", "square": "□", "blacksquare": "■",
|
||||
"triangle": "△", "blacktriangle": "▲", "blacktriangledown": "▼",
|
||||
"bigstar": "★", "lozenge": "◊", "blacklozenge": "⧫",
|
||||
"measuredangle": "∡", "sphericalangle": "∢",
|
||||
"therefore": "∴", "because": "∵", "maltese": "✠",
|
||||
"diagup": "╱", "diagdown": "╲", "complement": "∁", "eth": "ð",
|
||||
"Finv": "Ⅎ", "Game": "⅁", "circledS": "Ⓢ", "mho": "℧",
|
||||
}
|
||||
|
||||
// translateSymbols replaces known symbol commands with Unicode. Font
|
||||
// commands such as \mathbb are left untouched: client math engines render
|
||||
// them themselves and their Unicode math-alphanumeric fallbacks are not
|
||||
// supported everywhere.
|
||||
func translateSymbols(source string) string {
|
||||
if !strings.Contains(source, `\`) {
|
||||
return source
|
||||
}
|
||||
var out strings.Builder
|
||||
for index := 0; index < len(source); {
|
||||
if source[index] != '\\' {
|
||||
out.WriteByte(source[index])
|
||||
index++
|
||||
continue
|
||||
}
|
||||
name, next := readCommand(source, index)
|
||||
if name == "" {
|
||||
out.WriteString(source[index:next])
|
||||
index = next
|
||||
continue
|
||||
}
|
||||
if replacement, ok := symbolCommands[name]; ok {
|
||||
out.WriteString(replacement)
|
||||
index = skipEmptyGroup(source, next)
|
||||
continue
|
||||
}
|
||||
out.WriteString(source[index:next])
|
||||
index = next
|
||||
}
|
||||
return out.String()
|
||||
}
|
||||
|
||||
// Document converts a LaTeX source into a rich message. Math segments
|
||||
// ($...$, $$...$$, \(...\), \[...\]) become math nodes; text paragraphs keep
|
||||
// their line structure and map the common text-mode formatting commands to
|
||||
// rich text nodes. With defaultMath, a source with no math delimiters and
|
||||
// no known text-mode command is treated as a single math expression;
|
||||
// otherwise such a source is parsed as text and formulas need delimiters.
|
||||
func Document(source string, defaultMath bool) *Message {
|
||||
source = documentBody(source)
|
||||
source = translateSymbols(source)
|
||||
trimmed := strings.TrimSpace(source)
|
||||
if trimmed == "" {
|
||||
return &Message{Blocks: []Block{{Type: "paragraph", Text: " "}}}
|
||||
}
|
||||
if defaultMath && isPureMath(trimmed) {
|
||||
if isSymbolsOnly(trimmed) {
|
||||
// Client math fonts often lack exotic symbols; the regular
|
||||
// text font renders them everywhere.
|
||||
return &Message{Blocks: []Block{{Type: "paragraph", Text: trimmed}}}
|
||||
}
|
||||
return &Message{Blocks: []Block{{Type: "mathematical_expression", Expression: trimmed}}}
|
||||
}
|
||||
parser := &parser{src: source}
|
||||
blocks := parser.parseBlocks()
|
||||
if len(blocks) == 0 {
|
||||
blocks = []Block{{Type: "paragraph", Text: " "}}
|
||||
}
|
||||
return &Message{Blocks: blocks}
|
||||
}
|
||||
|
||||
// IsMathExpression reports whether the source is a math expression: bare
|
||||
// math commands with no text-mode command, a delimited math segment, or a
|
||||
// math environment. Used to decide whether a failed compile may be retried
|
||||
// with the expression forced into math mode.
|
||||
func IsMathExpression(source string) bool {
|
||||
trimmed := strings.TrimSpace(source)
|
||||
if trimmed == "" {
|
||||
return false
|
||||
}
|
||||
if stripOuterMathDelimiters(trimmed) != trimmed {
|
||||
return true
|
||||
}
|
||||
// Delimiters somewhere inside mean mixed text and math: not entirely
|
||||
// math, so no retry.
|
||||
if strings.Contains(trimmed, "$") || strings.Contains(trimmed, `\(`) || strings.Contains(trimmed, `\[`) {
|
||||
return false
|
||||
}
|
||||
for _, command := range commandsIn(trimmed) {
|
||||
if textCommands[command] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
for _, environment := range environmentsIn(trimmed) {
|
||||
if !mathEnvironments[environment] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// ContainsMathEnvironment reports whether the source contains a math
|
||||
// environment such as align or matrix; those compile standalone and must
|
||||
// not be wrapped again.
|
||||
func ContainsMathEnvironment(source string) bool {
|
||||
for _, environment := range environmentsIn(source) {
|
||||
if mathEnvironments[environment] {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// stripOuterMathDelimiters removes matching outer math delimiters.
|
||||
func stripOuterMathDelimiters(expr string) string {
|
||||
for _, pair := range [][2]string{{`$$`, `$$`}, {`$`, `$`}, {`\[`, `\]`}, {`\(`, `\)`}} {
|
||||
if len(expr) >= len(pair[0])+len(pair[1]) &&
|
||||
strings.HasPrefix(expr, pair[0]) && strings.HasSuffix(expr, pair[1]) {
|
||||
return strings.TrimSpace(expr[len(pair[0]) : len(expr)-len(pair[1])])
|
||||
}
|
||||
}
|
||||
return expr
|
||||
}
|
||||
|
||||
// isSymbolsOnly reports whether the source consists only of symbols and
|
||||
// punctuation: no letters, digits or LaTeX math structure.
|
||||
func isSymbolsOnly(source string) bool {
|
||||
for _, r := range source {
|
||||
switch r {
|
||||
case '\\', '_', '{', '}', '$', '&', '#', '^':
|
||||
return false
|
||||
}
|
||||
if r < 128 && !unicode.IsPunct(r) && !unicode.IsSpace(r) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// InlineDocument is Document prepared for inline answers: a leading math
|
||||
// expression renders inline one space after the bot mention instead of as a
|
||||
// display block, and a leading paragraph starts with one space.
|
||||
func InlineDocument(source string, defaultMath bool) *Message {
|
||||
message := Document(source, defaultMath)
|
||||
if len(message.Blocks) == 0 {
|
||||
return message
|
||||
}
|
||||
first := &message.Blocks[0]
|
||||
switch first.Type {
|
||||
case "mathematical_expression":
|
||||
*first = Block{Type: "paragraph", Text: []any{
|
||||
" ",
|
||||
Inline{Type: "mathematical_expression", Expression: first.Expression},
|
||||
}}
|
||||
case "paragraph":
|
||||
first.Text = prependSpace(first.Text)
|
||||
}
|
||||
return message
|
||||
}
|
||||
|
||||
func prependSpace(text any) any {
|
||||
switch value := text.(type) {
|
||||
case string:
|
||||
return " " + value
|
||||
case []any:
|
||||
return append([]any{" "}, value...)
|
||||
default:
|
||||
return text
|
||||
}
|
||||
}
|
||||
|
||||
// documentBody reduces a full document to the content of its document
|
||||
// environment.
|
||||
func documentBody(source string) string {
|
||||
const begin = `\begin{document}`
|
||||
const end = `\end{document}`
|
||||
index := strings.Index(source, begin)
|
||||
if index < 0 {
|
||||
return source
|
||||
}
|
||||
source = source[index+len(begin):]
|
||||
if index = strings.Index(source, end); index >= 0 {
|
||||
source = source[:index]
|
||||
}
|
||||
return source
|
||||
}
|
||||
|
||||
// isPureMath reports whether the source should be rendered as one math
|
||||
// expression: no math delimiters, no known text-mode command and no
|
||||
// non-math environment.
|
||||
func isPureMath(source string) bool {
|
||||
if strings.Contains(source, "$") || strings.Contains(source, `\(`) || strings.Contains(source, `\[`) {
|
||||
return false
|
||||
}
|
||||
for _, command := range commandsIn(source) {
|
||||
if textCommands[command] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
for _, env := range environmentsIn(source) {
|
||||
if !mathEnvironments[env] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func commandsIn(source string) []string {
|
||||
var commands []string
|
||||
for i := 0; i < len(source); i++ {
|
||||
if source[i] != '\\' || i+1 >= len(source) {
|
||||
continue
|
||||
}
|
||||
if source[i+1] == '\\' {
|
||||
commands = append(commands, `\\`)
|
||||
i++
|
||||
continue
|
||||
}
|
||||
name := ""
|
||||
next := i + 1
|
||||
for next < len(source) && isCommandLetter(source[next]) {
|
||||
name += string(source[next])
|
||||
next++
|
||||
}
|
||||
if name != "" {
|
||||
commands = append(commands, name)
|
||||
i = next - 1
|
||||
}
|
||||
}
|
||||
return commands
|
||||
}
|
||||
|
||||
func environmentsIn(source string) []string {
|
||||
var envs []string
|
||||
for index := 0; index < len(source); {
|
||||
position := strings.Index(source[index:], `\begin{`)
|
||||
if position < 0 {
|
||||
break
|
||||
}
|
||||
position += index
|
||||
start := position + len(`\begin{`)
|
||||
end := strings.IndexByte(source[start:], '}')
|
||||
if end < 0 {
|
||||
break
|
||||
}
|
||||
envs = append(envs, source[start:start+end])
|
||||
index = start
|
||||
}
|
||||
return envs
|
||||
}
|
||||
|
||||
// readEnvironment reads \begin{name} or \end{name} at index and returns the
|
||||
// environment name and the position after the closing brace.
|
||||
func readEnvironment(source string, index int) (string, int, bool) {
|
||||
prefix := `\begin{`
|
||||
if strings.HasPrefix(source[index:], `\end{`) {
|
||||
prefix = `\end{`
|
||||
} else if !strings.HasPrefix(source[index:], prefix) {
|
||||
return "", index, false
|
||||
}
|
||||
start := index + len(prefix)
|
||||
end := strings.IndexByte(source[start:], '}')
|
||||
if end < 0 {
|
||||
return "", index, false
|
||||
}
|
||||
return source[start : start+end], start + end + 1, true
|
||||
}
|
||||
|
||||
func isCommandLetter(c byte) bool {
|
||||
return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
|
||||
}
|
||||
|
||||
type parser struct {
|
||||
src string
|
||||
pos int
|
||||
}
|
||||
|
||||
func (p *parser) eof() bool { return p.pos >= len(p.src) }
|
||||
|
||||
func (p *parser) startsWith(prefix string) bool {
|
||||
return strings.HasPrefix(p.src[p.pos:], prefix)
|
||||
}
|
||||
|
||||
func (p *parser) parseBlocks() []Block {
|
||||
var blocks []Block
|
||||
var text strings.Builder
|
||||
flush := func() {
|
||||
trimmed := strings.TrimSpace(text.String())
|
||||
text.Reset()
|
||||
if trimmed == "" {
|
||||
return
|
||||
}
|
||||
blocks = append(blocks, Block{Type: "paragraph", Text: parseInline(trimmed)})
|
||||
}
|
||||
|
||||
for !p.eof() {
|
||||
switch {
|
||||
case p.startsWith(`\[`):
|
||||
flush()
|
||||
p.pos += 2
|
||||
expression := strings.TrimSpace(p.readUntil(`\]`))
|
||||
if isSymbolsOnly(expression) {
|
||||
blocks = append(blocks, Block{Type: "paragraph", Text: strings.TrimSpace(expression)})
|
||||
} else {
|
||||
blocks = append(blocks, Block{Type: "mathematical_expression", Expression: expression})
|
||||
}
|
||||
case p.startsWith("$$"):
|
||||
flush()
|
||||
p.pos += 2
|
||||
expression := strings.TrimSpace(p.readUntil("$$"))
|
||||
if isSymbolsOnly(expression) {
|
||||
blocks = append(blocks, Block{Type: "paragraph", Text: strings.TrimSpace(expression)})
|
||||
} else {
|
||||
blocks = append(blocks, Block{Type: "mathematical_expression", Expression: expression})
|
||||
}
|
||||
case p.startsWith(`\section`):
|
||||
flush()
|
||||
p.pos += len(`\section`)
|
||||
if title, ok := p.readOptionalStarAndGroup(); ok {
|
||||
blocks = append(blocks, Block{Type: "heading", Text: parseInline(title), Size: 2})
|
||||
}
|
||||
case p.startsWith(`\subsection`):
|
||||
flush()
|
||||
p.pos += len(`\subsection`)
|
||||
if title, ok := p.readOptionalStarAndGroup(); ok {
|
||||
blocks = append(blocks, Block{Type: "heading", Text: parseInline(title), Size: 3})
|
||||
}
|
||||
case p.startsWith(`\subsubsection`):
|
||||
flush()
|
||||
p.pos += len(`\subsubsection`)
|
||||
if title, ok := p.readOptionalStarAndGroup(); ok {
|
||||
blocks = append(blocks, Block{Type: "heading", Text: parseInline(title), Size: 4})
|
||||
}
|
||||
case p.startsWith(`\item`):
|
||||
flush()
|
||||
text.WriteString("• ")
|
||||
p.pos += len(`\item`)
|
||||
for p.pos < len(p.src) && (p.src[p.pos] == ' ' || p.src[p.pos] == '\t') {
|
||||
p.pos++
|
||||
}
|
||||
case p.startsWith(`\begin{`), p.startsWith(`\end{`):
|
||||
environment, after, ok := readEnvironment(p.src, p.pos)
|
||||
if ok && wrapperEnvironments[environment] {
|
||||
flush()
|
||||
p.pos = after
|
||||
} else {
|
||||
text.WriteByte(p.src[p.pos])
|
||||
p.pos++
|
||||
}
|
||||
case p.startsWith(`\par`), p.startsWith(`\\`):
|
||||
flush()
|
||||
if p.startsWith(`\par`) {
|
||||
p.pos += len(`\par`)
|
||||
} else {
|
||||
p.pos += 2
|
||||
}
|
||||
case p.src[p.pos] == '\n' && p.blankLineAhead():
|
||||
flush()
|
||||
for !p.eof() && (p.src[p.pos] == '\n' || p.src[p.pos] == ' ' || p.src[p.pos] == '\t' || p.src[p.pos] == '\r') {
|
||||
p.pos++
|
||||
}
|
||||
default:
|
||||
text.WriteByte(p.src[p.pos])
|
||||
p.pos++
|
||||
}
|
||||
}
|
||||
flush()
|
||||
return blocks
|
||||
}
|
||||
|
||||
// blankLineAhead reports whether the newline at the current position is
|
||||
// followed by a blank line (two or more newlines, spaces allowed between).
|
||||
func (p *parser) blankLineAhead() bool {
|
||||
index := p.pos + 1
|
||||
if index < len(p.src) && p.src[index] == '\r' {
|
||||
index++
|
||||
}
|
||||
for index < len(p.src) && (p.src[index] == ' ' || p.src[index] == '\t') {
|
||||
index++
|
||||
}
|
||||
return index < len(p.src) && p.src[index] == '\n'
|
||||
}
|
||||
|
||||
func (p *parser) readUntil(terminator string) string {
|
||||
index := strings.Index(p.src[p.pos:], terminator)
|
||||
if index < 0 {
|
||||
result := p.src[p.pos:]
|
||||
p.pos = len(p.src)
|
||||
return result
|
||||
}
|
||||
result := p.src[p.pos : p.pos+index]
|
||||
p.pos += index + len(terminator)
|
||||
return result
|
||||
}
|
||||
|
||||
// readOptionalStarAndGroup consumes an optional star and the following
|
||||
// {group}; spaces are allowed before the brace.
|
||||
func (p *parser) readOptionalStarAndGroup() (string, bool) {
|
||||
if p.pos < len(p.src) && p.src[p.pos] == '*' {
|
||||
p.pos++
|
||||
}
|
||||
return p.readGroup()
|
||||
}
|
||||
|
||||
// readGroup reads a balanced {group} at the current position, skipping
|
||||
// leading spaces.
|
||||
func (p *parser) readGroup() (string, bool) {
|
||||
for p.pos < len(p.src) && (p.src[p.pos] == ' ' || p.src[p.pos] == '\t') {
|
||||
p.pos++
|
||||
}
|
||||
if p.pos >= len(p.src) || p.src[p.pos] != '{' {
|
||||
return "", false
|
||||
}
|
||||
depth := 0
|
||||
start := p.pos + 1
|
||||
for index := p.pos; index < len(p.src); index++ {
|
||||
switch p.src[index] {
|
||||
case '\\':
|
||||
index++ // skip the escaped character
|
||||
case '{':
|
||||
depth++
|
||||
case '}':
|
||||
depth--
|
||||
if depth == 0 {
|
||||
result := p.src[start:index]
|
||||
p.pos = index + 1
|
||||
return result, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// parseInline converts a text-mode fragment into a RichText value: a plain
|
||||
// string when nothing needs formatting, or a slice of strings and inline
|
||||
// nodes otherwise.
|
||||
func parseInline(source string) any {
|
||||
var parts []any
|
||||
var text strings.Builder
|
||||
flush := func() {
|
||||
if text.Len() > 0 {
|
||||
parts = append(parts, text.String())
|
||||
text.Reset()
|
||||
}
|
||||
}
|
||||
appendParsed := func(value string) {
|
||||
flush()
|
||||
switch parsed := parseInline(value).(type) {
|
||||
case string:
|
||||
if parsed != "" {
|
||||
parts = append(parts, parsed)
|
||||
}
|
||||
case []any:
|
||||
parts = append(parts, parsed...)
|
||||
}
|
||||
}
|
||||
|
||||
for index := 0; index < len(source); {
|
||||
char := source[index]
|
||||
switch {
|
||||
case char == '\\':
|
||||
if index+1 >= len(source) {
|
||||
text.WriteByte(char)
|
||||
index++
|
||||
continue
|
||||
}
|
||||
switch source[index+1] {
|
||||
case '&', '%', '$', '#', '_', '{', '}':
|
||||
text.WriteByte(source[index+1])
|
||||
index += 2
|
||||
continue
|
||||
case '\\':
|
||||
text.WriteByte('\n')
|
||||
index += 2
|
||||
continue
|
||||
case '(':
|
||||
end := strings.Index(source[index+2:], `\)`)
|
||||
if end >= 0 {
|
||||
expression := source[index+2 : index+2+end]
|
||||
if isSymbolsOnly(expression) {
|
||||
text.WriteString(strings.TrimSpace(expression))
|
||||
} else {
|
||||
flush()
|
||||
parts = append(parts, Inline{Type: "mathematical_expression", Expression: expression})
|
||||
}
|
||||
index += 2 + end + 2
|
||||
continue
|
||||
}
|
||||
}
|
||||
name, next := readCommand(source, index)
|
||||
if name == "" {
|
||||
text.WriteByte(char)
|
||||
index++
|
||||
continue
|
||||
}
|
||||
arg, after, hasArg := readGroupAt(source, next)
|
||||
switch name {
|
||||
case "textbf":
|
||||
if hasArg {
|
||||
flush()
|
||||
parts = append(parts, Inline{Type: "bold", Text: parseInline(arg)})
|
||||
index = after
|
||||
continue
|
||||
}
|
||||
case "textit", "emph":
|
||||
if hasArg {
|
||||
flush()
|
||||
parts = append(parts, Inline{Type: "italic", Text: parseInline(arg)})
|
||||
index = after
|
||||
continue
|
||||
}
|
||||
case "underline", "uline":
|
||||
if hasArg {
|
||||
flush()
|
||||
parts = append(parts, Inline{Type: "underline", Text: parseInline(arg)})
|
||||
index = after
|
||||
continue
|
||||
}
|
||||
case "texttt":
|
||||
if hasArg {
|
||||
flush()
|
||||
parts = append(parts, Inline{Type: "code", Text: parseInline(arg)})
|
||||
index = after
|
||||
continue
|
||||
}
|
||||
case "textsuperscript":
|
||||
if hasArg {
|
||||
flush()
|
||||
parts = append(parts, Inline{Type: "superscript", Text: parseInline(arg)})
|
||||
index = after
|
||||
continue
|
||||
}
|
||||
case "textsubscript":
|
||||
if hasArg {
|
||||
flush()
|
||||
parts = append(parts, Inline{Type: "subscript", Text: parseInline(arg)})
|
||||
index = after
|
||||
continue
|
||||
}
|
||||
case "textrm", "textsf", "textnormal", "textup", "textmd", "text":
|
||||
if hasArg {
|
||||
appendParsed(arg)
|
||||
index = after
|
||||
continue
|
||||
}
|
||||
case "textcolor":
|
||||
if _, afterColor, ok := readGroupAt(source, next); ok {
|
||||
if value, afterValue, ok := readGroupAt(source, afterColor); ok {
|
||||
appendParsed(value)
|
||||
index = afterValue
|
||||
continue
|
||||
}
|
||||
}
|
||||
case "href":
|
||||
if url, afterURL, ok := readGroupAt(source, next); ok {
|
||||
if label, afterLabel, ok := readGroupAt(source, afterURL); ok {
|
||||
flush()
|
||||
parts = append(parts, Inline{Type: "url", Text: parseInline(label), URL: url})
|
||||
index = afterLabel
|
||||
continue
|
||||
}
|
||||
}
|
||||
case "hyperlink", "hypertarget":
|
||||
// Telegram rich text has no named-anchor equivalent. Keep
|
||||
// the visible label and discard the internal PDF target; the
|
||||
// PNG/PDF renderer preserves the actual hyperref command.
|
||||
if _, afterTarget, ok := readGroupAt(source, next); ok {
|
||||
if label, afterLabel, ok := readGroupAt(source, afterTarget); ok {
|
||||
appendParsed(label)
|
||||
index = afterLabel
|
||||
continue
|
||||
}
|
||||
}
|
||||
case "url":
|
||||
if hasArg {
|
||||
flush()
|
||||
parts = append(parts, Inline{Type: "url", Text: arg, URL: arg})
|
||||
index = after
|
||||
continue
|
||||
}
|
||||
case "item":
|
||||
text.WriteString("• ")
|
||||
index = skipEmptyGroup(source, next)
|
||||
continue
|
||||
case "quad", "qquad", ",", ";", ":", "!", " ":
|
||||
text.WriteByte(' ')
|
||||
index = skipEmptyGroup(source, next)
|
||||
continue
|
||||
case "ldots", "dots":
|
||||
text.WriteString("…")
|
||||
index = skipEmptyGroup(source, next)
|
||||
continue
|
||||
case "LaTeX":
|
||||
text.WriteString("LaTeX")
|
||||
index = skipEmptyGroup(source, next)
|
||||
continue
|
||||
case "TeX":
|
||||
text.WriteString("TeX")
|
||||
index = skipEmptyGroup(source, next)
|
||||
continue
|
||||
}
|
||||
text.WriteString(`\` + name)
|
||||
index = next
|
||||
case char == '$':
|
||||
end := strings.IndexByte(source[index+1:], '$')
|
||||
if end >= 0 && end > 0 {
|
||||
expression := source[index+1 : index+1+end]
|
||||
if isSymbolsOnly(expression) {
|
||||
// Client math fonts often lack exotic symbols; the
|
||||
// regular text font renders them everywhere.
|
||||
text.WriteString(strings.TrimSpace(expression))
|
||||
} else {
|
||||
flush()
|
||||
parts = append(parts, Inline{Type: "mathematical_expression", Expression: expression})
|
||||
}
|
||||
index += end + 2
|
||||
continue
|
||||
}
|
||||
text.WriteByte(char)
|
||||
index++
|
||||
case char == '~':
|
||||
text.WriteByte(' ')
|
||||
index++
|
||||
default:
|
||||
text.WriteByte(char)
|
||||
index++
|
||||
}
|
||||
}
|
||||
flush()
|
||||
|
||||
if len(parts) == 1 {
|
||||
if value, ok := parts[0].(string); ok {
|
||||
return value
|
||||
}
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
return ""
|
||||
}
|
||||
return parts
|
||||
}
|
||||
|
||||
// skipEmptyGroup skips an empty {} group immediately following a control
|
||||
// word, as in \LaTeX{}.
|
||||
func skipEmptyGroup(source string, index int) int {
|
||||
if strings.HasPrefix(source[index:], "{}") {
|
||||
return index + 2
|
||||
}
|
||||
return index
|
||||
}
|
||||
|
||||
// readCommand reads a backslash command name at position index.
|
||||
func readCommand(source string, index int) (string, int) {
|
||||
next := index + 1
|
||||
name := ""
|
||||
for next < len(source) && isCommandLetter(source[next]) {
|
||||
name += string(source[next])
|
||||
next++
|
||||
}
|
||||
if name == "" && next < len(source) {
|
||||
// Single-character commands such as \, or \; .
|
||||
name = string(source[next])
|
||||
next++
|
||||
}
|
||||
return name, next
|
||||
}
|
||||
|
||||
// readGroupAt reads a balanced {group} starting at position index, skipping
|
||||
// spaces.
|
||||
func readGroupAt(source string, index int) (string, int, bool) {
|
||||
for index < len(source) && (source[index] == ' ' || source[index] == '\t') {
|
||||
index++
|
||||
}
|
||||
if index >= len(source) || source[index] != '{' {
|
||||
return "", index, false
|
||||
}
|
||||
depth := 0
|
||||
start := index + 1
|
||||
for position := index; position < len(source); position++ {
|
||||
switch source[position] {
|
||||
case '\\':
|
||||
position++ // skip the escaped character
|
||||
case '{':
|
||||
depth++
|
||||
case '}':
|
||||
depth--
|
||||
if depth == 0 {
|
||||
return source[start:position], position + 1, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return "", index, false
|
||||
}
|
||||
404
internal/rich/latex_test.go
Normal file
404
internal/rich/latex_test.go
Normal file
@@ -0,0 +1,404 @@
|
||||
package rich
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDocumentPureMath(t *testing.T) {
|
||||
message := Document(`\int_0^\infty e^{-x^2}\,dx`, true)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "mathematical_expression" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
if message.Blocks[0].Expression != `\int_0^\infty e^{-x^2}\,dx` {
|
||||
t.Fatalf("expression = %q", message.Blocks[0].Expression)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentDollarMath(t *testing.T) {
|
||||
message := Document(`$F=ma$`, true)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "paragraph" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentDisplayMath(t *testing.T) {
|
||||
message := Document("before\n\n$$F=ma$$\n\nafter", true)
|
||||
if len(message.Blocks) != 3 {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
if message.Blocks[0].Type != "paragraph" || message.Blocks[1].Type != "mathematical_expression" || message.Blocks[2].Type != "paragraph" {
|
||||
t.Fatalf("block types = %s/%s/%s", message.Blocks[0].Type, message.Blocks[1].Type, message.Blocks[2].Type)
|
||||
}
|
||||
if message.Blocks[1].Expression != "F=ma" {
|
||||
t.Fatalf("expression = %q", message.Blocks[1].Expression)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentTextFormatting(t *testing.T) {
|
||||
message := Document(`Hello \textbf{world} and \emph{italics}`, true)
|
||||
if len(message.Blocks) != 1 {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
parts, ok := message.Blocks[0].Text.([]any)
|
||||
if !ok || len(parts) != 4 {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
if parts[0] != "Hello " {
|
||||
t.Errorf("parts[0] = %#v", parts[0])
|
||||
}
|
||||
bold, ok := parts[1].(Inline)
|
||||
if !ok || bold.Type != "bold" || bold.Text != "world" {
|
||||
t.Errorf("parts[1] = %#v", parts[1])
|
||||
}
|
||||
if parts[2] != " and " {
|
||||
t.Errorf("parts[2] = %#v", parts[2])
|
||||
}
|
||||
italic, ok := parts[3].(Inline)
|
||||
if !ok || italic.Type != "italic" || italic.Text != "italics" {
|
||||
t.Errorf("parts[3] = %#v", parts[3])
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentNestedFormatting(t *testing.T) {
|
||||
message := Document(`\textbf{bold \emph{and italic}}`, true)
|
||||
parts, ok := message.Blocks[0].Text.([]any)
|
||||
if !ok || len(parts) != 1 {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
bold, ok := parts[0].(Inline)
|
||||
if !ok || bold.Type != "bold" {
|
||||
t.Fatalf("parts[0] = %#v", parts[0])
|
||||
}
|
||||
inner, ok := bold.Text.([]any)
|
||||
if !ok || len(inner) != 2 {
|
||||
t.Fatalf("bold text = %#v", bold.Text)
|
||||
}
|
||||
italic, ok := inner[1].(Inline)
|
||||
if !ok || italic.Type != "italic" || italic.Text != "and italic" {
|
||||
t.Fatalf("inner = %#v", inner[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentMixedTextAndMath(t *testing.T) {
|
||||
message := Document(`The area is $\pi r^2$ exactly`, true)
|
||||
parts, ok := message.Blocks[0].Text.([]any)
|
||||
if !ok || len(parts) != 3 {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
math, ok := parts[1].(Inline)
|
||||
if !ok || math.Type != "mathematical_expression" || math.Expression != `\pi r^2` {
|
||||
t.Fatalf("math = %#v", parts[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentSections(t *testing.T) {
|
||||
message := Document("\\section{Intro}\n\nSome text", true)
|
||||
if len(message.Blocks) != 2 {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
if message.Blocks[0].Type != "heading" || message.Blocks[0].Size != 2 {
|
||||
t.Fatalf("heading = %+v", message.Blocks[0])
|
||||
}
|
||||
if message.Blocks[0].Text != "Intro" {
|
||||
t.Fatalf("heading text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentList(t *testing.T) {
|
||||
message := Document("\\begin{itemize}\n\\item one\n\\item two\n\\end{itemize}", true)
|
||||
if len(message.Blocks) != 2 {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
if message.Blocks[0].Text != "• one" || message.Blocks[1].Text != "• two" {
|
||||
t.Fatalf("items = %#v / %#v", message.Blocks[0].Text, message.Blocks[1].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentEscapesAndCommands(t *testing.T) {
|
||||
message := Document(`100\% \& more \LaTeX{} \url{https://example.com}`, true)
|
||||
parts, ok := message.Blocks[0].Text.([]any)
|
||||
if !ok {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
joined := ""
|
||||
for _, part := range parts {
|
||||
if text, ok := part.(string); ok {
|
||||
joined += text
|
||||
}
|
||||
}
|
||||
if joined != "100% & more LaTeX " {
|
||||
t.Fatalf("joined = %q", joined)
|
||||
}
|
||||
url, ok := parts[len(parts)-1].(Inline)
|
||||
if !ok || url.Type != "url" || url.URL != "https://example.com" {
|
||||
t.Fatalf("url = %#v", parts[len(parts)-1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentHyperrefCommands(t *testing.T) {
|
||||
message := Document(`See \hyperlink{details}{the \textbf{details}}.`, true)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "paragraph" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
parts, ok := message.Blocks[0].Text.([]any)
|
||||
if !ok || len(parts) != 4 {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
if parts[0] != "See " || parts[1] != "the " || parts[3] != "." {
|
||||
t.Fatalf("parts = %#v", parts)
|
||||
}
|
||||
bold, ok := parts[2].(Inline)
|
||||
if !ok || bold.Type != "bold" || bold.Text != "details" {
|
||||
t.Fatalf("bold = %#v", parts[2])
|
||||
}
|
||||
|
||||
target := Document(`\hypertarget{details}{Details}`, true)
|
||||
if len(target.Blocks) != 1 || target.Blocks[0].Type != "paragraph" || target.Blocks[0].Text != "Details" {
|
||||
t.Fatalf("target blocks = %+v", target.Blocks)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentFullDocumentBody(t *testing.T) {
|
||||
source := "\\documentclass{article}\n\\begin{document}\nHello \\textbf{world}\n\\end{document}"
|
||||
message := Document(source, true)
|
||||
if len(message.Blocks) != 1 {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
if message.Blocks[0].Type != "paragraph" {
|
||||
t.Fatalf("block = %+v", message.Blocks[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestInlineDocumentLeadingMath(t *testing.T) {
|
||||
message := InlineDocument(`$F=ma$`, true)
|
||||
if len(message.Blocks) != 1 {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
parts, ok := message.Blocks[0].Text.([]any)
|
||||
if !ok || len(parts) != 2 || parts[0] != " " {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
math, ok := parts[1].(Inline)
|
||||
if !ok || math.Type != "mathematical_expression" || math.Expression != "F=ma" {
|
||||
t.Fatalf("math = %#v", parts[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestInlineDocumentDisplayMath(t *testing.T) {
|
||||
message := InlineDocument("$$F=ma$$", true)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "paragraph" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
parts, ok := message.Blocks[0].Text.([]any)
|
||||
if !ok || len(parts) != 2 || parts[0] != " " {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInlineDocumentLeadingParagraph(t *testing.T) {
|
||||
message := InlineDocument(`Hello \textbf{world}`, true)
|
||||
parts, ok := message.Blocks[0].Text.([]any)
|
||||
if !ok || len(parts) == 0 || parts[0] != " " {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentJSONMarshal(t *testing.T) {
|
||||
data, err := json.Marshal(Document(`Hello \textbf{world}`, true))
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal: %v", err)
|
||||
}
|
||||
expected := `{"blocks":[{"type":"paragraph","text":["Hello ",{"type":"bold","text":"world"}]}]}`
|
||||
if string(data) != expected {
|
||||
t.Fatalf("unexpected JSON:\n got %s\nwant %s", data, expected)
|
||||
}
|
||||
}
|
||||
|
||||
// blockText returns the math expression or the paragraph text of a block.
|
||||
func blockText(block Block) string {
|
||||
if block.Expression != "" {
|
||||
return block.Expression
|
||||
}
|
||||
if text, ok := block.Text.(string); ok {
|
||||
return text
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func TestDocumentSymbolCommands(t *testing.T) {
|
||||
// A lone symbol renders as text: client math fonts often lack exotic
|
||||
// symbols, while the regular text font has them.
|
||||
message := Document(`\heartsuit`, true)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "paragraph" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
if blockText(message.Blocks[0]) != "♡" {
|
||||
t.Fatalf("text = %q", blockText(message.Blocks[0]))
|
||||
}
|
||||
|
||||
// With surrounding letters it stays a math expression.
|
||||
message = Document(`I \heartsuit{} you`, true)
|
||||
if message.Blocks[0].Type != "mathematical_expression" || blockText(message.Blocks[0]) != "I ♡ you" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
|
||||
// Standard math commands are left to the client engine, which renders
|
||||
// them natively (Unicode replacements are not drawn by every font).
|
||||
message = Document(`\int_0^\infty e^{-x^2}\,dx`, true)
|
||||
if blockText(message.Blocks[0]) != `\int_0^\infty e^{-x^2}\,dx` {
|
||||
t.Fatalf("expression = %q", blockText(message.Blocks[0]))
|
||||
}
|
||||
|
||||
for source, want := range map[string]string{
|
||||
`\alpha + \beta = \gamma`: `\alpha + \beta = \gamma`,
|
||||
`\sum_{i=1}^n i`: `\sum_{i=1}^n i`,
|
||||
`a \leq b \to c`: `a \leq b \to c`,
|
||||
`\mathbb{R}^n`: `\mathbb{R}^n`,
|
||||
`\mathcal{L}(f)`: `\mathcal{L}(f)`,
|
||||
`\mathfrak{g} \otimes \mathfrak{h}`: `\mathfrak{g} \otimes \mathfrak{h}`,
|
||||
`\mathbf{v} \cdot \nabla f`: `\mathbf{v} \cdot \nabla f`,
|
||||
`\heartsuit`: "♡",
|
||||
} {
|
||||
message := Document(source, true)
|
||||
if got := blockText(message.Blocks[0]); got != want {
|
||||
t.Errorf("source %q: got %q, want %q", source, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentIntegralCommandPreserved(t *testing.T) {
|
||||
source := `$$\int_0^\infty e^{-x^{porcodio}}\,dx$$`
|
||||
message := Document(source, true)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "mathematical_expression" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
if got := message.Blocks[0].Expression; got != `\int_0^\infty e^{-x^{porcodio}}\,dx` {
|
||||
t.Fatalf("expression = %q, integral command must stay for the client engine", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentSymbolMathSegments(t *testing.T) {
|
||||
// A math segment containing only symbols is hoisted to text: the math
|
||||
// font does not carry exotic glyphs, the regular font does.
|
||||
for _, source := range []string{`$\heartsuit$`, `$$\heartsuit$$`, `\( \heartsuit \)`, `\[ \heartsuit \]`} {
|
||||
message := Document(source, true)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "paragraph" {
|
||||
t.Errorf("source %q: blocks = %+v", source, message.Blocks)
|
||||
continue
|
||||
}
|
||||
if got := blockText(message.Blocks[0]); got != "♡" {
|
||||
t.Errorf("source %q: text = %q, want ♡", source, got)
|
||||
}
|
||||
}
|
||||
|
||||
// Symbols mixed with real math stay in a math node.
|
||||
message := Document(`$x \heartsuit y$`, true)
|
||||
if message.Blocks[0].Type != "paragraph" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
parts, ok := message.Blocks[0].Text.([]any)
|
||||
if !ok || len(parts) != 1 {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
math, ok := parts[0].(Inline)
|
||||
if !ok || math.Type != "mathematical_expression" || math.Expression != "x ♡ y" {
|
||||
t.Fatalf("math = %#v", parts[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsMathExpression(t *testing.T) {
|
||||
for source, want := range map[string]bool{
|
||||
`\heartsuit`: true,
|
||||
`\int_0^\infty x`: true,
|
||||
`F = ma`: true,
|
||||
`$F=ma$`: true,
|
||||
`\begin{matrix} a & b \end{matrix}`: true,
|
||||
`Hello \textbf{x}`: false,
|
||||
`\hyperlink{x}{go}`: false,
|
||||
`\begin{itemize} \item a`: false,
|
||||
`\begin{itemize} \item $x$`: false,
|
||||
} {
|
||||
if got := IsMathExpression(source); got != want {
|
||||
t.Errorf("IsMathExpression(%q) = %v, want %v", source, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestContainsMathEnvironment(t *testing.T) {
|
||||
for source, want := range map[string]bool{
|
||||
`\begin{align} a &= b \end{align}`: true,
|
||||
`\begin{matrix} a & b \end{matrix}`: true,
|
||||
`\begin{itemize} a`: false,
|
||||
`x^2`: false,
|
||||
} {
|
||||
if got := ContainsMathEnvironment(source); got != want {
|
||||
t.Errorf("ContainsMathEnvironment(%q) = %v, want %v", source, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentAlphabetFallback(t *testing.T) {
|
||||
// A font command with non-alphanumeric content is left to the engine.
|
||||
message := Document(`\mathbb{R_1}`, true)
|
||||
if message.Blocks[0].Expression != `\mathbb{R_1}` {
|
||||
t.Fatalf("expression = %q", message.Blocks[0].Expression)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPureMathKeepsPlainText(t *testing.T) {
|
||||
message := Document("F = ma", true)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "mathematical_expression" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentTextDefault(t *testing.T) {
|
||||
// With defaultMath disabled a bare source is text, not math.
|
||||
message := Document("F = ma", false)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "paragraph" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
if message.Blocks[0].Text != "F = ma" {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
|
||||
// Delimited math still renders as math.
|
||||
message = Document("$F = ma$", false)
|
||||
parts, ok := message.Blocks[0].Text.([]any)
|
||||
if !ok || len(parts) != 1 {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
if math, ok := parts[0].(Inline); !ok || math.Type != "mathematical_expression" {
|
||||
t.Fatalf("parts[0] = %#v", parts[0])
|
||||
}
|
||||
|
||||
// Text-mode commands still work.
|
||||
message = Document(`Hello \textbf{world}`, false)
|
||||
if _, ok := message.Blocks[0].Text.([]any); !ok {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInlineDocumentTextDefault(t *testing.T) {
|
||||
message := InlineDocument("F = ma", false)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "paragraph" {
|
||||
t.Fatalf("blocks = %+v", message.Blocks)
|
||||
}
|
||||
if message.Blocks[0].Text != " F = ma" {
|
||||
t.Fatalf("text = %#v", message.Blocks[0].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocumentBareMathCommands(t *testing.T) {
|
||||
for _, source := range []string{
|
||||
`\frac{1}{2}`,
|
||||
`\begin{align} a &= b \\ c &= d \end{align}`,
|
||||
`x = \text{hello}`,
|
||||
} {
|
||||
message := Document(source, true)
|
||||
if len(message.Blocks) != 1 || message.Blocks[0].Type != "mathematical_expression" {
|
||||
t.Errorf("source %q: blocks = %+v", source, message.Blocks)
|
||||
}
|
||||
}
|
||||
}
|
||||
28
internal/rich/rich.go
Normal file
28
internal/rich/rich.go
Normal file
@@ -0,0 +1,28 @@
|
||||
// Package rich builds rich-message payloads for the Bot API 10.x rich
|
||||
// message endpoints (sendRichMessage and the InputRichMessageContent used in
|
||||
// inline answers). The pinned telegram-bot-api library predates the rich
|
||||
// message API, so these are plain JSON-encodable values sent through
|
||||
// api.MakeRequest.
|
||||
package rich
|
||||
|
||||
// Message is the payload for InputRichMessage (and structurally for
|
||||
// RichMessage). Only the blocks representation is used.
|
||||
type Message struct {
|
||||
Blocks []Block `json:"blocks"`
|
||||
}
|
||||
|
||||
// Block is a RichBlock/InputRichBlock.
|
||||
type Block struct {
|
||||
Type string `json:"type"`
|
||||
Text any `json:"text,omitempty"`
|
||||
Expression string `json:"expression,omitempty"`
|
||||
Size int `json:"size,omitempty"`
|
||||
}
|
||||
|
||||
// Inline is a RichText inline node (mathematical_expression, bold, url...).
|
||||
type Inline struct {
|
||||
Type string `json:"type"`
|
||||
Text any `json:"text,omitempty"`
|
||||
Expression string `json:"expression,omitempty"`
|
||||
URL string `json:"url,omitempty"`
|
||||
}
|
||||
296
internal/store/store.go
Normal file
296
internal/store/store.go
Normal file
@@ -0,0 +1,296 @@
|
||||
// Package store persists settings, users, group memberships and render
|
||||
// statistics in the shared MongoDB cluster.
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
"go.mongodb.org/mongo-driver/mongo/readpref"
|
||||
)
|
||||
|
||||
// Settings are the per-chat overrides. Zero values mean "use the default".
|
||||
type Settings struct {
|
||||
Preamble string `json:"preamble,omitempty" bson:"preamble,omitempty"`
|
||||
DPI int `json:"dpi,omitempty" bson:"dpi,omitempty"`
|
||||
}
|
||||
|
||||
// User is the caller information recorded on every interaction.
|
||||
type User struct {
|
||||
ID int64
|
||||
Username string
|
||||
FirstName string
|
||||
LastName string
|
||||
}
|
||||
|
||||
// Chat is the group information recorded on membership changes.
|
||||
type Chat struct {
|
||||
ID int64
|
||||
Type string
|
||||
Title string
|
||||
Username string
|
||||
}
|
||||
|
||||
// RenderStat is one render event, kept for statistics.
|
||||
type RenderStat struct {
|
||||
Source string // chat | inline
|
||||
Format string // rich | png | pdf | inline
|
||||
ChatID int64
|
||||
UserID int64
|
||||
DPI int
|
||||
ExpressionLen int
|
||||
OK bool
|
||||
Error string
|
||||
DurationMs int64
|
||||
}
|
||||
|
||||
// Store provides all persistence for the bot.
|
||||
type Store struct {
|
||||
client *mongo.Client
|
||||
settings *mongo.Collection
|
||||
users *mongo.Collection
|
||||
groups *mongo.Collection
|
||||
stats *mongo.Collection
|
||||
files *mongo.Collection
|
||||
}
|
||||
|
||||
// New connects a Store to the given database.
|
||||
func New(client *mongo.Client, database string) *Store {
|
||||
db := client.Database(database)
|
||||
return &Store{
|
||||
client: client,
|
||||
settings: db.Collection("settings"),
|
||||
users: db.Collection("users"),
|
||||
groups: db.Collection("groups"),
|
||||
stats: db.Collection("stats"),
|
||||
files: db.Collection("files"),
|
||||
}
|
||||
}
|
||||
|
||||
// Ping reports whether the cluster is reachable.
|
||||
func (s *Store) Ping(ctx context.Context) error {
|
||||
return s.client.Ping(ctx, readpref.PrimaryPreferred())
|
||||
}
|
||||
|
||||
// EnsureIndexes creates the indexes the bot relies on.
|
||||
func (s *Store) EnsureIndexes(ctx context.Context) error {
|
||||
_, err := s.stats.Indexes().CreateMany(ctx, []mongo.IndexModel{
|
||||
{Keys: bson.D{{Key: "ts", Value: -1}}, Options: options.Index().SetName("stats_ts")},
|
||||
{Keys: bson.D{{Key: "user_id", Value: 1}, {Key: "ts", Value: -1}}, Options: options.Index().SetName("stats_user_ts")},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// Get returns the settings for chatID, or the zero value.
|
||||
func (s *Store) Get(chatID int64) Settings {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var settings Settings
|
||||
err := s.settings.FindOne(ctx, bson.M{"_id": chatID}).Decode(&settings)
|
||||
if err != nil {
|
||||
if err != mongo.ErrNoDocuments {
|
||||
log.Printf("store: read settings for %d: %v", chatID, err)
|
||||
}
|
||||
return Settings{}
|
||||
}
|
||||
return settings
|
||||
}
|
||||
|
||||
// Set stores settings for chatID.
|
||||
func (s *Store) Set(chatID int64, settings Settings) error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_, err := s.settings.UpdateOne(ctx,
|
||||
bson.M{"_id": chatID},
|
||||
bson.M{"$set": settings, "$setOnInsert": bson.M{"created_at": time.Now().UTC()}},
|
||||
options.Update().SetUpsert(true),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("store: write settings for %d: %w", chatID, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RecordUser upserts the interacting user and refreshes last_seen_at.
|
||||
func (s *Store) RecordUser(user User) {
|
||||
if user.ID == 0 {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
now := time.Now().UTC()
|
||||
set := bson.M{"last_seen_at": now}
|
||||
for key, value := range map[string]string{
|
||||
"username": user.Username, "first_name": user.FirstName, "last_name": user.LastName,
|
||||
} {
|
||||
if strings.TrimSpace(value) != "" {
|
||||
set[key] = value
|
||||
}
|
||||
}
|
||||
_, err := s.users.UpdateOne(ctx,
|
||||
bson.M{"_id": user.ID},
|
||||
bson.M{"$set": set, "$setOnInsert": bson.M{"first_seen_at": now, "renders": int64(0)}},
|
||||
options.Update().SetUpsert(true),
|
||||
)
|
||||
if err != nil {
|
||||
log.Printf("store: record user %d: %v", user.ID, err)
|
||||
}
|
||||
}
|
||||
|
||||
// RecordGroup upserts a chat the bot is a member of (or left).
|
||||
func (s *Store) RecordGroup(chat Chat, member bool) {
|
||||
if chat.ID == 0 {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
now := time.Now().UTC()
|
||||
status := "left"
|
||||
if member {
|
||||
status = "member"
|
||||
}
|
||||
set := bson.M{"status": status, "updated_at": now}
|
||||
for key, value := range map[string]string{
|
||||
"type": chat.Type, "title": chat.Title, "username": chat.Username,
|
||||
} {
|
||||
if strings.TrimSpace(value) != "" {
|
||||
set[key] = value
|
||||
}
|
||||
}
|
||||
_, err := s.groups.UpdateOne(ctx,
|
||||
bson.M{"_id": chat.ID},
|
||||
bson.M{"$set": set, "$setOnInsert": bson.M{"added_at": now}},
|
||||
options.Update().SetUpsert(true),
|
||||
)
|
||||
if err != nil {
|
||||
log.Printf("store: record group %d: %v", chat.ID, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Stats are the aggregate numbers reported by the /stats command.
|
||||
type Stats struct {
|
||||
Renders int64
|
||||
ChatRenders int64
|
||||
InlineRenders int64
|
||||
FailedRenders int64
|
||||
Renders24h int64
|
||||
Users int64
|
||||
Groups int64
|
||||
}
|
||||
|
||||
// Stats computes the aggregate counts shown by /stats. Counting is cheap
|
||||
// enough for a command that is used rarely.
|
||||
func (s *Store) Stats() Stats {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
var stats Stats
|
||||
if count, err := s.stats.CountDocuments(ctx, bson.M{}); err == nil {
|
||||
stats.Renders = count
|
||||
} else {
|
||||
log.Printf("store: count renders: %v", err)
|
||||
}
|
||||
if count, err := s.stats.CountDocuments(ctx, bson.M{"source": "chat"}); err == nil {
|
||||
stats.ChatRenders = count
|
||||
}
|
||||
if count, err := s.stats.CountDocuments(ctx, bson.M{"source": "inline"}); err == nil {
|
||||
stats.InlineRenders = count
|
||||
}
|
||||
if count, err := s.stats.CountDocuments(ctx, bson.M{"ok": false}); err == nil {
|
||||
stats.FailedRenders = count
|
||||
}
|
||||
if count, err := s.stats.CountDocuments(ctx, bson.M{"ts": bson.M{"$gte": time.Now().UTC().Add(-24 * time.Hour)}}); err == nil {
|
||||
stats.Renders24h = count
|
||||
}
|
||||
if count, err := s.users.CountDocuments(ctx, bson.M{}); err == nil {
|
||||
stats.Users = count
|
||||
}
|
||||
if count, err := s.groups.CountDocuments(ctx, bson.M{"status": "member"}); err == nil {
|
||||
stats.Groups = count
|
||||
}
|
||||
return stats
|
||||
}
|
||||
|
||||
// RecordRender appends a render event to the statistics collection.
|
||||
func (s *Store) RecordRender(stat RenderStat) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_, err := s.stats.InsertOne(ctx, bson.M{
|
||||
"ts": time.Now().UTC(),
|
||||
"source": stat.Source,
|
||||
"format": stat.Format,
|
||||
"chat_id": stat.ChatID,
|
||||
"user_id": stat.UserID,
|
||||
"dpi": stat.DPI,
|
||||
"expression_len": stat.ExpressionLen,
|
||||
"ok": stat.OK,
|
||||
"error": stat.Error,
|
||||
"duration_ms": stat.DurationMs,
|
||||
})
|
||||
if err != nil {
|
||||
log.Printf("store: record render stat: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// FileRefs are correctly typed Telegram file_ids captured from normal chat
|
||||
// sends. They can be reused in inline results without another upload.
|
||||
type FileRefs struct {
|
||||
PhotoID string
|
||||
DocID string
|
||||
}
|
||||
|
||||
// FileRefs returns cached Telegram references for a rendered content hash.
|
||||
func (s *Store) FileRefs(hash string) FileRefs {
|
||||
if hash == "" {
|
||||
return FileRefs{}
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var doc struct {
|
||||
PhotoFileID string `bson:"photo_file_id"`
|
||||
DocFileID string `bson:"document_file_id"`
|
||||
}
|
||||
err := s.files.FindOne(ctx, bson.M{"_id": hash}).Decode(&doc)
|
||||
if err != nil {
|
||||
if err != mongo.ErrNoDocuments {
|
||||
log.Printf("store: read file refs for %s: %v", hash, err)
|
||||
}
|
||||
return FileRefs{}
|
||||
}
|
||||
return FileRefs{PhotoID: doc.PhotoFileID, DocID: doc.DocFileID}
|
||||
}
|
||||
|
||||
// SavePhotoID caches a Photo-type file_id for a content hash.
|
||||
func (s *Store) SavePhotoID(hash, fileID string) {
|
||||
s.saveFileID(hash, "photo_file_id", fileID)
|
||||
}
|
||||
|
||||
// SaveDocID caches a Document-type file_id for a content hash.
|
||||
func (s *Store) SaveDocID(hash, fileID string) {
|
||||
s.saveFileID(hash, "document_file_id", fileID)
|
||||
}
|
||||
|
||||
func (s *Store) saveFileID(hash, field, fileID string) {
|
||||
if hash == "" || fileID == "" {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
now := time.Now().UTC()
|
||||
_, err := s.files.UpdateOne(ctx,
|
||||
bson.M{"_id": hash},
|
||||
bson.M{
|
||||
"$set": bson.M{field: fileID, "updated_at": now},
|
||||
"$setOnInsert": bson.M{"created_at": now},
|
||||
},
|
||||
options.Update().SetUpsert(true),
|
||||
)
|
||||
if err != nil {
|
||||
log.Printf("store: save %s for %s: %v", field, hash, err)
|
||||
}
|
||||
}
|
||||
55
internal/store/store_test.go
Normal file
55
internal/store/store_test.go
Normal file
@@ -0,0 +1,55 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
// integrationStore connects to the cluster from MONGODB_URI (integration
|
||||
// tests only; CI runs without a database and skips them).
|
||||
func integrationStore(t *testing.T) *Store {
|
||||
t.Helper()
|
||||
uri := os.Getenv("MONGODB_URI")
|
||||
if uri == "" {
|
||||
t.Skip("MONGODB_URI not set; skipping MongoDB integration test")
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
client, err := mongo.Connect(ctx, options.Client().ApplyURI(uri))
|
||||
if err != nil {
|
||||
t.Fatalf("connect: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = client.Disconnect(context.Background()) })
|
||||
if err := client.Ping(ctx, nil); err != nil {
|
||||
t.Skipf("MongoDB not reachable, skipping integration test: %v", err)
|
||||
}
|
||||
return New(client, "latexgram_test_"+time.Now().Format("20060102150405"))
|
||||
}
|
||||
|
||||
func TestSettingsRoundTrip(t *testing.T) {
|
||||
s := integrationStore(t)
|
||||
if got := s.Get(42); got != (Settings{}) {
|
||||
t.Fatalf("unexpected initial settings: %+v", got)
|
||||
}
|
||||
settings := Settings{Preamble: `\usepackage{mhchem}`, DPI: 200}
|
||||
if err := s.Set(42, settings); err != nil {
|
||||
t.Fatalf("Set: %v", err)
|
||||
}
|
||||
if got := s.Get(42); got != settings {
|
||||
t.Fatalf("got %+v, want %+v", got, settings)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordUserAndGroup(t *testing.T) {
|
||||
s := integrationStore(t)
|
||||
userID := time.Now().UnixNano()
|
||||
s.RecordUser(User{ID: userID, Username: "someone", FirstName: "A", LastName: "B"})
|
||||
s.RecordUser(User{ID: userID})
|
||||
s.RecordGroup(Chat{ID: -1001, Type: "supergroup", Title: "Test", Username: "test"}, true)
|
||||
s.RecordGroup(Chat{ID: -1001}, false)
|
||||
}
|
||||
65
internal/telemetry/telemetry.go
Normal file
65
internal/telemetry/telemetry.go
Normal file
@@ -0,0 +1,65 @@
|
||||
// Package telemetry wires the OpenTelemetry SDK to an OTLP collector
|
||||
// (OTEL_EXPORTER_OTLP_ENDPOINT). With no endpoint configured the global
|
||||
// tracer provider stays a no-op.
|
||||
package telemetry
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"go.opentelemetry.io/otel"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
|
||||
"go.opentelemetry.io/otel/propagation"
|
||||
"go.opentelemetry.io/otel/sdk/resource"
|
||||
sdktrace "go.opentelemetry.io/otel/sdk/trace"
|
||||
)
|
||||
|
||||
// Setup installs the global tracer provider backed by the OTLP collector.
|
||||
// The returned shutdown flushes pending spans; call it on process exit.
|
||||
func Setup(ctx context.Context) (shutdown func(context.Context) error, err error) {
|
||||
endpoint := strings.TrimSpace(os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT"))
|
||||
if endpoint == "" {
|
||||
otel.SetTracerProvider(sdktrace.NewTracerProvider())
|
||||
return func(context.Context) error { return nil }, nil
|
||||
}
|
||||
insecure, parseErr := strconv.ParseBool(getenv("OTEL_EXPORTER_OTLP_INSECURE", "true"))
|
||||
if parseErr != nil {
|
||||
return nil, fmt.Errorf("invalid OTEL_EXPORTER_OTLP_INSECURE: %w", parseErr)
|
||||
}
|
||||
options := []otlptracegrpc.Option{otlptracegrpc.WithEndpoint(endpoint)}
|
||||
if insecure {
|
||||
options = append(options, otlptracegrpc.WithInsecure())
|
||||
}
|
||||
exporter, err := otlptracegrpc.New(ctx, options...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create OTLP exporter: %w", err)
|
||||
}
|
||||
res, err := resource.New(ctx,
|
||||
resource.WithFromEnv(),
|
||||
resource.WithAttributes(attribute.String("service.name", getenv("OTEL_SERVICE_NAME", "latexgrambot"))),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create resource: %w", err)
|
||||
}
|
||||
provider := sdktrace.NewTracerProvider(
|
||||
sdktrace.WithBatcher(exporter),
|
||||
sdktrace.WithResource(res),
|
||||
)
|
||||
otel.SetTracerProvider(provider)
|
||||
otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(
|
||||
propagation.TraceContext{},
|
||||
propagation.Baggage{},
|
||||
))
|
||||
return provider.Shutdown, nil
|
||||
}
|
||||
|
||||
func getenv(key, fallback string) string {
|
||||
if value := strings.TrimSpace(os.Getenv(key)); value != "" {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
Reference in New Issue
Block a user