chore: import upstream snapshot with attribution
Backend Tests / Tests (other) (push) Failing after 1s
Backend Tests / Static Checks (push) Failing after 2s
Backend Tests / Tests (internal) (push) Failing after 2s
Backend Tests / Tests (server) (push) Failing after 1s
Backend Tests / Tests (store) (push) Failing after 1s
Build Canary Image / build-frontend (push) Failing after 1s
Frontend Tests / Lint (push) Failing after 1s
Build Canary Image / build-push (linux/amd64) (push) Has been skipped
Build Canary Image / build-push (linux/arm64) (push) Has been skipped
Build Canary Image / merge (push) Has been skipped
Frontend Tests / Build (push) Failing after 1s
Release Please / release-please (push) Failing after 0s
Proto Linter / Lint Protos (push) Failing after 2s
Backend Tests / Tests (other) (push) Failing after 1s
Backend Tests / Static Checks (push) Failing after 2s
Backend Tests / Tests (internal) (push) Failing after 2s
Backend Tests / Tests (server) (push) Failing after 1s
Backend Tests / Tests (store) (push) Failing after 1s
Build Canary Image / build-frontend (push) Failing after 1s
Frontend Tests / Lint (push) Failing after 1s
Build Canary Image / build-push (linux/amd64) (push) Has been skipped
Build Canary Image / build-push (linux/arm64) (push) Has been skipped
Build Canary Image / merge (push) Has been skipped
Frontend Tests / Build (push) Failing after 1s
Release Please / release-please (push) Failing after 0s
Proto Linter / Lint Protos (push) Failing after 2s
This commit is contained in:
@@ -0,0 +1,20 @@
|
||||
package ai
|
||||
|
||||
// ProviderType identifies an AI provider implementation.
|
||||
type ProviderType string
|
||||
|
||||
const (
|
||||
// ProviderOpenAI is OpenAI's hosted API.
|
||||
ProviderOpenAI ProviderType = "OPENAI"
|
||||
// ProviderGemini is Google's Gemini API.
|
||||
ProviderGemini ProviderType = "GEMINI"
|
||||
)
|
||||
|
||||
// ProviderConfig configures a callable AI provider connection.
|
||||
type ProviderConfig struct {
|
||||
ID string
|
||||
Title string
|
||||
Type ProviderType
|
||||
Endpoint string
|
||||
APIKey string
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
// Package audio provides audio container/codec helpers for AI providers.
|
||||
//
|
||||
// The motivating use case is Gemini transcription: Gemini's audio inputs
|
||||
// require WAV/MP3/AIFF/AAC/OGG/FLAC, but browser MediaRecorder defaults to
|
||||
// WebM/Opus. This package converts WebM/Opus into 16-bit PCM WAV using
|
||||
// pure-Go decoders — no ffmpeg or other system dependency.
|
||||
package audio
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/at-wat/ebml-go"
|
||||
"github.com/at-wat/ebml-go/webm"
|
||||
"github.com/pion/opus"
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
const (
|
||||
opusOutputSampleRate = 48000
|
||||
// maxOpusPacketSamples is Opus's spec maximum: 120 ms at 48 kHz.
|
||||
maxOpusPacketSamples = 5760
|
||||
// opusCodecID is the WebM TrackEntry CodecID for an Opus audio track.
|
||||
opusCodecID = "A_OPUS"
|
||||
// opusHeadMinLength is the minimum size of the OpusHead identification
|
||||
// header stored in TrackEntry.CodecPrivate.
|
||||
opusHeadMinLength = 19
|
||||
)
|
||||
|
||||
// WebMOpusToWAV decodes a WebM/Opus file into 16-bit PCM WAV bytes.
|
||||
//
|
||||
// The output is mono or stereo at 48 kHz (Opus's native decode rate),
|
||||
// regardless of the original encoder's hint. Pre-skip samples declared in
|
||||
// the OpusHead are discarded to avoid the encoder's startup padding.
|
||||
//
|
||||
// The function reads the entire WebM document into memory; callers should
|
||||
// enforce their own size limits before invoking it.
|
||||
func WebMOpusToWAV(input []byte) ([]byte, error) {
|
||||
var doc struct {
|
||||
Header webm.EBMLHeader `ebml:"EBML"`
|
||||
Segment webm.Segment `ebml:"Segment"`
|
||||
}
|
||||
if err := ebml.Unmarshal(bytes.NewReader(input), &doc); err != nil && !errors.Is(err, io.EOF) {
|
||||
return nil, errors.Wrap(err, "parse webm")
|
||||
}
|
||||
|
||||
track := findOpusTrack(doc.Segment.Tracks.TrackEntry)
|
||||
if track == nil {
|
||||
return nil, errors.New("webm has no Opus audio track")
|
||||
}
|
||||
if len(track.CodecPrivate) < opusHeadMinLength {
|
||||
return nil, errors.Errorf("invalid OpusHead: expected at least %d bytes, got %d", opusHeadMinLength, len(track.CodecPrivate))
|
||||
}
|
||||
|
||||
channels := int(track.Audio.Channels)
|
||||
if channels < 1 || channels > 2 {
|
||||
return nil, errors.Errorf("unsupported Opus channel count: %d", channels)
|
||||
}
|
||||
preSkip := int(binary.LittleEndian.Uint16(track.CodecPrivate[10:12]))
|
||||
|
||||
decoder := opus.NewDecoder()
|
||||
if err := decoder.Init(opusOutputSampleRate, channels); err != nil {
|
||||
return nil, errors.Wrap(err, "init opus decoder")
|
||||
}
|
||||
|
||||
pcm := make([]int16, 0, 1<<16)
|
||||
frame := make([]int16, maxOpusPacketSamples*channels)
|
||||
|
||||
decodeBlock := func(block ebml.Block) error {
|
||||
if block.TrackNumber != track.TrackNumber {
|
||||
return nil
|
||||
}
|
||||
for _, packet := range block.Data {
|
||||
if len(packet) == 0 {
|
||||
continue
|
||||
}
|
||||
n, err := decoder.DecodeToInt16(packet, frame)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "decode opus packet")
|
||||
}
|
||||
pcm = append(pcm, frame[:n*channels]...)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, cluster := range doc.Segment.Cluster {
|
||||
for _, sb := range cluster.SimpleBlock {
|
||||
if err := decodeBlock(sb); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
for _, bg := range cluster.BlockGroup {
|
||||
if err := decodeBlock(bg.Block); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
skip := min(preSkip*channels, len(pcm))
|
||||
pcm = pcm[skip:]
|
||||
|
||||
return encodeWAV(pcm, opusOutputSampleRate, channels), nil
|
||||
}
|
||||
|
||||
// IsWebMContentType reports whether the MIME type is WebM audio.
|
||||
// Both "audio/webm" and "audio/webm; codecs=opus" return true.
|
||||
func IsWebMContentType(contentType string) bool {
|
||||
contentType = strings.TrimSpace(contentType)
|
||||
if contentType == "" {
|
||||
return false
|
||||
}
|
||||
if i := strings.IndexByte(contentType, ';'); i >= 0 {
|
||||
contentType = contentType[:i]
|
||||
}
|
||||
return strings.EqualFold(strings.TrimSpace(contentType), "audio/webm")
|
||||
}
|
||||
|
||||
func findOpusTrack(entries []webm.TrackEntry) *webm.TrackEntry {
|
||||
for i := range entries {
|
||||
entry := &entries[i]
|
||||
if entry.CodecID == opusCodecID && entry.Audio != nil {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// encodeWAV writes a standard RIFF/WAVE container around 16-bit PCM samples.
|
||||
// Reference layout: http://soundfile.sapp.org/doc/WaveFormat/
|
||||
func encodeWAV(samples []int16, sampleRate, channels int) []byte {
|
||||
const bitsPerSample = 16
|
||||
const bytesPerSample = bitsPerSample / 8
|
||||
blockAlign := channels * bytesPerSample
|
||||
byteRate := sampleRate * blockAlign
|
||||
dataSize := len(samples) * bytesPerSample
|
||||
|
||||
buf := bytes.NewBuffer(make([]byte, 0, 44+dataSize))
|
||||
buf.WriteString("RIFF")
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint32(36+dataSize))
|
||||
buf.WriteString("WAVE")
|
||||
|
||||
buf.WriteString("fmt ")
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint32(16))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint16(1)) // PCM
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint16(channels))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint32(sampleRate))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint32(byteRate))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint16(blockAlign))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint16(bitsPerSample))
|
||||
|
||||
buf.WriteString("data")
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint32(dataSize))
|
||||
_ = binary.Write(buf, binary.LittleEndian, samples)
|
||||
|
||||
return buf.Bytes()
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package audio
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestIsWebMContentType(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
want bool
|
||||
}{
|
||||
{"audio/webm", true},
|
||||
{"audio/webm;codecs=opus", true},
|
||||
{"audio/webm; codecs=opus", true},
|
||||
{"AUDIO/WEBM", true},
|
||||
{" audio/webm ", true},
|
||||
{"audio/wav", false},
|
||||
{"audio/mp4", false},
|
||||
{"video/webm", false},
|
||||
{"", false},
|
||||
{"webm", false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.in, func(t *testing.T) {
|
||||
require.Equal(t, tc.want, IsWebMContentType(tc.in))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebMOpusToWAV_RejectsInvalidInput(t *testing.T) {
|
||||
t.Run("empty", func(t *testing.T) {
|
||||
_, err := WebMOpusToWAV(nil)
|
||||
require.Error(t, err)
|
||||
})
|
||||
|
||||
t.Run("not webm", func(t *testing.T) {
|
||||
_, err := WebMOpusToWAV([]byte("hello world this is not webm"))
|
||||
require.Error(t, err)
|
||||
})
|
||||
|
||||
t.Run("truncated webm header bytes", func(t *testing.T) {
|
||||
// Valid EBML magic but no Segment.
|
||||
_, err := WebMOpusToWAV([]byte{0x1A, 0x45, 0xDF, 0xA3})
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
// Package audiollm defines the multimodal-audio capability for AI providers.
|
||||
// Implementations call chat-completions or generate-content style APIs that
|
||||
// accept audio as input. For deterministic transcription, prefer internal/ai/stt
|
||||
// where a dedicated STT endpoint exists.
|
||||
package audiollm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
)
|
||||
|
||||
// Model invokes a multimodal LLM with audio input.
|
||||
type Model interface {
|
||||
GenerateFromAudio(ctx context.Context, req Request) (*Response, error)
|
||||
}
|
||||
|
||||
// Request is the input to a multimodal-audio call.
|
||||
type Request struct {
|
||||
Audio io.Reader
|
||||
Size int64
|
||||
ContentType string
|
||||
Model string
|
||||
Instructions string // literal instruction the model is expected to follow
|
||||
Temperature *float32 // optional; nil leaves the provider default in place
|
||||
}
|
||||
|
||||
// Response is the output of a multimodal-audio call.
|
||||
type Response struct {
|
||||
Text string
|
||||
FinishReason FinishReason
|
||||
}
|
||||
|
||||
// FinishReason describes why the model stopped generating.
|
||||
type FinishReason string
|
||||
|
||||
const (
|
||||
FinishStop FinishReason = "stop" // model finished normally
|
||||
FinishLength FinishReason = "length" // truncated by max-tokens
|
||||
FinishSafety FinishReason = "safety" // safety filter blocked output
|
||||
FinishOther FinishReason = "other" // anything else, including unknown
|
||||
)
|
||||
@@ -0,0 +1,202 @@
|
||||
// Package gemini implements audiollm.Model against the Gemini generateContent
|
||||
// endpoint. Used by Memos transcription when the user picks a Gemini provider:
|
||||
// the handler issues a transcription instruction via audiollm.Request.Instructions.
|
||||
package gemini
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"mime"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
"google.golang.org/genai"
|
||||
|
||||
"github.com/usememos/memos/internal/ai"
|
||||
"github.com/usememos/memos/internal/ai/audio"
|
||||
"github.com/usememos/memos/internal/ai/audiollm"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultEndpoint = "https://generativelanguage.googleapis.com/v1beta"
|
||||
defaultAPIVersion = "v1beta"
|
||||
maxInlineSize = 14 * 1024 * 1024
|
||||
providerName = "Gemini"
|
||||
)
|
||||
|
||||
var supportedContentTypes = map[string]string{
|
||||
"audio/wav": "audio/wav",
|
||||
"audio/x-wav": "audio/wav",
|
||||
"audio/mp3": "audio/mp3",
|
||||
"audio/mpeg": "audio/mp3",
|
||||
"audio/aiff": "audio/aiff",
|
||||
"audio/aac": "audio/aac",
|
||||
"audio/ogg": "audio/ogg",
|
||||
"audio/flac": "audio/flac",
|
||||
"audio/x-flac": "audio/flac",
|
||||
}
|
||||
|
||||
// Model implements audiollm.Model for Gemini generateContent.
|
||||
type Model struct {
|
||||
client *genai.Client
|
||||
}
|
||||
|
||||
// New constructs a Model from a provider config.
|
||||
func New(cfg ai.ProviderConfig, options audiollm.Options) (*Model, error) {
|
||||
endpoint, err := normalizeEndpoint(cfg.Endpoint)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cfg.APIKey == "" {
|
||||
return nil, errors.Errorf("%s API key is required", providerName)
|
||||
}
|
||||
baseURL, apiVersion, err := splitEndpoint(endpoint)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
httpOptions := genai.HTTPOptions{BaseURL: baseURL, APIVersion: apiVersion}
|
||||
if options.HTTPClient != nil && options.HTTPClient.Timeout > 0 {
|
||||
timeout := options.HTTPClient.Timeout
|
||||
httpOptions.Timeout = &timeout
|
||||
}
|
||||
client, err := genai.NewClient(context.Background(), &genai.ClientConfig{
|
||||
APIKey: cfg.APIKey,
|
||||
Backend: genai.BackendGeminiAPI,
|
||||
HTTPClient: options.HTTPClient,
|
||||
HTTPOptions: httpOptions,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to create Gemini client")
|
||||
}
|
||||
return &Model{client: client}, nil
|
||||
}
|
||||
|
||||
// GenerateFromAudio calls Gemini generateContent with the audio attached.
|
||||
func (m *Model) GenerateFromAudio(ctx context.Context, req audiollm.Request) (*audiollm.Response, error) {
|
||||
if strings.TrimSpace(req.Model) == "" {
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
if req.Audio == nil {
|
||||
return nil, errors.New("audio is required")
|
||||
}
|
||||
if strings.TrimSpace(req.Instructions) == "" {
|
||||
return nil, errors.New("instructions are required")
|
||||
}
|
||||
|
||||
audioBytes, err := io.ReadAll(req.Audio)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to read audio")
|
||||
}
|
||||
if len(audioBytes) == 0 {
|
||||
return nil, errors.New("audio is required")
|
||||
}
|
||||
|
||||
contentType := req.ContentType
|
||||
if audio.IsWebMContentType(contentType) {
|
||||
wav, err := audio.WebMOpusToWAV(audioBytes)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to transcode webm audio for Gemini")
|
||||
}
|
||||
audioBytes = wav
|
||||
contentType = "audio/wav"
|
||||
}
|
||||
|
||||
if len(audioBytes) > maxInlineSize {
|
||||
return nil, errors.Errorf("audio is too large for Gemini inline request; maximum size is %d bytes", maxInlineSize)
|
||||
}
|
||||
|
||||
contentType, err = normalizeContentType(contentType)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cfg := &genai.GenerateContentConfig{}
|
||||
if req.Temperature != nil {
|
||||
t := *req.Temperature
|
||||
cfg.Temperature = &t
|
||||
}
|
||||
|
||||
resp, err := m.client.Models.GenerateContent(ctx, normalizeModelName(req.Model), []*genai.Content{
|
||||
genai.NewContentFromParts([]*genai.Part{
|
||||
genai.NewPartFromBytes(audioBytes, contentType),
|
||||
genai.NewPartFromText(req.Instructions),
|
||||
}, genai.RoleUser),
|
||||
}, cfg)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to send Gemini request")
|
||||
}
|
||||
|
||||
return &audiollm.Response{
|
||||
Text: strings.TrimSpace(resp.Text()),
|
||||
FinishReason: mapFinishReason(resp),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func mapFinishReason(resp *genai.GenerateContentResponse) audiollm.FinishReason {
|
||||
if resp == nil || len(resp.Candidates) == 0 {
|
||||
return audiollm.FinishOther
|
||||
}
|
||||
switch resp.Candidates[0].FinishReason {
|
||||
case genai.FinishReasonStop:
|
||||
return audiollm.FinishStop
|
||||
case genai.FinishReasonMaxTokens:
|
||||
return audiollm.FinishLength
|
||||
case genai.FinishReasonSafety,
|
||||
genai.FinishReasonRecitation,
|
||||
genai.FinishReasonProhibitedContent,
|
||||
genai.FinishReasonSPII,
|
||||
genai.FinishReasonBlocklist,
|
||||
genai.FinishReasonImageSafety,
|
||||
genai.FinishReasonImageProhibitedContent,
|
||||
genai.FinishReasonImageRecitation:
|
||||
return audiollm.FinishSafety
|
||||
default:
|
||||
return audiollm.FinishOther
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeEndpoint(endpoint string) (string, error) {
|
||||
endpoint = strings.TrimSpace(endpoint)
|
||||
if endpoint == "" {
|
||||
endpoint = defaultEndpoint
|
||||
}
|
||||
if _, err := url.ParseRequestURI(endpoint); err != nil {
|
||||
return "", errors.Wrapf(err, "invalid %s endpoint", providerName)
|
||||
}
|
||||
return strings.TrimRight(endpoint, "/"), nil
|
||||
}
|
||||
|
||||
func splitEndpoint(endpoint string) (string, string, error) {
|
||||
parsed, err := url.Parse(endpoint)
|
||||
if err != nil {
|
||||
return "", "", errors.Wrap(err, "invalid Gemini endpoint")
|
||||
}
|
||||
path := strings.TrimRight(parsed.Path, "/")
|
||||
apiVersion := defaultAPIVersion
|
||||
for _, supported := range []string{"v1alpha", "v1beta", "v1"} {
|
||||
if path == "/"+supported || strings.HasSuffix(path, "/"+supported) {
|
||||
apiVersion = supported
|
||||
parsed.Path = strings.TrimSuffix(path, "/"+supported)
|
||||
break
|
||||
}
|
||||
}
|
||||
return strings.TrimRight(parsed.String(), "/"), apiVersion, nil
|
||||
}
|
||||
|
||||
func normalizeContentType(contentType string) (string, error) {
|
||||
mediaType, _, err := mime.ParseMediaType(strings.TrimSpace(contentType))
|
||||
if err != nil {
|
||||
return "", errors.Wrap(err, "invalid audio content type")
|
||||
}
|
||||
mediaType = strings.ToLower(mediaType)
|
||||
normalized, ok := supportedContentTypes[mediaType]
|
||||
if !ok {
|
||||
return "", errors.Errorf("audio content type %q is not supported by Gemini", mediaType)
|
||||
}
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
func normalizeModelName(model string) string {
|
||||
return strings.TrimPrefix(strings.TrimSpace(model), "models/")
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
package gemini_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/usememos/memos/internal/ai"
|
||||
"github.com/usememos/memos/internal/ai/audiollm"
|
||||
audiollmgemini "github.com/usememos/memos/internal/ai/audiollm/gemini"
|
||||
)
|
||||
|
||||
func TestGenerateFromAudio(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, http.MethodPost, r.Method)
|
||||
require.Equal(t, "/v1beta/models/gemini-2.5-flash:generateContent", r.URL.Path)
|
||||
require.Equal(t, "test-key", r.Header.Get("x-goog-api-key"))
|
||||
require.Equal(t, "application/json", r.Header.Get("Content-Type"))
|
||||
|
||||
var request struct {
|
||||
Contents []struct {
|
||||
Parts []struct {
|
||||
Text string `json:"text"`
|
||||
InlineData *struct {
|
||||
MIMEType string `json:"mimeType"`
|
||||
Data string `json:"data"`
|
||||
} `json:"inlineData"`
|
||||
} `json:"parts"`
|
||||
} `json:"contents"`
|
||||
GenerationConfig map[string]json.Number `json:"generationConfig"`
|
||||
}
|
||||
require.NoError(t, json.NewDecoder(r.Body).Decode(&request))
|
||||
require.Len(t, request.Contents, 1)
|
||||
require.Len(t, request.Contents[0].Parts, 2)
|
||||
require.NotNil(t, request.Contents[0].Parts[0].InlineData)
|
||||
require.Equal(t, "audio/mp3", request.Contents[0].Parts[0].InlineData.MIMEType)
|
||||
audio, err := base64.StdEncoding.DecodeString(request.Contents[0].Parts[0].InlineData.Data)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "audio bytes", string(audio))
|
||||
require.Equal(t, "transcribe please", request.Contents[0].Parts[1].Text)
|
||||
require.Equal(t, json.Number("0"), request.GenerationConfig["temperature"])
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
require.NoError(t, json.NewEncoder(w).Encode(map[string]any{
|
||||
"candidates": []map[string]any{
|
||||
{
|
||||
"finishReason": "STOP",
|
||||
"content": map[string]any{
|
||||
"parts": []map[string]string{{"text": "hello from gemini"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
model, err := audiollmgemini.New(ai.ProviderConfig{
|
||||
Type: ai.ProviderGemini,
|
||||
Endpoint: server.URL + "/v1beta",
|
||||
APIKey: "test-key",
|
||||
}, audiollm.ApplyOptions(nil))
|
||||
require.NoError(t, err)
|
||||
|
||||
temp := float32(0)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
resp, err := model.GenerateFromAudio(ctx, audiollm.Request{
|
||||
Model: "models/gemini-2.5-flash",
|
||||
ContentType: "audio/mpeg",
|
||||
Audio: strings.NewReader("audio bytes"),
|
||||
Instructions: "transcribe please",
|
||||
Temperature: &temp,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "hello from gemini", resp.Text)
|
||||
require.Equal(t, audiollm.FinishStop, resp.FinishReason)
|
||||
}
|
||||
|
||||
func TestGenerateFromAudioRejectsUnsupportedContentType(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
model, err := audiollmgemini.New(ai.ProviderConfig{
|
||||
Type: ai.ProviderGemini,
|
||||
Endpoint: "https://example.com/v1beta",
|
||||
APIKey: "test-key",
|
||||
}, audiollm.ApplyOptions(nil))
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = model.GenerateFromAudio(context.Background(), audiollm.Request{
|
||||
Model: "gemini-2.5-flash",
|
||||
ContentType: "video/mp4",
|
||||
Audio: strings.NewReader("video bytes"),
|
||||
Instructions: "transcribe please",
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "not supported by Gemini")
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package audiollm
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
const defaultHTTPTimeout = 2 * time.Minute
|
||||
|
||||
// Options is the resolved option set passed to provider implementations.
|
||||
type Options struct {
|
||||
HTTPClient *http.Client
|
||||
}
|
||||
|
||||
// ModelOption customizes a Model.
|
||||
type ModelOption func(*Options)
|
||||
|
||||
// WithHTTPClient overrides the HTTP client used by the model.
|
||||
func WithHTTPClient(client *http.Client) ModelOption {
|
||||
return func(o *Options) {
|
||||
if client != nil {
|
||||
o.HTTPClient = client
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ApplyOptions resolves a ModelOption slice into Options with defaults.
|
||||
func ApplyOptions(opts []ModelOption) Options {
|
||||
resolved := Options{HTTPClient: &http.Client{Timeout: defaultHTTPTimeout}}
|
||||
for _, apply := range opts {
|
||||
apply(&resolved)
|
||||
}
|
||||
return resolved
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package ai
|
||||
|
||||
import "github.com/pkg/errors"
|
||||
|
||||
var (
|
||||
// ErrProviderNotFound indicates that a requested provider ID does not exist.
|
||||
ErrProviderNotFound = errors.New("AI provider not found")
|
||||
// ErrCapabilityUnsupported indicates that the provider does not support the requested capability.
|
||||
ErrCapabilityUnsupported = errors.New("AI provider capability unsupported")
|
||||
// ErrSTTNotSupported indicates that the provider does not have a dedicated
|
||||
// speech-to-text endpoint. Use the audiollm package for multimodal audio
|
||||
// understanding when this is returned.
|
||||
ErrSTTNotSupported = errors.New("provider does not support speech-to-text capability")
|
||||
// ErrAudioLLMNotSupported indicates that the provider does not have a
|
||||
// multimodal-audio LLM available in this codebase.
|
||||
ErrAudioLLMNotSupported = errors.New("provider does not support multimodal audio capability")
|
||||
)
|
||||
@@ -0,0 +1,22 @@
|
||||
package ai
|
||||
|
||||
import "github.com/pkg/errors"
|
||||
|
||||
const (
|
||||
// DefaultOpenAITranscriptionModel is the built-in OpenAI transcription model.
|
||||
DefaultOpenAITranscriptionModel = "whisper-1"
|
||||
// DefaultGeminiTranscriptionModel is the built-in Gemini transcription model.
|
||||
DefaultGeminiTranscriptionModel = "gemini-2.5-flash"
|
||||
)
|
||||
|
||||
// DefaultTranscriptionModel returns the built-in transcription model for a provider.
|
||||
func DefaultTranscriptionModel(providerType ProviderType) (string, error) {
|
||||
switch providerType {
|
||||
case ProviderOpenAI:
|
||||
return DefaultOpenAITranscriptionModel, nil
|
||||
case ProviderGemini:
|
||||
return DefaultGeminiTranscriptionModel, nil
|
||||
default:
|
||||
return "", errors.Wrapf(ErrCapabilityUnsupported, "provider type %q", providerType)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package ai
|
||||
|
||||
import "github.com/pkg/errors"
|
||||
|
||||
// FindProvider returns the provider with the given ID.
|
||||
func FindProvider(providers []ProviderConfig, providerID string) (*ProviderConfig, error) {
|
||||
if providerID == "" {
|
||||
return nil, errors.Wrap(ErrProviderNotFound, "provider ID is required")
|
||||
}
|
||||
for _, provider := range providers {
|
||||
if provider.ID == providerID {
|
||||
return &provider, nil
|
||||
}
|
||||
}
|
||||
return nil, errors.Wrapf(ErrProviderNotFound, "provider ID %q", providerID)
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
// Package openai implements stt.Transcriber against the OpenAI
|
||||
// /audio/transcriptions endpoint (and any compatible third-party endpoint
|
||||
// such as Groq Whisper, faster-whisper self-hosted, or Azure Whisper).
|
||||
package openai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"mime"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
openaisdk "github.com/openai/openai-go/v3"
|
||||
openaioption "github.com/openai/openai-go/v3/option"
|
||||
"github.com/pkg/errors"
|
||||
|
||||
"github.com/usememos/memos/internal/ai"
|
||||
"github.com/usememos/memos/internal/ai/stt"
|
||||
)
|
||||
|
||||
const defaultEndpoint = "https://api.openai.com/v1"
|
||||
|
||||
// Transcriber implements stt.Transcriber for OpenAI-compatible STT endpoints.
|
||||
type Transcriber struct {
|
||||
client openaisdk.Client
|
||||
}
|
||||
|
||||
// New constructs a Transcriber from a provider config.
|
||||
func New(cfg ai.ProviderConfig, options stt.Options) (*Transcriber, error) {
|
||||
endpoint, err := normalizeEndpoint(cfg.Endpoint)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cfg.APIKey == "" {
|
||||
return nil, errors.New("OpenAI API key is required")
|
||||
}
|
||||
return &Transcriber{
|
||||
client: openaisdk.NewClient(
|
||||
openaioption.WithAPIKey(cfg.APIKey),
|
||||
openaioption.WithBaseURL(endpoint),
|
||||
openaioption.WithHTTPClient(options.HTTPClient),
|
||||
),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Transcribe sends the audio to /audio/transcriptions.
|
||||
func (t *Transcriber) Transcribe(ctx context.Context, req stt.Request) (*stt.Response, error) {
|
||||
if strings.TrimSpace(req.Model) == "" {
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
if req.Audio == nil {
|
||||
return nil, errors.New("audio is required")
|
||||
}
|
||||
|
||||
filename, contentType, err := normalizeAudioMetadata(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
params := openaisdk.AudioTranscriptionNewParams{
|
||||
File: openaisdk.File(req.Audio, filename, contentType),
|
||||
Model: openaisdk.AudioModel(req.Model),
|
||||
ResponseFormat: openaisdk.AudioResponseFormatJSON,
|
||||
}
|
||||
if req.Prompt != "" {
|
||||
params.Prompt = openaisdk.String(req.Prompt)
|
||||
}
|
||||
if req.Language != "" {
|
||||
params.Language = openaisdk.String(req.Language)
|
||||
}
|
||||
|
||||
resp, err := t.client.Audio.Transcriptions.New(ctx, params)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to send OpenAI transcription request")
|
||||
}
|
||||
return &stt.Response{
|
||||
Text: resp.Text,
|
||||
Language: resp.Language,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func normalizeEndpoint(endpoint string) (string, error) {
|
||||
endpoint = strings.TrimSpace(endpoint)
|
||||
if endpoint == "" {
|
||||
endpoint = defaultEndpoint
|
||||
}
|
||||
if _, err := url.ParseRequestURI(endpoint); err != nil {
|
||||
return "", errors.Wrap(err, "invalid OpenAI endpoint")
|
||||
}
|
||||
return strings.TrimRight(endpoint, "/"), nil
|
||||
}
|
||||
|
||||
func normalizeAudioMetadata(req stt.Request) (string, string, error) {
|
||||
filename := strings.TrimSpace(req.Filename)
|
||||
if filename == "" {
|
||||
filename = "audio"
|
||||
}
|
||||
contentType := strings.TrimSpace(req.ContentType)
|
||||
if contentType == "" {
|
||||
contentType = "application/octet-stream"
|
||||
} else {
|
||||
mediaType, _, err := mime.ParseMediaType(contentType)
|
||||
if err != nil {
|
||||
return "", "", errors.Wrap(err, "invalid audio content type")
|
||||
}
|
||||
contentType = mediaType
|
||||
}
|
||||
return sanitizeFilename(filename), contentType, nil
|
||||
}
|
||||
|
||||
func sanitizeFilename(filename string) string {
|
||||
filename = strings.NewReplacer("\r", "_", "\n", "_").Replace(filename)
|
||||
if strings.TrimSpace(filename) == "" {
|
||||
return "audio"
|
||||
}
|
||||
return filename
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package openai_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/usememos/memos/internal/ai"
|
||||
"github.com/usememos/memos/internal/ai/stt"
|
||||
sttopenai "github.com/usememos/memos/internal/ai/stt/openai"
|
||||
)
|
||||
|
||||
func TestTranscribe(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, http.MethodPost, r.Method)
|
||||
require.Equal(t, "/audio/transcriptions", r.URL.Path)
|
||||
require.Equal(t, "Bearer test-key", r.Header.Get("Authorization"))
|
||||
require.NoError(t, r.ParseMultipartForm(10<<20))
|
||||
require.Equal(t, "gpt-4o-transcribe", r.FormValue("model"))
|
||||
require.Equal(t, "json", r.FormValue("response_format"))
|
||||
require.Equal(t, "domain words", r.FormValue("prompt"))
|
||||
require.Equal(t, "en", r.FormValue("language"))
|
||||
|
||||
file, header, err := r.FormFile("file")
|
||||
require.NoError(t, err)
|
||||
defer file.Close()
|
||||
require.Equal(t, "voice.wav", header.Filename)
|
||||
require.Equal(t, "audio/wav", header.Header.Get("Content-Type"))
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
require.NoError(t, json.NewEncoder(w).Encode(map[string]any{
|
||||
"text": "hello world",
|
||||
"language": "en",
|
||||
"duration": 1.5,
|
||||
}))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
transcriber, err := sttopenai.New(ai.ProviderConfig{
|
||||
Type: ai.ProviderOpenAI,
|
||||
Endpoint: server.URL,
|
||||
APIKey: "test-key",
|
||||
}, stt.ApplyOptions(nil))
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
response, err := transcriber.Transcribe(ctx, stt.Request{
|
||||
Model: "gpt-4o-transcribe",
|
||||
Filename: "voice.wav",
|
||||
ContentType: "audio/wav",
|
||||
Audio: strings.NewReader("RIFF"),
|
||||
Prompt: "domain words",
|
||||
Language: "en",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "hello world", response.Text)
|
||||
require.Equal(t, "en", response.Language)
|
||||
// Note: Duration intentionally omitted from stt.Response — not exposed in the new contract.
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package stt
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
const defaultHTTPTimeout = 2 * time.Minute
|
||||
|
||||
// Options is the resolved option set passed to provider implementations.
|
||||
type Options struct {
|
||||
HTTPClient *http.Client
|
||||
}
|
||||
|
||||
// TranscriberOption customizes a Transcriber.
|
||||
type TranscriberOption func(*Options)
|
||||
|
||||
// WithHTTPClient overrides the HTTP client used by the transcriber.
|
||||
func WithHTTPClient(client *http.Client) TranscriberOption {
|
||||
return func(o *Options) {
|
||||
if client != nil {
|
||||
o.HTTPClient = client
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ApplyOptions resolves a TranscriberOption slice into Options with defaults.
|
||||
func ApplyOptions(opts []TranscriberOption) Options {
|
||||
resolved := Options{HTTPClient: &http.Client{Timeout: defaultHTTPTimeout}}
|
||||
for _, apply := range opts {
|
||||
apply(&resolved)
|
||||
}
|
||||
return resolved
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
// Package stt defines the speech-to-text capability for AI providers.
|
||||
// Implementations call dedicated STT endpoints (e.g. OpenAI /audio/transcriptions)
|
||||
// and return deterministic transcription output. For multimodal LLMs that
|
||||
// happen to accept audio input, see internal/ai/audiollm.
|
||||
package stt
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
)
|
||||
|
||||
// Transcriber transcribes audio to text using a provider's dedicated STT endpoint.
|
||||
type Transcriber interface {
|
||||
Transcribe(ctx context.Context, req Request) (*Response, error)
|
||||
}
|
||||
|
||||
// Request is the input to a transcription call.
|
||||
type Request struct {
|
||||
Audio io.Reader
|
||||
Size int64
|
||||
Filename string
|
||||
ContentType string // IANA media type, e.g. "audio/wav"
|
||||
Model string // provider-specific model id (e.g. "whisper-1", "gpt-4o-transcribe")
|
||||
Prompt string // soft spelling/vocabulary hint (Whisper "prompt" parameter)
|
||||
Language string // ISO 639-1, optional
|
||||
}
|
||||
|
||||
// Response is the output of a transcription call.
|
||||
type Response struct {
|
||||
Text string
|
||||
Language string // empty if provider did not return it
|
||||
Segments []Segment // empty unless provider returned timestamps
|
||||
}
|
||||
|
||||
// Segment is a timestamped portion of the transcript.
|
||||
type Segment struct {
|
||||
Text string
|
||||
Start float64
|
||||
End float64
|
||||
Speaker string // empty unless using a diarization-capable model
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
package base
|
||||
|
||||
import "regexp"
|
||||
|
||||
var (
|
||||
UIDMatcher = regexp.MustCompile("^[a-zA-Z0-9]([a-zA-Z0-9-]{0,34}[a-zA-Z0-9])?$")
|
||||
)
|
||||
@@ -0,0 +1,37 @@
|
||||
package base
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestUIDMatcher(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
expected bool
|
||||
}{
|
||||
{"", false},
|
||||
{"-abc123", false},
|
||||
{"012345678901234567890123456789", true},
|
||||
{"1abc-123", true},
|
||||
{"A123B456C789", true},
|
||||
{"a", true},
|
||||
{"ab", true},
|
||||
{"a*b&c", false},
|
||||
{"a--b", true},
|
||||
{"a-1b-2c", true},
|
||||
{"a1234567890123456789012345678901", true},
|
||||
{"abc123", true},
|
||||
{"abc123-", false},
|
||||
{"123e4567-e89b-12d3-a456-426614174000", true}, // UUID v4 from IDP
|
||||
{"a123456789012345678901234567890123456", false}, // 37 characters (too long)
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.input, func(*testing.T) {
|
||||
result := UIDMatcher.MatchString(test.input)
|
||||
if result != test.expected {
|
||||
t.Errorf("For input '%s', expected %v but got %v", test.input, test.expected, result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,507 @@
|
||||
# Email Plugin
|
||||
|
||||
SMTP email sending functionality for self-hosted Memos instances.
|
||||
|
||||
## Overview
|
||||
|
||||
This plugin provides a simple, reliable email sending interface following industry-standard SMTP protocols. It's designed for self-hosted environments where instance administrators configure their own email service, similar to platforms like GitHub, GitLab, and Discourse.
|
||||
|
||||
## Features
|
||||
|
||||
- Standard SMTP protocol support
|
||||
- TLS/STARTTLS and SSL/TLS encryption
|
||||
- HTML and plain text emails
|
||||
- Multiple recipients (To, Cc, Bcc)
|
||||
- Synchronous and asynchronous sending
|
||||
- Detailed error reporting with context
|
||||
- Works with all major email providers
|
||||
- Reply-To header support
|
||||
- RFC 5322 compliant message formatting
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Configure SMTP Settings
|
||||
|
||||
```go
|
||||
import "github.com/usememos/memos/internal/email"
|
||||
|
||||
config := &email.Config{
|
||||
SMTPHost: "smtp.gmail.com",
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "your-email@gmail.com",
|
||||
SMTPPassword: "your-app-password",
|
||||
FromEmail: "noreply@yourdomain.com",
|
||||
FromName: "Memos",
|
||||
UseTLS: true,
|
||||
}
|
||||
```
|
||||
|
||||
### 2. Create and Send Email
|
||||
|
||||
```go
|
||||
message := &email.Message{
|
||||
To: []string{"user@example.com"},
|
||||
Subject: "Welcome to Memos!",
|
||||
Body: "Thanks for signing up.",
|
||||
IsHTML: false,
|
||||
}
|
||||
|
||||
// Synchronous send (waits for result)
|
||||
err := email.Send(config, message)
|
||||
if err != nil {
|
||||
log.Printf("Failed to send email: %v", err)
|
||||
}
|
||||
|
||||
// Asynchronous send (returns immediately)
|
||||
email.SendAsync(config, message)
|
||||
```
|
||||
|
||||
## Provider Configuration
|
||||
|
||||
### Gmail
|
||||
|
||||
Requires an [App Password](https://support.google.com/accounts/answer/185833) (2FA must be enabled):
|
||||
|
||||
```go
|
||||
config := &email.Config{
|
||||
SMTPHost: "smtp.gmail.com",
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "your-email@gmail.com",
|
||||
SMTPPassword: "your-16-char-app-password",
|
||||
FromEmail: "your-email@gmail.com",
|
||||
FromName: "Memos",
|
||||
UseTLS: true,
|
||||
}
|
||||
```
|
||||
|
||||
**Alternative (SSL):**
|
||||
```go
|
||||
config := &email.Config{
|
||||
SMTPHost: "smtp.gmail.com",
|
||||
SMTPPort: 465,
|
||||
SMTPUsername: "your-email@gmail.com",
|
||||
SMTPPassword: "your-16-char-app-password",
|
||||
FromEmail: "your-email@gmail.com",
|
||||
FromName: "Memos",
|
||||
UseSSL: true,
|
||||
}
|
||||
```
|
||||
|
||||
### SendGrid
|
||||
|
||||
```go
|
||||
config := &email.Config{
|
||||
SMTPHost: "smtp.sendgrid.net",
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "apikey",
|
||||
SMTPPassword: "your-sendgrid-api-key",
|
||||
FromEmail: "noreply@yourdomain.com",
|
||||
FromName: "Memos",
|
||||
UseTLS: true,
|
||||
}
|
||||
```
|
||||
|
||||
### AWS SES
|
||||
|
||||
```go
|
||||
config := &email.Config{
|
||||
SMTPHost: "email-smtp.us-east-1.amazonaws.com",
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "your-smtp-username",
|
||||
SMTPPassword: "your-smtp-password",
|
||||
FromEmail: "verified@yourdomain.com",
|
||||
FromName: "Memos",
|
||||
UseTLS: true,
|
||||
}
|
||||
```
|
||||
|
||||
**Note:** Replace `us-east-1` with your AWS region. Email must be verified in SES.
|
||||
|
||||
### Mailgun
|
||||
|
||||
```go
|
||||
config := &email.Config{
|
||||
SMTPHost: "smtp.mailgun.org",
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "postmaster@yourdomain.com",
|
||||
SMTPPassword: "your-mailgun-smtp-password",
|
||||
FromEmail: "noreply@yourdomain.com",
|
||||
FromName: "Memos",
|
||||
UseTLS: true,
|
||||
}
|
||||
```
|
||||
|
||||
### Self-Hosted SMTP (Postfix, Exim, etc.)
|
||||
|
||||
```go
|
||||
config := &email.Config{
|
||||
SMTPHost: "mail.yourdomain.com",
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "username",
|
||||
SMTPPassword: "password",
|
||||
FromEmail: "noreply@yourdomain.com",
|
||||
FromName: "Memos",
|
||||
UseTLS: true,
|
||||
}
|
||||
```
|
||||
|
||||
## HTML Emails
|
||||
|
||||
```go
|
||||
message := &email.Message{
|
||||
To: []string{"user@example.com"},
|
||||
Subject: "Welcome to Memos!",
|
||||
Body: `
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
</head>
|
||||
<body style="font-family: Arial, sans-serif;">
|
||||
<h1 style="color: #333;">Welcome to Memos!</h1>
|
||||
<p>We're excited to have you on board.</p>
|
||||
<a href="https://yourdomain.com" style="background-color: #4CAF50; color: white; padding: 10px 20px; text-decoration: none; border-radius: 5px;">Get Started</a>
|
||||
</body>
|
||||
</html>
|
||||
`,
|
||||
IsHTML: true,
|
||||
}
|
||||
|
||||
email.Send(config, message)
|
||||
```
|
||||
|
||||
## Multiple Recipients
|
||||
|
||||
```go
|
||||
message := &email.Message{
|
||||
To: []string{"user1@example.com", "user2@example.com"},
|
||||
Cc: []string{"manager@example.com"},
|
||||
Bcc: []string{"admin@example.com"},
|
||||
Subject: "Team Update",
|
||||
Body: "Important team announcement...",
|
||||
ReplyTo: "support@yourdomain.com",
|
||||
}
|
||||
|
||||
email.Send(config, message)
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
### Run Tests
|
||||
|
||||
```bash
|
||||
# All tests
|
||||
go test ./internal/email/... -v
|
||||
|
||||
# With coverage
|
||||
go test ./internal/email/... -v -cover
|
||||
|
||||
# With race detector
|
||||
go test ./internal/email/... -race
|
||||
```
|
||||
|
||||
### Manual Testing
|
||||
|
||||
Create a simple test program:
|
||||
|
||||
```go
|
||||
package main
|
||||
|
||||
import (
|
||||
"log"
|
||||
"github.com/usememos/memos/internal/email"
|
||||
)
|
||||
|
||||
func main() {
|
||||
config := &email.Config{
|
||||
SMTPHost: "smtp.gmail.com",
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "your-email@gmail.com",
|
||||
SMTPPassword: "your-app-password",
|
||||
FromEmail: "your-email@gmail.com",
|
||||
FromName: "Test",
|
||||
UseTLS: true,
|
||||
}
|
||||
|
||||
message := &email.Message{
|
||||
To: []string{"recipient@example.com"},
|
||||
Subject: "Test Email",
|
||||
Body: "This is a test email from Memos email plugin.",
|
||||
}
|
||||
|
||||
if err := email.Send(config, message); err != nil {
|
||||
log.Fatalf("Failed to send email: %v", err)
|
||||
}
|
||||
|
||||
log.Println("Email sent successfully!")
|
||||
}
|
||||
```
|
||||
|
||||
## Security Best Practices
|
||||
|
||||
### 1. Use TLS/SSL Encryption
|
||||
|
||||
Always enable encryption in production:
|
||||
|
||||
```go
|
||||
// STARTTLS (port 587) - Recommended
|
||||
config.UseTLS = true
|
||||
|
||||
// SSL/TLS (port 465)
|
||||
config.UseSSL = true
|
||||
```
|
||||
|
||||
### 2. Secure Credential Storage
|
||||
|
||||
Never hardcode credentials. Use environment variables:
|
||||
|
||||
```go
|
||||
import "os"
|
||||
|
||||
config := &email.Config{
|
||||
SMTPHost: os.Getenv("SMTP_HOST"),
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: os.Getenv("SMTP_USERNAME"),
|
||||
SMTPPassword: os.Getenv("SMTP_PASSWORD"),
|
||||
FromEmail: os.Getenv("SMTP_FROM_EMAIL"),
|
||||
UseTLS: true,
|
||||
}
|
||||
```
|
||||
|
||||
### 3. Use App-Specific Passwords
|
||||
|
||||
For Gmail and similar services, use app passwords instead of your main account password.
|
||||
|
||||
### 4. Validate and Sanitize Input
|
||||
|
||||
Always validate email addresses and sanitize content:
|
||||
|
||||
```go
|
||||
// Validate before sending
|
||||
if err := message.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
```
|
||||
|
||||
### 5. Implement Rate Limiting
|
||||
|
||||
Prevent abuse by limiting email sending:
|
||||
|
||||
```go
|
||||
// Example using golang.org/x/time/rate
|
||||
limiter := rate.NewLimiter(rate.Every(time.Second), 10) // 10 emails per second
|
||||
|
||||
if !limiter.Allow() {
|
||||
return errors.New("rate limit exceeded")
|
||||
}
|
||||
```
|
||||
|
||||
### 6. Monitor and Log
|
||||
|
||||
Log email sending activity for security monitoring:
|
||||
|
||||
```go
|
||||
if err := email.Send(config, message); err != nil {
|
||||
slog.Error("Email send failed",
|
||||
slog.String("recipient", message.To[0]),
|
||||
slog.Any("error", err))
|
||||
}
|
||||
```
|
||||
|
||||
## Common Ports
|
||||
|
||||
| Port | Protocol | Security | Use Case |
|
||||
|------|----------|----------|----------|
|
||||
| **587** | SMTP + STARTTLS | Encrypted | **Recommended** for most providers |
|
||||
| **465** | SMTP over SSL/TLS | Encrypted | Alternative secure option |
|
||||
| **25** | SMTP | Unencrypted | Legacy, often blocked by ISPs |
|
||||
| **2525** | SMTP + STARTTLS | Encrypted | Alternative when 587 is blocked |
|
||||
|
||||
**Port 587 (STARTTLS)** is the recommended standard for modern SMTP:
|
||||
```go
|
||||
config := &email.Config{
|
||||
SMTPPort: 587,
|
||||
UseTLS: true,
|
||||
}
|
||||
```
|
||||
|
||||
**Port 465 (SSL/TLS)** is the alternative:
|
||||
```go
|
||||
config := &email.Config{
|
||||
SMTPPort: 465,
|
||||
UseSSL: true,
|
||||
}
|
||||
```
|
||||
|
||||
## Error Handling
|
||||
|
||||
The package provides detailed, contextual errors:
|
||||
|
||||
```go
|
||||
err := email.Send(config, message)
|
||||
if err != nil {
|
||||
// Error messages include context:
|
||||
switch {
|
||||
case strings.Contains(err.Error(), "invalid email configuration"):
|
||||
// Configuration error (missing host, invalid port, etc.)
|
||||
log.Printf("Configuration error: %v", err)
|
||||
|
||||
case strings.Contains(err.Error(), "invalid email message"):
|
||||
// Message validation error (missing recipients, subject, body)
|
||||
log.Printf("Message error: %v", err)
|
||||
|
||||
case strings.Contains(err.Error(), "authentication failed"):
|
||||
// SMTP authentication failed (wrong credentials)
|
||||
log.Printf("Auth error: %v", err)
|
||||
|
||||
case strings.Contains(err.Error(), "failed to connect"):
|
||||
// Network/connection error
|
||||
log.Printf("Connection error: %v", err)
|
||||
|
||||
case strings.Contains(err.Error(), "recipient rejected"):
|
||||
// SMTP server rejected recipient
|
||||
log.Printf("Recipient error: %v", err)
|
||||
|
||||
default:
|
||||
log.Printf("Unknown error: %v", err)
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Common Error Messages
|
||||
|
||||
```
|
||||
❌ "invalid email configuration: SMTP host is required"
|
||||
→ Fix: Set config.SMTPHost
|
||||
|
||||
❌ "invalid email configuration: SMTP port must be between 1 and 65535"
|
||||
→ Fix: Set valid config.SMTPPort (usually 587 or 465)
|
||||
|
||||
❌ "invalid email configuration: from email is required"
|
||||
→ Fix: Set config.FromEmail
|
||||
|
||||
❌ "invalid email message: at least one recipient is required"
|
||||
→ Fix: Set message.To with at least one email address
|
||||
|
||||
❌ "invalid email message: subject is required"
|
||||
→ Fix: Set message.Subject
|
||||
|
||||
❌ "invalid email message: body is required"
|
||||
→ Fix: Set message.Body
|
||||
|
||||
❌ "SMTP authentication failed"
|
||||
→ Fix: Check credentials (username/password)
|
||||
|
||||
❌ "failed to connect to SMTP server"
|
||||
→ Fix: Verify host/port, check firewall, ensure TLS/SSL settings match server
|
||||
```
|
||||
|
||||
### Async Error Handling
|
||||
|
||||
For async sending, errors are logged automatically:
|
||||
|
||||
```go
|
||||
email.SendAsync(config, message)
|
||||
// Errors logged as:
|
||||
// [WARN] Failed to send email asynchronously recipients=user@example.com error=...
|
||||
```
|
||||
|
||||
## Dependencies
|
||||
|
||||
### Required
|
||||
|
||||
- **Go 1.25+**
|
||||
- Standard library: `net/smtp`, `crypto/tls`
|
||||
- `github.com/pkg/errors` - Error wrapping with context
|
||||
|
||||
### No External SMTP Libraries
|
||||
|
||||
This plugin uses Go's standard `net/smtp` library for maximum compatibility and minimal dependencies.
|
||||
|
||||
## API Reference
|
||||
|
||||
### Types
|
||||
|
||||
#### `Config`
|
||||
```go
|
||||
type Config struct {
|
||||
SMTPHost string // SMTP server hostname
|
||||
SMTPPort int // SMTP server port
|
||||
SMTPUsername string // SMTP auth username
|
||||
SMTPPassword string // SMTP auth password
|
||||
FromEmail string // From email address
|
||||
FromName string // From display name (optional)
|
||||
UseTLS bool // Enable STARTTLS (port 587)
|
||||
UseSSL bool // Enable SSL/TLS (port 465)
|
||||
}
|
||||
```
|
||||
|
||||
#### `Message`
|
||||
```go
|
||||
type Message struct {
|
||||
To []string // Recipients
|
||||
Cc []string // CC recipients (optional)
|
||||
Bcc []string // BCC recipients (optional)
|
||||
Subject string // Email subject
|
||||
Body string // Email body (plain text or HTML)
|
||||
IsHTML bool // true for HTML, false for plain text
|
||||
ReplyTo string // Reply-To address (optional)
|
||||
}
|
||||
```
|
||||
|
||||
### Functions
|
||||
|
||||
#### `Send(config *Config, message *Message) error`
|
||||
Sends an email synchronously. Blocks until email is sent or error occurs.
|
||||
|
||||
#### `SendAsync(config *Config, message *Message)`
|
||||
Sends an email asynchronously in a goroutine. Returns immediately. Errors are logged.
|
||||
|
||||
#### `NewClient(config *Config) *Client`
|
||||
Creates a new SMTP client for advanced usage.
|
||||
|
||||
#### `Client.Send(message *Message) error`
|
||||
Sends email using the client's configuration.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
internal/email/
|
||||
├── config.go # SMTP configuration types
|
||||
├── message.go # Email message types and formatting
|
||||
├── client.go # SMTP client implementation
|
||||
├── email.go # High-level Send/SendAsync API
|
||||
├── doc.go # Package documentation
|
||||
└── *_test.go # Unit tests
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
Part of the Memos project. See main repository for license details.
|
||||
|
||||
## Contributing
|
||||
|
||||
This package follows the Memos contribution guidelines. Please ensure:
|
||||
|
||||
1. All code is tested (TDD approach)
|
||||
2. Tests pass: `go test ./internal/email/... -v`
|
||||
3. Code is formatted: `go fmt ./internal/email/...`
|
||||
4. No linting errors: `golangci-lint run ./internal/email/...`
|
||||
|
||||
## Support
|
||||
|
||||
For issues and questions:
|
||||
|
||||
- Memos GitHub Issues: https://github.com/usememos/memos/issues
|
||||
- Memos Documentation: https://usememos.com/docs
|
||||
|
||||
## Roadmap
|
||||
|
||||
Future enhancements may include:
|
||||
|
||||
- Email template system
|
||||
- Attachment support
|
||||
- Inline image embedding
|
||||
- Email queuing system
|
||||
- Delivery status tracking
|
||||
- Bounce handling
|
||||
@@ -0,0 +1,173 @@
|
||||
package email
|
||||
|
||||
import (
|
||||
"crypto/tls"
|
||||
"net"
|
||||
"net/smtp"
|
||||
"time"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
const smtpOperationTimeout = 15 * time.Second
|
||||
|
||||
// Client represents an SMTP email client.
|
||||
type Client struct {
|
||||
config *Config
|
||||
}
|
||||
|
||||
// NewClient creates a new email client with the given configuration.
|
||||
func NewClient(config *Config) *Client {
|
||||
return &Client{
|
||||
config: config,
|
||||
}
|
||||
}
|
||||
|
||||
// validateConfig validates the client configuration.
|
||||
func (c *Client) validateConfig() error {
|
||||
if c.config == nil {
|
||||
return errors.New("email configuration is required")
|
||||
}
|
||||
return c.config.Validate()
|
||||
}
|
||||
|
||||
// createAuth creates an SMTP auth mechanism if credentials are provided.
|
||||
func (c *Client) createAuth() smtp.Auth {
|
||||
if c.config.SMTPUsername == "" && c.config.SMTPPassword == "" {
|
||||
return nil
|
||||
}
|
||||
return smtp.PlainAuth("", c.config.SMTPUsername, c.config.SMTPPassword, c.config.SMTPHost)
|
||||
}
|
||||
|
||||
// createTLSConfig creates a TLS configuration for secure connections.
|
||||
func (c *Client) createTLSConfig() *tls.Config {
|
||||
return &tls.Config{
|
||||
ServerName: c.config.SMTPHost,
|
||||
MinVersion: tls.VersionTLS12,
|
||||
}
|
||||
}
|
||||
|
||||
// Send sends an email message via SMTP.
|
||||
func (c *Client) Send(message *Message) error {
|
||||
// Validate configuration
|
||||
if err := c.validateConfig(); err != nil {
|
||||
return errors.Wrap(err, "invalid email configuration")
|
||||
}
|
||||
|
||||
// Validate message
|
||||
if message == nil {
|
||||
return errors.New("message is required")
|
||||
}
|
||||
if err := message.Validate(); err != nil {
|
||||
return errors.Wrap(err, "invalid email message")
|
||||
}
|
||||
|
||||
// Format the message
|
||||
body := message.Format(c.config.FromEmail, c.config.FromName)
|
||||
|
||||
// Get all recipients
|
||||
recipients := message.GetAllRecipients()
|
||||
|
||||
// Create auth
|
||||
auth := c.createAuth()
|
||||
|
||||
// Send based on encryption type
|
||||
if c.config.UseSSL {
|
||||
return c.sendWithSSL(auth, recipients, body)
|
||||
}
|
||||
return c.sendWithTLS(auth, recipients, body)
|
||||
}
|
||||
|
||||
// sendWithTLS sends email using STARTTLS (port 587).
|
||||
func (c *Client) sendWithTLS(auth smtp.Auth, recipients []string, body string) error {
|
||||
serverAddr := c.config.GetServerAddress()
|
||||
|
||||
dialer := &net.Dialer{Timeout: smtpOperationTimeout}
|
||||
conn, err := dialer.Dial("tcp", serverAddr)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "failed to connect to SMTP server: %s", serverAddr)
|
||||
}
|
||||
defer conn.Close()
|
||||
if err := conn.SetDeadline(time.Now().Add(smtpOperationTimeout)); err != nil {
|
||||
return errors.Wrap(err, "failed to set SMTP connection deadline")
|
||||
}
|
||||
|
||||
client, err := smtp.NewClient(conn, c.config.SMTPHost)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "failed to create SMTP client")
|
||||
}
|
||||
defer client.Quit()
|
||||
|
||||
if c.config.UseTLS {
|
||||
if ok, _ := client.Extension("STARTTLS"); !ok {
|
||||
return errors.New("SMTP server does not support STARTTLS")
|
||||
}
|
||||
if err := client.StartTLS(c.createTLSConfig()); err != nil {
|
||||
return errors.Wrap(err, "failed to start SMTP STARTTLS")
|
||||
}
|
||||
}
|
||||
|
||||
return c.sendWithClient(client, auth, recipients, body)
|
||||
}
|
||||
|
||||
// sendWithSSL sends email using SSL/TLS (port 465).
|
||||
func (c *Client) sendWithSSL(auth smtp.Auth, recipients []string, body string) error {
|
||||
serverAddr := c.config.GetServerAddress()
|
||||
|
||||
// Create TLS connection
|
||||
tlsConfig := c.createTLSConfig()
|
||||
dialer := &net.Dialer{Timeout: smtpOperationTimeout}
|
||||
conn, err := tls.DialWithDialer(dialer, "tcp", serverAddr, tlsConfig)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "failed to connect to SMTP server with SSL: %s", serverAddr)
|
||||
}
|
||||
defer conn.Close()
|
||||
if err := conn.SetDeadline(time.Now().Add(smtpOperationTimeout)); err != nil {
|
||||
return errors.Wrap(err, "failed to set SMTP connection deadline")
|
||||
}
|
||||
|
||||
// Create SMTP client
|
||||
client, err := smtp.NewClient(conn, c.config.SMTPHost)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "failed to create SMTP client")
|
||||
}
|
||||
defer client.Quit()
|
||||
|
||||
return c.sendWithClient(client, auth, recipients, body)
|
||||
}
|
||||
|
||||
func (c *Client) sendWithClient(client *smtp.Client, auth smtp.Auth, recipients []string, body string) error {
|
||||
if auth != nil {
|
||||
if err := client.Auth(auth); err != nil {
|
||||
return errors.Wrap(err, "SMTP authentication failed")
|
||||
}
|
||||
}
|
||||
|
||||
// Set sender
|
||||
if err := client.Mail(c.config.FromEmail); err != nil {
|
||||
return errors.Wrap(err, "failed to set sender")
|
||||
}
|
||||
|
||||
// Set recipients
|
||||
for _, recipient := range recipients {
|
||||
if err := client.Rcpt(recipient); err != nil {
|
||||
return errors.Wrapf(err, "failed to set recipient: %s", recipient)
|
||||
}
|
||||
}
|
||||
|
||||
// Send message body
|
||||
writer, err := client.Data()
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "failed to send DATA command")
|
||||
}
|
||||
|
||||
if _, err := writer.Write([]byte(body)); err != nil {
|
||||
return errors.Wrap(err, "failed to write message body")
|
||||
}
|
||||
|
||||
if err := writer.Close(); err != nil {
|
||||
return errors.Wrap(err, "failed to close message writer")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
package email
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestNewClient(t *testing.T) {
|
||||
config := &Config{
|
||||
SMTPHost: "smtp.example.com",
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "user@example.com",
|
||||
SMTPPassword: "password",
|
||||
FromEmail: "noreply@example.com",
|
||||
FromName: "Test App",
|
||||
UseTLS: true,
|
||||
}
|
||||
|
||||
client := NewClient(config)
|
||||
|
||||
assert.NotNil(t, client)
|
||||
assert.Equal(t, config, client.config)
|
||||
}
|
||||
|
||||
func TestClientValidateConfig(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
config *Config
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid config",
|
||||
config: &Config{
|
||||
SMTPHost: "smtp.example.com",
|
||||
SMTPPort: 587,
|
||||
FromEmail: "test@example.com",
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "nil config",
|
||||
config: nil,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "invalid config",
|
||||
config: &Config{
|
||||
SMTPHost: "",
|
||||
SMTPPort: 587,
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
client := NewClient(tt.config)
|
||||
err := client.validateConfig()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientSendValidation(t *testing.T) {
|
||||
config := &Config{
|
||||
SMTPHost: "smtp.example.com",
|
||||
SMTPPort: 587,
|
||||
FromEmail: "test@example.com",
|
||||
}
|
||||
client := NewClient(config)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
message *Message
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid message",
|
||||
message: &Message{
|
||||
To: []string{"recipient@example.com"},
|
||||
Subject: "Test",
|
||||
Body: "Test body",
|
||||
},
|
||||
wantErr: false, // Will fail on actual send, but passes validation
|
||||
},
|
||||
{
|
||||
name: "nil message",
|
||||
message: nil,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "invalid message",
|
||||
message: &Message{
|
||||
To: []string{},
|
||||
Subject: "Test",
|
||||
Body: "Test",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := client.Send(tt.message)
|
||||
// We expect validation errors for invalid messages
|
||||
// For valid messages, we'll get connection errors (which is expected in tests)
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
// Should fail validation before attempting connection
|
||||
assert.NotContains(t, err.Error(), "dial")
|
||||
}
|
||||
// Note: We don't assert NoError for valid messages because
|
||||
// we don't have a real SMTP server in tests
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package email
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
// Config represents the SMTP configuration for email sending.
|
||||
// These settings should be provided by the self-hosted instance administrator.
|
||||
type Config struct {
|
||||
// SMTPHost is the SMTP server hostname (e.g., "smtp.gmail.com")
|
||||
SMTPHost string
|
||||
// SMTPPort is the SMTP server port (common: 587 for TLS, 465 for SSL, 25 for unencrypted)
|
||||
SMTPPort int
|
||||
// SMTPUsername is the SMTP authentication username (usually the email address)
|
||||
SMTPUsername string
|
||||
// SMTPPassword is the SMTP authentication password or app-specific password
|
||||
SMTPPassword string
|
||||
// FromEmail is the email address that will appear in the "From" field
|
||||
FromEmail string
|
||||
// FromName is the display name that will appear in the "From" field
|
||||
FromName string
|
||||
// UseTLS enables STARTTLS encryption (recommended for port 587)
|
||||
UseTLS bool
|
||||
// UseSSL enables SSL/TLS encryption (for port 465)
|
||||
UseSSL bool
|
||||
}
|
||||
|
||||
// Validate checks if the configuration is valid.
|
||||
func (c *Config) Validate() error {
|
||||
if c.SMTPHost == "" {
|
||||
return errors.New("SMTP host is required")
|
||||
}
|
||||
if c.SMTPPort <= 0 || c.SMTPPort > 65535 {
|
||||
return errors.New("SMTP port must be between 1 and 65535")
|
||||
}
|
||||
if c.FromEmail == "" {
|
||||
return errors.New("from email is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetServerAddress returns the SMTP server address in the format "host:port".
|
||||
func (c *Config) GetServerAddress() string {
|
||||
return fmt.Sprintf("%s:%d", c.SMTPHost, c.SMTPPort)
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package email
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestConfigValidation(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
config *Config
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid config",
|
||||
config: &Config{
|
||||
SMTPHost: "smtp.gmail.com",
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "user@example.com",
|
||||
SMTPPassword: "password",
|
||||
FromEmail: "noreply@example.com",
|
||||
FromName: "Memos",
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "missing host",
|
||||
config: &Config{
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "user@example.com",
|
||||
SMTPPassword: "password",
|
||||
FromEmail: "noreply@example.com",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "invalid port",
|
||||
config: &Config{
|
||||
SMTPHost: "smtp.gmail.com",
|
||||
SMTPPort: 0,
|
||||
SMTPUsername: "user@example.com",
|
||||
SMTPPassword: "password",
|
||||
FromEmail: "noreply@example.com",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "missing from email",
|
||||
config: &Config{
|
||||
SMTPHost: "smtp.gmail.com",
|
||||
SMTPPort: 587,
|
||||
SMTPUsername: "user@example.com",
|
||||
SMTPPassword: "password",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := tt.config.Validate()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfigGetServerAddress(t *testing.T) {
|
||||
config := &Config{
|
||||
SMTPHost: "smtp.gmail.com",
|
||||
SMTPPort: 587,
|
||||
}
|
||||
|
||||
expected := "smtp.gmail.com:587"
|
||||
assert.Equal(t, expected, config.GetServerAddress())
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
// Package email provides SMTP email sending functionality for self-hosted Memos instances.
|
||||
//
|
||||
// This package is designed for self-hosted environments where instance administrators
|
||||
// configure their own SMTP servers. It follows industry-standard patterns used by
|
||||
// platforms like GitHub, GitLab, and Discourse.
|
||||
//
|
||||
// # Configuration
|
||||
//
|
||||
// The package requires SMTP server configuration provided by the instance administrator:
|
||||
//
|
||||
// config := &email.Config{
|
||||
// SMTPHost: "smtp.gmail.com",
|
||||
// SMTPPort: 587,
|
||||
// SMTPUsername: "your-email@gmail.com",
|
||||
// SMTPPassword: "your-app-password",
|
||||
// FromEmail: "noreply@yourdomain.com",
|
||||
// FromName: "Memos Notifications",
|
||||
// UseTLS: true,
|
||||
// }
|
||||
//
|
||||
// # Common SMTP Settings
|
||||
//
|
||||
// Gmail (requires App Password):
|
||||
// - Host: smtp.gmail.com
|
||||
// - Port: 587 (TLS) or 465 (SSL)
|
||||
// - Username: your-email@gmail.com
|
||||
// - UseTLS: true (for port 587) or UseSSL: true (for port 465)
|
||||
//
|
||||
// SendGrid:
|
||||
// - Host: smtp.sendgrid.net
|
||||
// - Port: 587
|
||||
// - Username: apikey
|
||||
// - Password: your-sendgrid-api-key
|
||||
// - UseTLS: true
|
||||
//
|
||||
// AWS SES:
|
||||
// - Host: email-smtp.[region].amazonaws.com
|
||||
// - Port: 587
|
||||
// - Username: your-smtp-username
|
||||
// - Password: your-smtp-password
|
||||
// - UseTLS: true
|
||||
//
|
||||
// Mailgun:
|
||||
// - Host: smtp.mailgun.org
|
||||
// - Port: 587
|
||||
// - Username: your-mailgun-smtp-username
|
||||
// - Password: your-mailgun-smtp-password
|
||||
// - UseTLS: true
|
||||
//
|
||||
// # Sending Email
|
||||
//
|
||||
// Synchronous (waits for completion):
|
||||
//
|
||||
// message := &email.Message{
|
||||
// To: []string{"user@example.com"},
|
||||
// Subject: "Welcome to Memos",
|
||||
// Body: "Thank you for joining!",
|
||||
// IsHTML: false,
|
||||
// }
|
||||
//
|
||||
// err := email.Send(config, message)
|
||||
// if err != nil {
|
||||
// // Handle error
|
||||
// }
|
||||
//
|
||||
// Asynchronous (returns immediately):
|
||||
//
|
||||
// email.SendAsync(config, message)
|
||||
// // Errors are logged but not returned
|
||||
//
|
||||
// # HTML Email
|
||||
//
|
||||
// message := &email.Message{
|
||||
// To: []string{"user@example.com"},
|
||||
// Subject: "Welcome!",
|
||||
// Body: "<html><body><h1>Welcome to Memos!</h1></body></html>",
|
||||
// IsHTML: true,
|
||||
// }
|
||||
//
|
||||
// # Security Considerations
|
||||
//
|
||||
// - Always use TLS (port 587) or SSL (port 465) for production
|
||||
// - Store SMTP credentials securely (environment variables or secrets management)
|
||||
// - Use app-specific passwords for services like Gmail
|
||||
// - Validate and sanitize email content to prevent injection attacks
|
||||
// - Rate limit email sending to prevent abuse
|
||||
//
|
||||
// # Error Handling
|
||||
//
|
||||
// The package returns descriptive errors for common issues:
|
||||
// - Configuration validation errors (missing host, invalid port, etc.)
|
||||
// - Message validation errors (missing recipients, subject, or body)
|
||||
// - Connection errors (cannot reach SMTP server)
|
||||
// - Authentication errors (invalid credentials)
|
||||
// - SMTP protocol errors (recipient rejected, etc.)
|
||||
//
|
||||
// All errors are wrapped with context using github.com/pkg/errors for better debugging.
|
||||
package email
|
||||
@@ -0,0 +1,61 @@
|
||||
package email
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
type asyncEmailRequest struct {
|
||||
config *Config
|
||||
message *Message
|
||||
}
|
||||
|
||||
var asyncEmailQueue = make(chan asyncEmailRequest, 128)
|
||||
|
||||
func init() {
|
||||
for range 2 {
|
||||
go func() {
|
||||
for request := range asyncEmailQueue {
|
||||
if err := Send(request.config, request.message); err != nil {
|
||||
recipients := ""
|
||||
if request.message != nil && len(request.message.To) > 0 {
|
||||
recipients = request.message.To[0]
|
||||
if len(request.message.To) > 1 {
|
||||
recipients += " and others"
|
||||
}
|
||||
}
|
||||
|
||||
slog.Warn("Failed to send email asynchronously",
|
||||
slog.String("recipients", recipients),
|
||||
slog.Any("error", err))
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
// Send sends an email synchronously.
|
||||
// Returns an error if the email fails to send.
|
||||
func Send(config *Config, message *Message) error {
|
||||
if config == nil {
|
||||
return errors.New("email configuration is required")
|
||||
}
|
||||
if message == nil {
|
||||
return errors.New("email message is required")
|
||||
}
|
||||
|
||||
client := NewClient(config)
|
||||
return client.Send(message)
|
||||
}
|
||||
|
||||
// SendAsync sends an email asynchronously.
|
||||
// It enqueues the message for bounded asynchronous sending and does not wait for the response.
|
||||
// Any errors are logged but not returned.
|
||||
func SendAsync(config *Config, message *Message) {
|
||||
select {
|
||||
case asyncEmailQueue <- asyncEmailRequest{config: config, message: message}:
|
||||
default:
|
||||
slog.Warn("Dropped email because the async queue is full")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
package email
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
func TestSend(t *testing.T) {
|
||||
config := &Config{
|
||||
SMTPHost: "smtp.example.com",
|
||||
SMTPPort: 587,
|
||||
FromEmail: "test@example.com",
|
||||
}
|
||||
|
||||
message := &Message{
|
||||
To: []string{"recipient@example.com"},
|
||||
Subject: "Test",
|
||||
Body: "Test body",
|
||||
}
|
||||
|
||||
// This will fail to connect (no real server), but should validate inputs
|
||||
err := Send(config, message)
|
||||
// We expect an error because there's no real SMTP server
|
||||
// But it should be a connection error, not a validation error
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "dial")
|
||||
}
|
||||
|
||||
func TestSendValidation(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
config *Config
|
||||
message *Message
|
||||
wantErr bool
|
||||
errMsg string
|
||||
}{
|
||||
{
|
||||
name: "nil config",
|
||||
config: nil,
|
||||
message: &Message{To: []string{"test@example.com"}, Subject: "Test", Body: "Test"},
|
||||
wantErr: true,
|
||||
errMsg: "configuration is required",
|
||||
},
|
||||
{
|
||||
name: "nil message",
|
||||
config: &Config{SMTPHost: "smtp.example.com", SMTPPort: 587, FromEmail: "from@example.com"},
|
||||
message: nil,
|
||||
wantErr: true,
|
||||
errMsg: "message is required",
|
||||
},
|
||||
{
|
||||
name: "invalid config",
|
||||
config: &Config{
|
||||
SMTPHost: "",
|
||||
SMTPPort: 587,
|
||||
},
|
||||
message: &Message{To: []string{"test@example.com"}, Subject: "Test", Body: "Test"},
|
||||
wantErr: true,
|
||||
errMsg: "invalid email configuration",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := Send(tt.config, tt.message)
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), tt.errMsg)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendAsync(t *testing.T) {
|
||||
config := &Config{
|
||||
SMTPHost: "smtp.example.com",
|
||||
SMTPPort: 587,
|
||||
FromEmail: "test@example.com",
|
||||
}
|
||||
|
||||
message := &Message{
|
||||
To: []string{"recipient@example.com"},
|
||||
Subject: "Test Async",
|
||||
Body: "Test async body",
|
||||
}
|
||||
|
||||
// SendAsync should not block
|
||||
start := time.Now()
|
||||
SendAsync(config, message)
|
||||
duration := time.Since(start)
|
||||
|
||||
// Should return almost immediately (< 100ms)
|
||||
assert.Less(t, duration, 100*time.Millisecond)
|
||||
|
||||
// Give goroutine time to start
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
|
||||
func TestSendAsyncConcurrent(t *testing.T) {
|
||||
config := &Config{
|
||||
SMTPHost: "smtp.example.com",
|
||||
SMTPPort: 587,
|
||||
FromEmail: "test@example.com",
|
||||
}
|
||||
|
||||
g := errgroup.Group{}
|
||||
count := 5
|
||||
|
||||
for i := 0; i < count; i++ {
|
||||
g.Go(func() error {
|
||||
message := &Message{
|
||||
To: []string{"recipient@example.com"},
|
||||
Subject: "Concurrent Test",
|
||||
Body: "Test body",
|
||||
}
|
||||
SendAsync(config, message)
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
if err := g.Wait(); err != nil {
|
||||
t.Fatalf("SendAsync calls failed: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package email
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Message represents an email message to be sent.
|
||||
type Message struct {
|
||||
To []string // Required: recipient email addresses
|
||||
Cc []string // Optional: carbon copy recipients
|
||||
Bcc []string // Optional: blind carbon copy recipients
|
||||
Subject string // Required: email subject
|
||||
Body string // Required: email body content
|
||||
IsHTML bool // Whether the body is HTML (default: false for plain text)
|
||||
ReplyTo string // Optional: reply-to address
|
||||
}
|
||||
|
||||
// Validate checks that the message has all required fields.
|
||||
func (m *Message) Validate() error {
|
||||
if len(m.To) == 0 {
|
||||
return errors.New("at least one recipient is required")
|
||||
}
|
||||
if m.Subject == "" {
|
||||
return errors.New("subject is required")
|
||||
}
|
||||
if m.Body == "" {
|
||||
return errors.New("body is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Format creates an RFC 5322 formatted email message.
|
||||
func (m *Message) Format(fromEmail, fromName string) string {
|
||||
var sb strings.Builder
|
||||
fromEmail = sanitizeEmailHeaderValue(fromEmail)
|
||||
fromName = sanitizeEmailHeaderValue(fromName)
|
||||
to := sanitizeEmailHeaderValues(m.To)
|
||||
cc := sanitizeEmailHeaderValues(m.Cc)
|
||||
replyTo := sanitizeEmailHeaderValue(m.ReplyTo)
|
||||
subject := sanitizeEmailHeaderValue(m.Subject)
|
||||
|
||||
// From header
|
||||
if fromName != "" {
|
||||
fmt.Fprintf(&sb, "From: %s <%s>\r\n", fromName, fromEmail)
|
||||
} else {
|
||||
fmt.Fprintf(&sb, "From: %s\r\n", fromEmail)
|
||||
}
|
||||
|
||||
// To header
|
||||
fmt.Fprintf(&sb, "To: %s\r\n", strings.Join(to, ", "))
|
||||
|
||||
// Cc header (optional)
|
||||
if len(cc) > 0 {
|
||||
fmt.Fprintf(&sb, "Cc: %s\r\n", strings.Join(cc, ", "))
|
||||
}
|
||||
|
||||
// Reply-To header (optional)
|
||||
if replyTo != "" {
|
||||
fmt.Fprintf(&sb, "Reply-To: %s\r\n", replyTo)
|
||||
}
|
||||
|
||||
// Subject header
|
||||
fmt.Fprintf(&sb, "Subject: %s\r\n", subject)
|
||||
|
||||
// Date header (RFC 5322 format)
|
||||
fmt.Fprintf(&sb, "Date: %s\r\n", time.Now().Format(time.RFC1123Z))
|
||||
|
||||
// MIME headers
|
||||
sb.WriteString("MIME-Version: 1.0\r\n")
|
||||
|
||||
// Content-Type header
|
||||
if m.IsHTML {
|
||||
sb.WriteString("Content-Type: text/html; charset=utf-8\r\n")
|
||||
} else {
|
||||
sb.WriteString("Content-Type: text/plain; charset=utf-8\r\n")
|
||||
}
|
||||
|
||||
// Empty line separating headers from body
|
||||
sb.WriteString("\r\n")
|
||||
|
||||
// Body
|
||||
sb.WriteString(m.Body)
|
||||
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func sanitizeEmailHeaderValue(value string) string {
|
||||
value = strings.NewReplacer("\r", " ", "\n", " ").Replace(value)
|
||||
return strings.Join(strings.Fields(value), " ")
|
||||
}
|
||||
|
||||
func sanitizeEmailHeaderValues(values []string) []string {
|
||||
sanitized := make([]string, 0, len(values))
|
||||
for _, value := range values {
|
||||
sanitized = append(sanitized, sanitizeEmailHeaderValue(value))
|
||||
}
|
||||
return sanitized
|
||||
}
|
||||
|
||||
// GetAllRecipients returns all recipients (To, Cc, Bcc) as a single slice.
|
||||
func (m *Message) GetAllRecipients() []string {
|
||||
var recipients []string
|
||||
recipients = append(recipients, m.To...)
|
||||
recipients = append(recipients, m.Cc...)
|
||||
recipients = append(recipients, m.Bcc...)
|
||||
return recipients
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
package email
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMessageValidation(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
msg Message
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid message",
|
||||
msg: Message{
|
||||
To: []string{"user@example.com"},
|
||||
Subject: "Test Subject",
|
||||
Body: "Test Body",
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "no recipients",
|
||||
msg: Message{
|
||||
To: []string{},
|
||||
Subject: "Test Subject",
|
||||
Body: "Test Body",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "no subject",
|
||||
msg: Message{
|
||||
To: []string{"user@example.com"},
|
||||
Subject: "",
|
||||
Body: "Test Body",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "no body",
|
||||
msg: Message{
|
||||
To: []string{"user@example.com"},
|
||||
Subject: "Test Subject",
|
||||
Body: "",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "multiple recipients",
|
||||
msg: Message{
|
||||
To: []string{"user1@example.com", "user2@example.com"},
|
||||
Cc: []string{"cc@example.com"},
|
||||
Subject: "Test Subject",
|
||||
Body: "Test Body",
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := tt.msg.Validate()
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("Validate() error = %v, wantErr %v", err, tt.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageFormatPlainText(t *testing.T) {
|
||||
msg := Message{
|
||||
To: []string{"user@example.com"},
|
||||
Subject: "Test Subject",
|
||||
Body: "Test Body",
|
||||
IsHTML: false,
|
||||
}
|
||||
|
||||
formatted := msg.Format("sender@example.com", "Sender Name")
|
||||
|
||||
// Check required headers
|
||||
if !strings.Contains(formatted, "From: Sender Name <sender@example.com>") {
|
||||
t.Error("Missing or incorrect From header")
|
||||
}
|
||||
if !strings.Contains(formatted, "To: user@example.com") {
|
||||
t.Error("Missing or incorrect To header")
|
||||
}
|
||||
if !strings.Contains(formatted, "Subject: Test Subject") {
|
||||
t.Error("Missing or incorrect Subject header")
|
||||
}
|
||||
if !strings.Contains(formatted, "Content-Type: text/plain; charset=utf-8") {
|
||||
t.Error("Missing or incorrect Content-Type header for plain text")
|
||||
}
|
||||
if !strings.Contains(formatted, "Test Body") {
|
||||
t.Error("Missing message body")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageFormatHTML(t *testing.T) {
|
||||
msg := Message{
|
||||
To: []string{"user@example.com"},
|
||||
Subject: "Test Subject",
|
||||
Body: "<html><body>Test Body</body></html>",
|
||||
IsHTML: true,
|
||||
}
|
||||
|
||||
formatted := msg.Format("sender@example.com", "Sender Name")
|
||||
|
||||
// Check HTML content-type
|
||||
if !strings.Contains(formatted, "Content-Type: text/html; charset=utf-8") {
|
||||
t.Error("Missing or incorrect Content-Type header for HTML")
|
||||
}
|
||||
if !strings.Contains(formatted, "<html><body>Test Body</body></html>") {
|
||||
t.Error("Missing HTML body")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageFormatMultipleRecipients(t *testing.T) {
|
||||
msg := Message{
|
||||
To: []string{"user1@example.com", "user2@example.com"},
|
||||
Cc: []string{"cc1@example.com", "cc2@example.com"},
|
||||
Bcc: []string{"bcc@example.com"},
|
||||
Subject: "Test Subject",
|
||||
Body: "Test Body",
|
||||
ReplyTo: "reply@example.com",
|
||||
}
|
||||
|
||||
formatted := msg.Format("sender@example.com", "Sender Name")
|
||||
|
||||
// Check To header formatting
|
||||
if !strings.Contains(formatted, "To: user1@example.com, user2@example.com") {
|
||||
t.Error("Missing or incorrect To header with multiple recipients")
|
||||
}
|
||||
// Check Cc header formatting
|
||||
if !strings.Contains(formatted, "Cc: cc1@example.com, cc2@example.com") {
|
||||
t.Error("Missing or incorrect Cc header")
|
||||
}
|
||||
// Bcc should NOT appear in the formatted message
|
||||
if strings.Contains(formatted, "Bcc:") {
|
||||
t.Error("Bcc header should not appear in formatted message")
|
||||
}
|
||||
// Check Reply-To header
|
||||
if !strings.Contains(formatted, "Reply-To: reply@example.com") {
|
||||
t.Error("Missing or incorrect Reply-To header")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageFormatSanitizesHeaderValues(t *testing.T) {
|
||||
msg := Message{
|
||||
To: []string{"user@example.com\r\nX-Injected-To: bad"},
|
||||
Cc: []string{"cc@example.com\r\nX-Injected-Cc: bad"},
|
||||
Subject: "Test\r\nX-Injected-Subject: bad",
|
||||
Body: "Test Body",
|
||||
ReplyTo: "reply@example.com\r\nX-Injected-Reply-To: bad",
|
||||
}
|
||||
|
||||
formatted := msg.Format("sender@example.com\r\nX-Injected-From: bad", "Sender\r\nX-Injected-Name: bad")
|
||||
headers := strings.SplitN(formatted, "\r\n\r\n", 2)[0]
|
||||
|
||||
if strings.Contains(headers, "\r\nX-Injected") {
|
||||
t.Fatalf("header value injection was not sanitized:\n%s", headers)
|
||||
}
|
||||
if !strings.Contains(headers, "Subject: Test X-Injected-Subject: bad") {
|
||||
t.Error("subject header was not normalized")
|
||||
}
|
||||
if !strings.Contains(headers, "From: Sender X-Injected-Name: bad <sender@example.com X-Injected-From: bad>") {
|
||||
t.Error("from header was not normalized")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAllRecipients(t *testing.T) {
|
||||
msg := Message{
|
||||
To: []string{"user1@example.com", "user2@example.com"},
|
||||
Cc: []string{"cc@example.com"},
|
||||
Bcc: []string{"bcc@example.com"},
|
||||
}
|
||||
|
||||
recipients := msg.GetAllRecipients()
|
||||
|
||||
// Should have all 4 recipients
|
||||
if len(recipients) != 4 {
|
||||
t.Errorf("GetAllRecipients() returned %d recipients, want 4", len(recipients))
|
||||
}
|
||||
|
||||
// Check all recipients are present
|
||||
expectedRecipients := map[string]bool{
|
||||
"user1@example.com": true,
|
||||
"user2@example.com": true,
|
||||
"cc@example.com": true,
|
||||
"bcc@example.com": true,
|
||||
}
|
||||
|
||||
for _, recipient := range recipients {
|
||||
if !expectedRecipients[recipient] {
|
||||
t.Errorf("Unexpected recipient: %s", recipient)
|
||||
}
|
||||
delete(expectedRecipients, recipient)
|
||||
}
|
||||
|
||||
if len(expectedRecipients) > 0 {
|
||||
t.Error("Not all expected recipients were returned")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
# Maintaining the Memo Filter Engine
|
||||
|
||||
The engine is memo-specific; any future field or behavior changes must stay
|
||||
consistent with the memo schema and store implementations. Use this guide when
|
||||
extending or debugging the package.
|
||||
|
||||
## Adding a New Memo Field
|
||||
|
||||
1. **Update the schema**
|
||||
- Add the field entry in `schema.go`.
|
||||
- Define the backing column (`Column`), JSON path (if applicable), type, and
|
||||
allowed operators.
|
||||
- Include the CEL variable in `EnvOptions`.
|
||||
2. **Adjust parser or renderer (if needed)**
|
||||
- For non-scalar fields (JSON booleans, lists), add handling in
|
||||
`parser.go` or extend the renderer helpers.
|
||||
- Keep validation in the parser (e.g., reject unsupported operators).
|
||||
3. **Write a golden test**
|
||||
- Extend the dialect-specific memo filter tests under
|
||||
`store/db/{sqlite,mysql,postgres}/memo_filter_test.go` with a case that
|
||||
exercises the new field.
|
||||
4. **Run `go test ./...`** to ensure the SQL output matches expectations across
|
||||
all dialects.
|
||||
|
||||
## Supporting Dialect Nuances
|
||||
|
||||
- Centralize differences inside `render.go`. If a new dialect-specific behavior
|
||||
emerges (e.g., JSON operators), add the logic there rather than leaking it
|
||||
into store code.
|
||||
- Use the renderer helpers (`jsonExtractExpr`, `jsonArrayExpr`, etc.) rather than
|
||||
sprinkling ad-hoc SQL strings.
|
||||
- When placeholders change, adjust `addArg` so that argument numbering stays in
|
||||
sync with store queries.
|
||||
|
||||
## Debugging Tips
|
||||
|
||||
- **Parser errors** – Most originate in `buildCondition` or schema validation.
|
||||
Enable logging around `parser.go` when diagnosing unknown identifier/operator
|
||||
messages.
|
||||
- **Renderer output** – Temporary printf/log statements in `renderCondition` help
|
||||
identify which IR node produced unexpected SQL.
|
||||
- **Store integration** – Ensure drivers call `filter.DefaultEngine()` exactly once
|
||||
per process; the singleton caches the parsed CEL environment.
|
||||
|
||||
## Testing Checklist
|
||||
|
||||
- `go test ./store/...` ensures all dialect tests consume the engine correctly.
|
||||
- Add targeted unit tests whenever new IR nodes or renderer paths are introduced.
|
||||
- When changing boolean or JSON handling, verify all three dialect test suites
|
||||
(SQLite, MySQL, Postgres) to avoid regression.
|
||||
@@ -0,0 +1,105 @@
|
||||
# Memo Filter Engine
|
||||
|
||||
This package houses the memo-only filter engine that turns standard CEL syntax
|
||||
into SQL fragments for the subset of expressions supported by the memo schema.
|
||||
The engine follows a three phase pipeline inspired by systems
|
||||
such as Calcite or Prisma:
|
||||
|
||||
1. **Parsing** – CEL expressions are parsed with `cel-go` and validated against
|
||||
the memo-specific environment declared in `schema.go`. Only fields that
|
||||
exist in the schema can surface in the filter, and non-standard legacy
|
||||
coercions are rejected.
|
||||
2. **Normalization** – the raw CEL AST is converted into an intermediate
|
||||
representation (IR) defined in `ir.go`. The IR is a dialect-agnostic tree of
|
||||
conditions (logical operators, comparisons, list membership, etc.). This
|
||||
step enforces schema rules (e.g. operator compatibility, type checks).
|
||||
3. **Rendering** – the renderer in `render.go` walks the IR and produces a SQL
|
||||
fragment plus placeholder arguments tailored to a target dialect
|
||||
(`sqlite`, `mysql`, or `postgres`). Dialect differences such as JSON access,
|
||||
boolean semantics, placeholders, and `LIKE` vs `ILIKE` are encapsulated in
|
||||
renderer helpers.
|
||||
|
||||
The entry point is `filter.DefaultEngine()` from `engine.go`. It lazily constructs
|
||||
an `Engine` configured with the memo schema and exposes:
|
||||
|
||||
```go
|
||||
engine, _ := filter.DefaultEngine()
|
||||
stmt, _ := engine.CompileToStatement(ctx, `has_task_list && visibility == "PUBLIC"`, filter.RenderOptions{
|
||||
Dialect: filter.DialectPostgres,
|
||||
})
|
||||
// stmt.SQL -> "((memo.payload->'property'->>'hasTaskList')::boolean IS TRUE AND memo.visibility = $1)"
|
||||
// stmt.Args -> ["PUBLIC"]
|
||||
```
|
||||
|
||||
## Core Files
|
||||
|
||||
| File | Responsibility |
|
||||
| ------------- | ------------------------------------------------------------------------------- |
|
||||
| `schema.go` | Declares memo fields, their types, backing columns, CEL environment options |
|
||||
| `ir.go` | IR node definitions used across the pipeline |
|
||||
| `parser.go` | Converts CEL `Expr` into IR while applying schema validation |
|
||||
| `render.go` | Translates IR into SQL, handling dialect-specific behavior |
|
||||
| `engine.go` | Glue between the phases; exposes `Compile`, `CompileToStatement`, and `DefaultEngine` |
|
||||
| `helpers.go` | Convenience helpers for store integration (appending conditions) |
|
||||
|
||||
## SQL Generation Notes
|
||||
|
||||
- **Placeholders** — `?` is used for SQLite/MySQL, `$n` for Postgres. The renderer
|
||||
tracks offsets to compose queries with pre-existing arguments.
|
||||
- **JSON Fields** — Memo metadata lives in `memo.payload`. The renderer handles
|
||||
`JSON_EXTRACT`/`json_extract`/`->`/`->>` variations and boolean coercion.
|
||||
- **Time Fields** — `created_ts`, `updated_ts`, and attachment `create_time` are
|
||||
CEL `timestamp` values. Express instants with the `now` variable,
|
||||
`duration("…")` (e.g. `created_ts >= now - duration("24h")`), or
|
||||
`timestamp("2006-01-02T15:04:05Z")` / `timestamp(<epoch-seconds>)`. These fold
|
||||
to epoch seconds at compile time — `now` is frozen once per compile (injectable
|
||||
for tests via the engine clock) — so the backing columns stay unchanged.
|
||||
- **Tag Operations** — `tag in [...]` and `"tag" in tags` become JSON array
|
||||
predicates. SQLite uses `LIKE` patterns, MySQL uses `JSON_CONTAINS`, and
|
||||
Postgres uses `@>`.
|
||||
- **Boolean Flags** — Fields such as `has_task_list` render as `IS TRUE` equality
|
||||
checks, or comparisons against `CAST('true' AS JSON)` depending on the dialect.
|
||||
- **String Matching** — `content.contains(x)`, `content.startsWith(x)`, and
|
||||
`content.endsWith(x)` render as case-insensitive `LIKE`/`ILIKE` with LIKE
|
||||
metacharacters (`%`, `_`, `\`) escaped. Available on scalar string fields whose
|
||||
schema sets `SupportsContains` (memo `content`; attachment `filename`,
|
||||
`mime_type`).
|
||||
- **Regex** — `field.matches("pattern")` renders to `~` (Postgres) or `REGEXP`
|
||||
(MySQL/SQLite). SQLite uses a Go-backed `regexp` function registered in
|
||||
`store/db/sqlite/functions.go`. Patterns are validated at compile time against
|
||||
Go's RE2 via `cel.ValidateRegexLiterals()`. **Caveat:** regex *syntax* differs
|
||||
per engine (Go RE2 on SQLite, POSIX ERE on Postgres, ICU on MySQL 8.0+), so
|
||||
engine-specific patterns may not be portable.
|
||||
- **Tag `all()` / `exists_one()`** — `tags.all(t, <pred>)` matches only non-empty
|
||||
tag sets where every element satisfies the predicate; `tags.exists_one(t,
|
||||
<pred>)` matches when exactly one element does (`COUNT(...) = 1`). Both iterate
|
||||
per-element (`json_each` / `jsonb_array_elements_text` / `JSON_TABLE`).
|
||||
- **Timestamp Accessors** — `created_ts.getFullYear()`, `getMonth()`, `getDate()`,
|
||||
`getDayOfMonth()`, `getDayOfWeek()`, `getDayOfYear()`, `getHours()`,
|
||||
`getMinutes()`, `getSeconds()` render to date-part extraction (`strftime` /
|
||||
`EXTRACT` / `YEAR`/`MONTH`/…). Results are normalized to CEL's base (0-based
|
||||
month, 0-based day-of-week with 0 = Sunday). The same accessors on `now` fold
|
||||
to literal date parts of the frozen evaluation time (UTC), so saved filters
|
||||
like `created_ts.getMonth() == now.getMonth() && created_ts.getDate() ==
|
||||
now.getDate()` ("on this day") re-resolve on every compile. Extraction is UTC
|
||||
on SQLite/Postgres (epoch columns); on MySQL the `TIMESTAMP` column is read in
|
||||
the session time zone. A timezone argument is not supported.
|
||||
- **Set Operations** — `ext.Sets()`: `sets.contains(tags, [...])`,
|
||||
`sets.intersects(tags, [...])`, and `sets.equivalent(tags, [...])` desugar to
|
||||
exact-membership checks (AND / OR of `"v" in tags`); `equivalent` adds a
|
||||
`size(tags)` length check (relies on tags being a set).
|
||||
- **`size()`** — `size(tags)` renders to JSON array length; `size(content)` (and
|
||||
other string fields) render to `LENGTH` / `CHAR_LENGTH` (MySQL) for code-point
|
||||
counts.
|
||||
- **Arithmetic** — `+`, `-`, `*`, `/`, `%` constant-fold on literal/`now`/`duration`
|
||||
operands (division and modulo guard against a zero divisor).
|
||||
|
||||
## Typical Integration
|
||||
|
||||
1. Fetch the engine with `filter.DefaultEngine()`.
|
||||
2. Call `CompileToStatement` using the appropriate dialect enum.
|
||||
3. Append the emitted SQL fragment/args to the existing `WHERE` clause.
|
||||
4. Execute the resulting query through the store driver.
|
||||
|
||||
The `helpers.AppendConditions` helper encapsulates steps 2–3 when a driver needs
|
||||
to process an array of filters.
|
||||
@@ -0,0 +1,124 @@
|
||||
package filter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/cel-go/cel"
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
// Engine parses CEL filters into a dialect-agnostic condition tree.
|
||||
type Engine struct {
|
||||
schema Schema
|
||||
env *cel.Env
|
||||
// nowFunc resolves the value of the `now` variable. It is frozen once per
|
||||
// Compile so a single filter sees a single instant, and is overridable in
|
||||
// tests for deterministic folding.
|
||||
nowFunc func() time.Time
|
||||
}
|
||||
|
||||
// NewEngine builds a new Engine for the provided schema.
|
||||
func NewEngine(schema Schema) (*Engine, error) {
|
||||
env, err := cel.NewEnv(schema.EnvOptions...)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to create CEL environment")
|
||||
}
|
||||
return &Engine{
|
||||
schema: schema,
|
||||
env: env,
|
||||
nowFunc: time.Now,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Program stores a compiled filter condition.
|
||||
type Program struct {
|
||||
schema Schema
|
||||
condition Condition
|
||||
}
|
||||
|
||||
// ConditionTree exposes the underlying condition tree.
|
||||
func (p *Program) ConditionTree() Condition {
|
||||
return p.condition
|
||||
}
|
||||
|
||||
// Compile parses the filter string into an executable program.
|
||||
func (e *Engine) Compile(_ context.Context, filter string) (*Program, error) {
|
||||
if strings.TrimSpace(filter) == "" {
|
||||
return nil, errors.New("filter expression is empty")
|
||||
}
|
||||
|
||||
ast, issues := e.env.Compile(filter)
|
||||
if issues != nil && issues.Err() != nil {
|
||||
return nil, errors.Wrap(issues.Err(), "failed to compile filter")
|
||||
}
|
||||
parsed, err := cel.AstToParsedExpr(ast)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to convert AST")
|
||||
}
|
||||
|
||||
cond, err := buildCondition(parsed.GetExpr(), parseContext{schema: e.schema, now: e.nowFunc()})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &Program{
|
||||
schema: e.schema,
|
||||
condition: cond,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// CompileToStatement compiles and renders the filter in a single step.
|
||||
func (e *Engine) CompileToStatement(ctx context.Context, filter string, opts RenderOptions) (Statement, error) {
|
||||
program, err := e.Compile(ctx, filter)
|
||||
if err != nil {
|
||||
return Statement{}, err
|
||||
}
|
||||
return program.Render(opts)
|
||||
}
|
||||
|
||||
// RenderOptions configure SQL rendering.
|
||||
type RenderOptions struct {
|
||||
Dialect DialectName
|
||||
PlaceholderOffset int
|
||||
DisableNullChecks bool
|
||||
}
|
||||
|
||||
// Statement contains the rendered SQL fragment and its args.
|
||||
type Statement struct {
|
||||
SQL string
|
||||
Args []any
|
||||
}
|
||||
|
||||
// Render converts the program into a dialect-specific SQL fragment.
|
||||
func (p *Program) Render(opts RenderOptions) (Statement, error) {
|
||||
renderer := newRenderer(p.schema, opts)
|
||||
return renderer.Render(p.condition)
|
||||
}
|
||||
|
||||
var (
|
||||
defaultOnce sync.Once
|
||||
defaultInst *Engine
|
||||
defaultErr error
|
||||
defaultAttachmentOnce sync.Once
|
||||
defaultAttachmentInst *Engine
|
||||
defaultAttachmentErr error
|
||||
)
|
||||
|
||||
// DefaultEngine returns the process-wide memo filter engine.
|
||||
func DefaultEngine() (*Engine, error) {
|
||||
defaultOnce.Do(func() {
|
||||
defaultInst, defaultErr = NewEngine(NewSchema())
|
||||
})
|
||||
return defaultInst, defaultErr
|
||||
}
|
||||
|
||||
// DefaultAttachmentEngine returns the process-wide attachment filter engine.
|
||||
func DefaultAttachmentEngine() (*Engine, error) {
|
||||
defaultAttachmentOnce.Do(func() {
|
||||
defaultAttachmentInst, defaultAttachmentErr = NewEngine(NewAttachmentSchema())
|
||||
})
|
||||
return defaultAttachmentInst, defaultAttachmentErr
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
package filter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCompileAcceptsStandardTagEqualityPredicate(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = engine.Compile(context.Background(), `tags.exists(t, t == "1231")`)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestCompileRejectsLegacyNumericLogicalOperand(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = engine.Compile(context.Background(), `pinned && 1`)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "failed to compile filter")
|
||||
}
|
||||
|
||||
func TestCompileRejectsNonBooleanTopLevelConstant(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = engine.Compile(context.Background(), `1`)
|
||||
require.EqualError(t, err, "filter must evaluate to a boolean value")
|
||||
}
|
||||
|
||||
func TestCompileRejectsMalformedRegex(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = engine.Compile(context.Background(), `content.matches("(")`)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestCompileMatchesRendersRegexOperator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `content.matches("v[0-9]+")`, RenderOptions{Dialect: DialectPostgres})
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, stmt.SQL, "~")
|
||||
require.Equal(t, []any{"v[0-9]+"}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileRejectsStartsWithOnUnsupportedField(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = engine.Compile(context.Background(), `visibility.startsWith("P")`)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "does not support text matching")
|
||||
}
|
||||
|
||||
func TestCompileContainsEscapesLikeWildcards(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `content.contains("50%_off")`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
// The % and _ in the value must be escaped so they are matched literally,
|
||||
// and SQLite needs an explicit ESCAPE clause.
|
||||
require.Contains(t, stmt.SQL, `ESCAPE '\'`)
|
||||
require.Equal(t, []any{`%50\%\_off%`}, stmt.Args)
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Cross-dialect rendering tests (no DB required; complements the SQLite-only
|
||||
// behavioral tests in store/test by asserting MySQL/Postgres SQL generation).
|
||||
// =============================================================================
|
||||
|
||||
func TestRenderStartsWithPerDialect(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
cases := []struct {
|
||||
dialect DialectName
|
||||
fragments []string
|
||||
}{
|
||||
{DialectSQLite, []string{"memos_unicode_lower(", "`memo`.`content`", `ESCAPE '\'`}},
|
||||
{DialectPostgres, []string{"memo.content ILIKE $1"}},
|
||||
{DialectMySQL, []string{"`memo`.`content` LIKE ?"}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `content.startsWith("TODO")`, RenderOptions{Dialect: tc.dialect})
|
||||
require.NoError(t, err, tc.dialect)
|
||||
for _, frag := range tc.fragments {
|
||||
require.Contains(t, stmt.SQL, frag, "dialect %s", tc.dialect)
|
||||
}
|
||||
require.Equal(t, []any{"TODO%"}, stmt.Args, "dialect %s", tc.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderEndsWithPerDialect(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
for _, dialect := range []DialectName{DialectSQLite, DialectPostgres, DialectMySQL} {
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `content.endsWith(".md")`, RenderOptions{Dialect: dialect})
|
||||
require.NoError(t, err, dialect)
|
||||
require.Equal(t, []any{"%.md"}, stmt.Args, "dialect %s", dialect)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderMatchesPerDialect(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
cases := []struct {
|
||||
dialect DialectName
|
||||
fragment string
|
||||
}{
|
||||
{DialectSQLite, "`memo`.`content` REGEXP ?"},
|
||||
{DialectMySQL, "`memo`.`content` REGEXP ?"},
|
||||
{DialectPostgres, "memo.content ~ $1"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `content.matches("v[0-9]+")`, RenderOptions{Dialect: tc.dialect})
|
||||
require.NoError(t, err, tc.dialect)
|
||||
require.Contains(t, stmt.SQL, tc.fragment, "dialect %s", tc.dialect)
|
||||
require.Equal(t, []any{"v[0-9]+"}, stmt.Args, "dialect %s", tc.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderTagsAllPerDialect(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
cases := []struct {
|
||||
dialect DialectName
|
||||
fragments []string
|
||||
}{
|
||||
{DialectSQLite, []string{"NOT EXISTS", "json_each(", "!= '[]'", "memos_unicode_lower(value)"}},
|
||||
{DialectPostgres, []string{"NOT EXISTS", "jsonb_array_elements_text(", "jsonb_array_length(", "value ILIKE"}},
|
||||
{DialectMySQL, []string{"NOT EXISTS", "JSON_TABLE(", "JSON_LENGTH(", "value LIKE"}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `tags.all(t, t.startsWith("work/"))`, RenderOptions{Dialect: tc.dialect})
|
||||
require.NoError(t, err, tc.dialect)
|
||||
for _, frag := range tc.fragments {
|
||||
require.Contains(t, stmt.SQL, frag, "dialect %s", tc.dialect)
|
||||
}
|
||||
require.Equal(t, []any{"work/%"}, stmt.Args, "dialect %s", tc.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderTextMatchEscaping(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
// Both % and _ in the value must be escaped so they match literally.
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `content.contains("a%b_c")`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []any{`%a\%b\_c%`}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestRenderAllRejectsUnsupportedPredicate(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
// size() is not a valid per-element predicate inside all().
|
||||
_, err = engine.CompileToStatement(context.Background(), `tags.all(t, size(t) > 2)`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.Error(t, err)
|
||||
}
|
||||
@@ -0,0 +1,377 @@
|
||||
package filter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Arithmetic folding: division and modulo
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestCompileDivisionFolds(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `creator_id == 100 / 10`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []any{int64(10)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileModuloFolds(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `creator_id == 17 % 5`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []any{int64(2)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileDivisionByZeroErrors(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
_, err = engine.Compile(context.Background(), `creator_id == 10 / 0`)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// size() on scalar string fields -> SQL length
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestCompileSizeOnContentRendersLengthPerDialect(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
cases := []struct {
|
||||
dialect DialectName
|
||||
fragment string
|
||||
}{
|
||||
{DialectSQLite, "LENGTH("},
|
||||
{DialectPostgres, "LENGTH("},
|
||||
{DialectMySQL, "CHAR_LENGTH("},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `size(content) > 5`, RenderOptions{Dialect: tc.dialect})
|
||||
require.NoError(t, err, tc.dialect)
|
||||
require.Contains(t, stmt.SQL, tc.fragment, "dialect %s", tc.dialect)
|
||||
require.Equal(t, []any{int64(5)}, stmt.Args, "dialect %s", tc.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Timestamp accessor methods (getFullYear, getMonth, ...)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestCompileTimestampAccessorsPerDialect(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
filter string
|
||||
dialect DialectName
|
||||
fragments []string
|
||||
arg int64
|
||||
}{
|
||||
// getFullYear == 2024
|
||||
{"sqlite year", `created_ts.getFullYear() == 2024`, DialectSQLite, []string{"strftime('%Y'", "'unixepoch'"}, 2024},
|
||||
{"pg year", `created_ts.getFullYear() == 2024`, DialectPostgres, []string{"EXTRACT(YEAR FROM to_timestamp(", "AT TIME ZONE 'UTC'"}, 2024},
|
||||
{"mysql year", `created_ts.getFullYear() == 2024`, DialectMySQL, []string{"YEAR(`memo`.`created_ts`)"}, 2024},
|
||||
// getMonth is 0-based -> SQL must subtract 1
|
||||
{"sqlite month", `created_ts.getMonth() == 5`, DialectSQLite, []string{"strftime('%m'", "- 1)"}, 5},
|
||||
{"pg month", `created_ts.getMonth() == 5`, DialectPostgres, []string{"EXTRACT(MONTH FROM", "- 1)"}, 5},
|
||||
{"mysql month", `created_ts.getMonth() == 5`, DialectMySQL, []string{"MONTH(`memo`.`created_ts`)", "- 1)"}, 5},
|
||||
// getDayOfWeek 0=Sunday -> MySQL DAYOFWEEK is 1-based and must subtract 1
|
||||
{"mysql dow", `created_ts.getDayOfWeek() == 0`, DialectMySQL, []string{"DAYOFWEEK(`memo`.`created_ts`)", "- 1)"}, 0},
|
||||
{"sqlite dow", `created_ts.getDayOfWeek() == 0`, DialectSQLite, []string{"strftime('%w'"}, 0},
|
||||
// getDate is 1-based -> no offset
|
||||
{"sqlite date", `created_ts.getDate() == 22`, DialectSQLite, []string{"strftime('%d'"}, 22},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
stmt, err := engine.CompileToStatement(context.Background(), tc.filter, RenderOptions{Dialect: tc.dialect})
|
||||
require.NoError(t, err, tc.name)
|
||||
for _, frag := range tc.fragments {
|
||||
require.Contains(t, stmt.SQL, frag, tc.name)
|
||||
}
|
||||
require.Equal(t, []any{tc.arg}, stmt.Args, tc.name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompileNowAccessorsFoldToInjectedClock(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// 2026-07-07T10:30:45Z, a Tuesday (year day 188).
|
||||
engine := memoEngineAt(t, time.Date(2026, time.July, 7, 10, 30, 45, 0, time.UTC).Unix())
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
filter string
|
||||
args []any
|
||||
}{
|
||||
{"on this day", `created_ts.getMonth() == now.getMonth() && created_ts.getDate() == now.getDate()`, []any{int64(6), int64(7)}},
|
||||
{"year", `created_ts.getFullYear() < now.getFullYear()`, []any{int64(2026)}},
|
||||
{"day of month", `created_ts.getDayOfMonth() == now.getDayOfMonth()`, []any{int64(6)}},
|
||||
{"day of week", `created_ts.getDayOfWeek() == now.getDayOfWeek()`, []any{int64(2)}},
|
||||
{"day of year", `created_ts.getDayOfYear() == now.getDayOfYear()`, []any{int64(187)}},
|
||||
{"clock parts", `created_ts.getHours() == now.getHours() || created_ts.getMinutes() == now.getMinutes() || created_ts.getSeconds() == now.getSeconds()`, []any{int64(10), int64(30), int64(45)}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
stmt, err := engine.CompileToStatement(context.Background(), tc.filter, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err, tc.name)
|
||||
require.Equal(t, tc.args, stmt.Args, tc.name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompileNowAccessorRejectsTimezoneArg(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
_, err = engine.Compile(context.Background(), `created_ts.getMonth() == now.getMonth("America/New_York")`)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestCompileOnThisDayPerDialect(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const onThisDay = `created_ts.getMonth() == now.getMonth() && created_ts.getDate() == now.getDate()`
|
||||
engine := memoEngineAt(t, time.Date(2026, time.July, 7, 10, 30, 45, 0, time.UTC).Unix())
|
||||
|
||||
cases := []struct {
|
||||
dialect DialectName
|
||||
fragments []string
|
||||
}{
|
||||
{DialectSQLite, []string{"strftime('%m'", "strftime('%d'", "'unixepoch'"}},
|
||||
{DialectPostgres, []string{"EXTRACT(MONTH FROM", "EXTRACT(DAY FROM"}},
|
||||
{DialectMySQL, []string{"MONTH(`memo`.`created_ts`)", "DAYOFMONTH(`memo`.`created_ts`)"}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
stmt, err := engine.CompileToStatement(context.Background(), onThisDay, RenderOptions{Dialect: tc.dialect})
|
||||
require.NoError(t, err, tc.dialect)
|
||||
for _, frag := range tc.fragments {
|
||||
require.Contains(t, stmt.SQL, frag, "dialect %s", tc.dialect)
|
||||
}
|
||||
require.Equal(t, []any{int64(6), int64(7)}, stmt.Args, "dialect %s", tc.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompileNowAccessorsAtBoundaryDates(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Every accessor folded at calendar edges: year rollover, leap day, Sunday.
|
||||
allAccessors := `created_ts.getFullYear() == now.getFullYear() ` +
|
||||
`&& created_ts.getMonth() == now.getMonth() ` +
|
||||
`&& created_ts.getDate() == now.getDate() ` +
|
||||
`&& created_ts.getDayOfMonth() == now.getDayOfMonth() ` +
|
||||
`&& created_ts.getDayOfWeek() == now.getDayOfWeek() ` +
|
||||
`&& created_ts.getDayOfYear() == now.getDayOfYear() ` +
|
||||
`&& created_ts.getHours() == now.getHours() ` +
|
||||
`&& created_ts.getMinutes() == now.getMinutes() ` +
|
||||
`&& created_ts.getSeconds() == now.getSeconds()`
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
clock time.Time
|
||||
// year, month, date, dayOfMonth, dayOfWeek, dayOfYear, hours, minutes, seconds
|
||||
args []any
|
||||
}{
|
||||
{
|
||||
"new year's eve",
|
||||
time.Date(2026, time.December, 31, 23, 59, 59, 0, time.UTC),
|
||||
[]any{int64(2026), int64(11), int64(31), int64(30), int64(4), int64(364), int64(23), int64(59), int64(59)},
|
||||
},
|
||||
{
|
||||
"new year's day",
|
||||
time.Date(2027, time.January, 1, 0, 0, 0, 0, time.UTC),
|
||||
[]any{int64(2027), int64(0), int64(1), int64(0), int64(5), int64(0), int64(0), int64(0), int64(0)},
|
||||
},
|
||||
{
|
||||
"leap day",
|
||||
time.Date(2028, time.February, 29, 12, 0, 0, 0, time.UTC),
|
||||
[]any{int64(2028), int64(1), int64(29), int64(28), int64(2), int64(59), int64(12), int64(0), int64(0)},
|
||||
},
|
||||
{
|
||||
"sunday",
|
||||
time.Date(2026, time.July, 5, 8, 15, 30, 0, time.UTC),
|
||||
[]any{int64(2026), int64(6), int64(5), int64(4), int64(0), int64(185), int64(8), int64(15), int64(30)},
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
engine := memoEngineAt(t, tc.clock.Unix())
|
||||
stmt, err := engine.CompileToStatement(context.Background(), allAccessors, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err, tc.name)
|
||||
require.Equal(t, tc.args, stmt.Args, tc.name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompileNowAccessorFoldsInUTC(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// A clock in UTC+8 at 05:00 on July 8 is still July 7, 21:00 in UTC;
|
||||
// folding must not leak the clock's zone.
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
engine.nowFunc = func() time.Time {
|
||||
return time.Date(2026, time.July, 8, 5, 0, 0, 0, time.FixedZone("UTC+8", 8*3600))
|
||||
}
|
||||
|
||||
stmt, err := engine.CompileToStatement(
|
||||
context.Background(),
|
||||
`created_ts.getDate() == now.getDate() && created_ts.getHours() == now.getHours()`,
|
||||
RenderOptions{Dialect: DialectSQLite},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []any{int64(7), int64(21)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileNowAccessorOnLeftSide(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine := memoEngineAt(t, time.Date(2026, time.July, 7, 10, 30, 45, 0, time.UTC).Unix())
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `now.getMonth() == created_ts.getMonth()`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, stmt.SQL, "strftime('%m'")
|
||||
require.Equal(t, []any{int64(6)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileNowAccessorMirrorsOrderingOperator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Swapping the literal to the right must flip < to > to keep the meaning.
|
||||
engine := memoEngineAt(t, time.Date(2026, time.July, 7, 10, 30, 45, 0, time.UTC).Unix())
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `now.getHours() < created_ts.getHours()`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, stmt.SQL, "> ?")
|
||||
require.Equal(t, []any{int64(10)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileNowAccessorAgainstLiteralFoldsToConstant(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Both sides fold to literals; the comparison folds to a constant condition.
|
||||
engine := memoEngineAt(t, time.Date(2026, time.July, 7, 10, 30, 45, 0, time.UTC).Unix())
|
||||
|
||||
// True → trivial filter (empty SQL, matches everything).
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `now.getFullYear() >= 2026`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, stmt.SQL)
|
||||
require.Empty(t, stmt.Args)
|
||||
|
||||
// False → unsatisfiable filter.
|
||||
stmt, err = engine.CompileToStatement(context.Background(), `now.getFullYear() < 2026`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "1 = 0", stmt.SQL)
|
||||
require.Empty(t, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileAttachmentNowAccessors(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewAttachmentSchema())
|
||||
require.NoError(t, err)
|
||||
engine.nowFunc = fixedClock(time.Date(2026, time.July, 7, 10, 30, 45, 0, time.UTC).Unix())
|
||||
|
||||
stmt, err := engine.CompileToStatement(
|
||||
context.Background(),
|
||||
`create_time.getMonth() == now.getMonth() && create_time.getDate() == now.getDate()`,
|
||||
RenderOptions{Dialect: DialectSQLite},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, stmt.SQL, "`attachment`.`created_ts`")
|
||||
require.Equal(t, []any{int64(6), int64(7)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileTimestampAccessorRejectsTimezoneArg(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
_, err = engine.Compile(context.Background(), `created_ts.getHours("America/New_York") == 9`)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestCompileTimestampAccessorRejectsNonTimestampField(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
// content is a string, not a timestamp.
|
||||
_, err = engine.Compile(context.Background(), `content.getFullYear() == 2024`)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ext.Sets(): sets.contains / sets.intersects / sets.equivalent over tags
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestCompileSetsContainsRendersAndOfMemberships(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `sets.contains(tags, ["a", "b"])`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, stmt.SQL, " AND ")
|
||||
require.Len(t, stmt.Args, 2)
|
||||
}
|
||||
|
||||
func TestCompileSetsIntersectsRendersOrOfMemberships(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `sets.intersects(tags, ["a", "b"])`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, stmt.SQL, " OR ")
|
||||
require.Len(t, stmt.Args, 2)
|
||||
}
|
||||
|
||||
func TestCompileSetsEquivalentAddsLengthCheck(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `sets.equivalent(tags, ["a", "b"])`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, stmt.SQL, "JSON_ARRAY_LENGTH")
|
||||
require.Contains(t, stmt.Args, int64(2))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// exists_one() comprehension on tags
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestCompileExistsOnePerDialect(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
cases := []struct {
|
||||
dialect DialectName
|
||||
fragments []string
|
||||
}{
|
||||
{DialectSQLite, []string{"COUNT(", "json_each(", ") = 1"}},
|
||||
{DialectPostgres, []string{"COUNT(", "jsonb_array_elements_text(", ") = 1"}},
|
||||
{DialectMySQL, []string{"COUNT(", "JSON_TABLE(", ") = 1"}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `tags.exists_one(t, t == "urgent")`, RenderOptions{Dialect: tc.dialect})
|
||||
require.NoError(t, err, tc.dialect)
|
||||
for _, frag := range tc.fragments {
|
||||
require.Contains(t, stmt.SQL, frag, "dialect %s", tc.dialect)
|
||||
}
|
||||
require.Equal(t, []any{"urgent"}, stmt.Args, "dialect %s", tc.dialect)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
package filter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// AppendConditions compiles the provided filters and appends the resulting SQL fragments and args.
|
||||
func AppendConditions(ctx context.Context, engine *Engine, filters []string, dialect DialectName, where *[]string, args *[]any) error {
|
||||
for _, filterStr := range filters {
|
||||
stmt, err := engine.CompileToStatement(ctx, filterStr, RenderOptions{
|
||||
Dialect: dialect,
|
||||
PlaceholderOffset: len(*args),
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if stmt.SQL == "" {
|
||||
continue
|
||||
}
|
||||
*where = append(*where, fmt.Sprintf("(%s)", stmt.SQL))
|
||||
*args = append(*args, stmt.Args...)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
package filter
|
||||
|
||||
// Condition represents a boolean expression derived from the CEL filter.
|
||||
type Condition interface {
|
||||
isCondition()
|
||||
}
|
||||
|
||||
// LogicalOperator enumerates the supported logical operators.
|
||||
type LogicalOperator string
|
||||
|
||||
const (
|
||||
LogicalAnd LogicalOperator = "AND"
|
||||
LogicalOr LogicalOperator = "OR"
|
||||
)
|
||||
|
||||
// LogicalCondition composes two conditions with a logical operator.
|
||||
type LogicalCondition struct {
|
||||
Operator LogicalOperator
|
||||
Left Condition
|
||||
Right Condition
|
||||
}
|
||||
|
||||
func (*LogicalCondition) isCondition() {}
|
||||
|
||||
// NotCondition negates a child condition.
|
||||
type NotCondition struct {
|
||||
Expr Condition
|
||||
}
|
||||
|
||||
func (*NotCondition) isCondition() {}
|
||||
|
||||
// FieldPredicateCondition asserts that a field evaluates to true.
|
||||
type FieldPredicateCondition struct {
|
||||
Field string
|
||||
}
|
||||
|
||||
func (*FieldPredicateCondition) isCondition() {}
|
||||
|
||||
// ComparisonOperator lists supported comparison operators.
|
||||
type ComparisonOperator string
|
||||
|
||||
const (
|
||||
CompareEq ComparisonOperator = "="
|
||||
CompareNeq ComparisonOperator = "!="
|
||||
CompareLt ComparisonOperator = "<"
|
||||
CompareLte ComparisonOperator = "<="
|
||||
CompareGt ComparisonOperator = ">"
|
||||
CompareGte ComparisonOperator = ">="
|
||||
)
|
||||
|
||||
// ComparisonCondition represents a binary comparison.
|
||||
type ComparisonCondition struct {
|
||||
Left ValueExpr
|
||||
Operator ComparisonOperator
|
||||
Right ValueExpr
|
||||
}
|
||||
|
||||
func (*ComparisonCondition) isCondition() {}
|
||||
|
||||
// InCondition represents an IN predicate with literal list values.
|
||||
type InCondition struct {
|
||||
Left ValueExpr
|
||||
Values []ValueExpr
|
||||
}
|
||||
|
||||
func (*InCondition) isCondition() {}
|
||||
|
||||
// ElementInCondition represents the CEL syntax `"value" in field`.
|
||||
type ElementInCondition struct {
|
||||
Element ValueExpr
|
||||
Field string
|
||||
}
|
||||
|
||||
func (*ElementInCondition) isCondition() {}
|
||||
|
||||
// TextMatchMode enumerates LIKE-based string match modes.
|
||||
type TextMatchMode string
|
||||
|
||||
const (
|
||||
TextMatchContains TextMatchMode = "contains"
|
||||
TextMatchPrefix TextMatchMode = "prefix"
|
||||
TextMatchSuffix TextMatchMode = "suffix"
|
||||
)
|
||||
|
||||
// TextMatchCondition models a case-insensitive LIKE match on a scalar string field
|
||||
// (content.contains/startsWith/endsWith).
|
||||
type TextMatchCondition struct {
|
||||
Field string
|
||||
Mode TextMatchMode
|
||||
Value string
|
||||
}
|
||||
|
||||
func (*TextMatchCondition) isCondition() {}
|
||||
|
||||
// RegexCondition models field.matches("pattern") on a string field.
|
||||
type RegexCondition struct {
|
||||
Field string
|
||||
Pattern string
|
||||
}
|
||||
|
||||
func (*RegexCondition) isCondition() {}
|
||||
|
||||
// ConstantCondition captures a literal boolean outcome.
|
||||
type ConstantCondition struct {
|
||||
Value bool
|
||||
}
|
||||
|
||||
func (*ConstantCondition) isCondition() {}
|
||||
|
||||
// ValueExpr models arithmetic or scalar expressions whose result feeds a comparison.
|
||||
type ValueExpr interface {
|
||||
isValueExpr()
|
||||
}
|
||||
|
||||
// FieldRef references a named schema field.
|
||||
type FieldRef struct {
|
||||
Name string
|
||||
}
|
||||
|
||||
func (*FieldRef) isValueExpr() {}
|
||||
|
||||
// LiteralValue holds a literal scalar.
|
||||
type LiteralValue struct {
|
||||
Value interface{}
|
||||
}
|
||||
|
||||
func (*LiteralValue) isValueExpr() {}
|
||||
|
||||
// FunctionValue captures simple function calls like size(tags).
|
||||
type FunctionValue struct {
|
||||
Name string
|
||||
Args []ValueExpr
|
||||
}
|
||||
|
||||
func (*FunctionValue) isValueExpr() {}
|
||||
|
||||
// FieldAccessorValue captures a CEL timestamp accessor on a field, such as
|
||||
// created_ts.getMonth(). It renders to a dialect-specific date-part extraction.
|
||||
type FieldAccessorValue struct {
|
||||
Field string
|
||||
Accessor string // e.g. "getFullYear", "getMonth"
|
||||
}
|
||||
|
||||
func (*FieldAccessorValue) isValueExpr() {}
|
||||
|
||||
// ListComprehensionCondition represents CEL macros like exists(), all(), filter().
|
||||
type ListComprehensionCondition struct {
|
||||
Kind ComprehensionKind
|
||||
Field string // The list field to iterate over (e.g., "tags")
|
||||
IterVar string // The iteration variable name (e.g., "t")
|
||||
Predicate PredicateExpr // The predicate to evaluate for each element
|
||||
}
|
||||
|
||||
func (*ListComprehensionCondition) isCondition() {}
|
||||
|
||||
// ComprehensionKind enumerates the types of list comprehensions.
|
||||
type ComprehensionKind string
|
||||
|
||||
const (
|
||||
ComprehensionExists ComprehensionKind = "exists"
|
||||
ComprehensionAll ComprehensionKind = "all"
|
||||
ComprehensionExistsOne ComprehensionKind = "exists_one"
|
||||
)
|
||||
|
||||
// PredicateExpr represents predicates used in comprehensions.
|
||||
type PredicateExpr interface {
|
||||
isPredicateExpr()
|
||||
}
|
||||
|
||||
// StartsWithPredicate represents t.startsWith("prefix").
|
||||
type StartsWithPredicate struct {
|
||||
Prefix string
|
||||
}
|
||||
|
||||
func (*StartsWithPredicate) isPredicateExpr() {}
|
||||
|
||||
// EndsWithPredicate represents t.endsWith("suffix").
|
||||
type EndsWithPredicate struct {
|
||||
Suffix string
|
||||
}
|
||||
|
||||
func (*EndsWithPredicate) isPredicateExpr() {}
|
||||
|
||||
// ContainsPredicate represents t.contains("substring").
|
||||
type ContainsPredicate struct {
|
||||
Substring string
|
||||
}
|
||||
|
||||
func (*ContainsPredicate) isPredicateExpr() {}
|
||||
|
||||
// EqualsPredicate represents t == "value".
|
||||
type EqualsPredicate struct {
|
||||
Value string
|
||||
}
|
||||
|
||||
func (*EqualsPredicate) isPredicateExpr() {}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,980 @@
|
||||
package filter
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
type renderer struct {
|
||||
schema Schema
|
||||
dialect DialectName
|
||||
placeholderOffset int
|
||||
placeholderCounter int
|
||||
args []any
|
||||
}
|
||||
|
||||
type renderResult struct {
|
||||
sql string
|
||||
trivial bool
|
||||
unsatisfiable bool
|
||||
}
|
||||
|
||||
func newRenderer(schema Schema, opts RenderOptions) *renderer {
|
||||
return &renderer{
|
||||
schema: schema,
|
||||
dialect: opts.Dialect,
|
||||
placeholderOffset: opts.PlaceholderOffset,
|
||||
}
|
||||
}
|
||||
|
||||
func (r *renderer) Render(cond Condition) (Statement, error) {
|
||||
result, err := r.renderCondition(cond)
|
||||
if err != nil {
|
||||
return Statement{}, err
|
||||
}
|
||||
args := r.args
|
||||
if args == nil {
|
||||
args = []any{}
|
||||
}
|
||||
|
||||
switch {
|
||||
case result.unsatisfiable:
|
||||
return Statement{
|
||||
SQL: "1 = 0",
|
||||
Args: args,
|
||||
}, nil
|
||||
case result.trivial:
|
||||
return Statement{
|
||||
SQL: "",
|
||||
Args: args,
|
||||
}, nil
|
||||
default:
|
||||
return Statement{
|
||||
SQL: result.sql,
|
||||
Args: args,
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (r *renderer) renderCondition(cond Condition) (renderResult, error) {
|
||||
switch c := cond.(type) {
|
||||
case *LogicalCondition:
|
||||
return r.renderLogicalCondition(c)
|
||||
case *NotCondition:
|
||||
return r.renderNotCondition(c)
|
||||
case *FieldPredicateCondition:
|
||||
return r.renderFieldPredicate(c)
|
||||
case *ComparisonCondition:
|
||||
return r.renderComparison(c)
|
||||
case *InCondition:
|
||||
return r.renderInCondition(c)
|
||||
case *ElementInCondition:
|
||||
return r.renderElementInCondition(c)
|
||||
case *TextMatchCondition:
|
||||
return r.renderTextMatch(c)
|
||||
case *RegexCondition:
|
||||
return r.renderRegex(c)
|
||||
case *ListComprehensionCondition:
|
||||
return r.renderListComprehension(c)
|
||||
case *ConstantCondition:
|
||||
if c.Value {
|
||||
return renderResult{trivial: true}, nil
|
||||
}
|
||||
return renderResult{sql: "1 = 0", unsatisfiable: true}, nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported condition type %T", c)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *renderer) renderLogicalCondition(cond *LogicalCondition) (renderResult, error) {
|
||||
left, err := r.renderCondition(cond.Left)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
right, err := r.renderCondition(cond.Right)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
|
||||
switch cond.Operator {
|
||||
case LogicalAnd:
|
||||
return combineAnd(left, right), nil
|
||||
case LogicalOr:
|
||||
return combineOr(left, right), nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported logical operator %s", cond.Operator)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *renderer) renderNotCondition(cond *NotCondition) (renderResult, error) {
|
||||
child, err := r.renderCondition(cond.Expr)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
|
||||
if child.trivial {
|
||||
return renderResult{sql: "1 = 0", unsatisfiable: true}, nil
|
||||
}
|
||||
if child.unsatisfiable {
|
||||
return renderResult{trivial: true}, nil
|
||||
}
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("NOT (%s)", child.sql),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *renderer) renderFieldPredicate(cond *FieldPredicateCondition) (renderResult, error) {
|
||||
field, ok := r.schema.Field(cond.Field)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("unknown field %q", cond.Field)
|
||||
}
|
||||
|
||||
switch field.Kind {
|
||||
case FieldKindBoolColumn:
|
||||
column := qualifyColumn(r.dialect, field.Column)
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("%s IS TRUE", column),
|
||||
}, nil
|
||||
case FieldKindJSONBool:
|
||||
sql, err := r.jsonBoolPredicate(field)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
return renderResult{sql: sql}, nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("field %q cannot be used as a predicate", cond.Field)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *renderer) renderComparison(cond *ComparisonCondition) (renderResult, error) {
|
||||
switch left := cond.Left.(type) {
|
||||
case *FieldRef:
|
||||
field, ok := r.schema.Field(left.Name)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("unknown field %q", left.Name)
|
||||
}
|
||||
switch field.Kind {
|
||||
case FieldKindBoolColumn:
|
||||
return r.renderBoolColumnComparison(field, cond.Operator, cond.Right)
|
||||
case FieldKindJSONBool:
|
||||
return r.renderJSONBoolComparison(field, cond.Operator, cond.Right)
|
||||
case FieldKindScalar:
|
||||
return r.renderScalarComparison(field, cond.Operator, cond.Right)
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("field %q does not support comparison", field.Name)
|
||||
}
|
||||
case *FunctionValue:
|
||||
return r.renderFunctionComparison(left, cond.Operator, cond.Right)
|
||||
case *FieldAccessorValue:
|
||||
return r.renderAccessorComparison(left, cond.Operator, cond.Right)
|
||||
default:
|
||||
return renderResult{}, errors.New("comparison must start with a field reference or supported function")
|
||||
}
|
||||
}
|
||||
|
||||
// accessorSpec maps a CEL timestamp accessor to per-dialect SQL date-part tokens
|
||||
// and the offset to subtract so the result matches CEL's base (e.g. CEL months
|
||||
// are 0-based but every dialect reports 1-based, so off=1). off is indexed
|
||||
// [sqlite, postgres, mysql].
|
||||
type accessorSpec struct {
|
||||
sqlite string // strftime format specifier
|
||||
pg string // EXTRACT field
|
||||
mysql string // function name
|
||||
off [3]int
|
||||
}
|
||||
|
||||
var accessorSpecs = map[string]accessorSpec{
|
||||
"getFullYear": {"%Y", "YEAR", "YEAR", [3]int{0, 0, 0}},
|
||||
"getMonth": {"%m", "MONTH", "MONTH", [3]int{1, 1, 1}},
|
||||
"getDate": {"%d", "DAY", "DAYOFMONTH", [3]int{0, 0, 0}},
|
||||
"getDayOfMonth": {"%d", "DAY", "DAYOFMONTH", [3]int{1, 1, 1}},
|
||||
"getDayOfWeek": {"%w", "DOW", "DAYOFWEEK", [3]int{0, 0, 1}},
|
||||
"getDayOfYear": {"%j", "DOY", "DAYOFYEAR", [3]int{1, 1, 1}},
|
||||
"getHours": {"%H", "HOUR", "HOUR", [3]int{0, 0, 0}},
|
||||
"getMinutes": {"%M", "MINUTE", "MINUTE", [3]int{0, 0, 0}},
|
||||
"getSeconds": {"%S", "SECOND", "SECOND", [3]int{0, 0, 0}},
|
||||
}
|
||||
|
||||
func (r *renderer) renderAccessorComparison(acc *FieldAccessorValue, op ComparisonOperator, right ValueExpr) (renderResult, error) {
|
||||
field, ok := r.schema.Field(acc.Field)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("unknown field %q", acc.Field)
|
||||
}
|
||||
value, err := expectNumericLiteral(right)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
expr, err := r.timestampAccessorExpr(field, acc.Accessor)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
placeholder := r.addArg(value)
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("%s %s %s", expr, sqlOperator(op), placeholder),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// timestampAccessorExpr builds a dialect-specific integer expression for a CEL
|
||||
// timestamp accessor. Extraction is UTC on SQLite/Postgres (epoch columns); on
|
||||
// MySQL the TIMESTAMP column is read in the server session time zone.
|
||||
func (r *renderer) timestampAccessorExpr(field Field, accessor string) (string, error) {
|
||||
spec, ok := accessorSpecs[accessor]
|
||||
if !ok {
|
||||
return "", errors.Errorf("unsupported timestamp accessor %q", accessor)
|
||||
}
|
||||
col := qualifyColumn(r.dialect, field.Column)
|
||||
var base string
|
||||
var off int
|
||||
switch r.dialect {
|
||||
case DialectSQLite:
|
||||
base = fmt.Sprintf("CAST(strftime('%s', %s, 'unixepoch') AS INTEGER)", spec.sqlite, col)
|
||||
off = spec.off[0]
|
||||
case DialectPostgres:
|
||||
base = fmt.Sprintf("CAST(EXTRACT(%s FROM to_timestamp(%s) AT TIME ZONE 'UTC') AS INTEGER)", spec.pg, col)
|
||||
off = spec.off[1]
|
||||
case DialectMySQL:
|
||||
base = fmt.Sprintf("%s(%s)", spec.mysql, col)
|
||||
off = spec.off[2]
|
||||
default:
|
||||
return "", errors.Errorf("unsupported dialect %q", r.dialect)
|
||||
}
|
||||
if off != 0 {
|
||||
base = fmt.Sprintf("(%s - %d)", base, off)
|
||||
}
|
||||
return base, nil
|
||||
}
|
||||
|
||||
func (r *renderer) renderFunctionComparison(fn *FunctionValue, op ComparisonOperator, right ValueExpr) (renderResult, error) {
|
||||
if fn.Name != "size" {
|
||||
return renderResult{}, errors.Errorf("unsupported function %s in comparison", fn.Name)
|
||||
}
|
||||
if len(fn.Args) != 1 {
|
||||
return renderResult{}, errors.New("size() expects one argument")
|
||||
}
|
||||
fieldArg, ok := fn.Args[0].(*FieldRef)
|
||||
if !ok {
|
||||
return renderResult{}, errors.New("size() argument must be a field")
|
||||
}
|
||||
|
||||
field, ok := r.schema.Field(fieldArg.Name)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("unknown field %q", fieldArg.Name)
|
||||
}
|
||||
if field.Kind == FieldKindVirtualAlias {
|
||||
field, ok = r.schema.ResolveAlias(fieldArg.Name)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("invalid alias %q", fieldArg.Name)
|
||||
}
|
||||
}
|
||||
|
||||
value, err := expectNumericLiteral(right)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
|
||||
var expr string
|
||||
switch {
|
||||
case field.Kind == FieldKindJSONList:
|
||||
expr = jsonArrayLengthExpr(r.dialect, field)
|
||||
case field.Kind == FieldKindScalar && field.Type == FieldTypeString:
|
||||
expr = stringLengthExpr(r.dialect, field.columnExpr(r.dialect))
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("size() does not support field %q", field.Name)
|
||||
}
|
||||
|
||||
placeholder := r.addArg(value)
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("%s %s %s", expr, sqlOperator(op), placeholder),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// stringLengthExpr returns the character-count expression for a string column.
|
||||
// MySQL's LENGTH counts bytes, so CHAR_LENGTH is used to count characters and
|
||||
// match CEL's size() code-point semantics; SQLite/Postgres LENGTH already counts
|
||||
// characters.
|
||||
func stringLengthExpr(d DialectName, colExpr string) string {
|
||||
if d == DialectMySQL {
|
||||
return fmt.Sprintf("CHAR_LENGTH(%s)", colExpr)
|
||||
}
|
||||
return fmt.Sprintf("LENGTH(%s)", colExpr)
|
||||
}
|
||||
|
||||
func (r *renderer) renderScalarComparison(field Field, op ComparisonOperator, right ValueExpr) (renderResult, error) {
|
||||
lit, err := expectLiteral(right)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
|
||||
columnExpr := field.columnExpr(r.dialect)
|
||||
if lit == nil {
|
||||
switch op {
|
||||
case CompareEq:
|
||||
return renderResult{sql: fmt.Sprintf("%s IS NULL", columnExpr)}, nil
|
||||
case CompareNeq:
|
||||
return renderResult{sql: fmt.Sprintf("%s IS NOT NULL", columnExpr)}, nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("operator %s not supported for null comparison", op)
|
||||
}
|
||||
}
|
||||
|
||||
placeholder := ""
|
||||
switch field.Type {
|
||||
case FieldTypeString:
|
||||
value, ok := lit.(string)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("field %q expects string value", field.Name)
|
||||
}
|
||||
placeholder = r.addArg(value)
|
||||
case FieldTypeInt, FieldTypeTimestamp:
|
||||
num, err := toInt64(lit)
|
||||
if err != nil {
|
||||
return renderResult{}, errors.Wrapf(err, "field %q expects integer value", field.Name)
|
||||
}
|
||||
placeholder = r.addArg(num)
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported data type %q for field %s", field.Type, field.Name)
|
||||
}
|
||||
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("%s %s %s", columnExpr, sqlOperator(op), placeholder),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *renderer) renderBoolColumnComparison(field Field, op ComparisonOperator, right ValueExpr) (renderResult, error) {
|
||||
value, err := expectBool(right)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
placeholder := r.addBoolArg(value)
|
||||
column := qualifyColumn(r.dialect, field.Column)
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("%s %s %s", column, sqlOperator(op), placeholder),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *renderer) renderJSONBoolComparison(field Field, op ComparisonOperator, right ValueExpr) (renderResult, error) {
|
||||
value, err := expectBool(right)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
|
||||
jsonExpr := jsonExtractExpr(r.dialect, field)
|
||||
switch r.dialect {
|
||||
case DialectSQLite:
|
||||
switch op {
|
||||
case CompareEq:
|
||||
if field.Name == "has_task_list" {
|
||||
target := "0"
|
||||
if value {
|
||||
target = "1"
|
||||
}
|
||||
return renderResult{sql: fmt.Sprintf("%s = %s", jsonExpr, target)}, nil
|
||||
}
|
||||
if value {
|
||||
return renderResult{sql: fmt.Sprintf("%s IS TRUE", jsonExpr)}, nil
|
||||
}
|
||||
return renderResult{sql: fmt.Sprintf("NOT(%s IS TRUE)", jsonExpr)}, nil
|
||||
case CompareNeq:
|
||||
if field.Name == "has_task_list" {
|
||||
target := "0"
|
||||
if value {
|
||||
target = "1"
|
||||
}
|
||||
return renderResult{sql: fmt.Sprintf("%s != %s", jsonExpr, target)}, nil
|
||||
}
|
||||
if value {
|
||||
return renderResult{sql: fmt.Sprintf("NOT(%s IS TRUE)", jsonExpr)}, nil
|
||||
}
|
||||
return renderResult{sql: fmt.Sprintf("%s IS TRUE", jsonExpr)}, nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("operator %s not supported for boolean JSON field", op)
|
||||
}
|
||||
case DialectMySQL:
|
||||
boolStr := "false"
|
||||
if value {
|
||||
boolStr = "true"
|
||||
}
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("%s %s CAST('%s' AS JSON)", jsonExpr, sqlOperator(op), boolStr),
|
||||
}, nil
|
||||
case DialectPostgres:
|
||||
placeholder := r.addArg(value)
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("(%s)::boolean %s %s", jsonExpr, sqlOperator(op), placeholder),
|
||||
}, nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported dialect %s", r.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *renderer) renderInCondition(cond *InCondition) (renderResult, error) {
|
||||
fieldRef, ok := cond.Left.(*FieldRef)
|
||||
if !ok {
|
||||
return renderResult{}, errors.New("IN operator requires a field on the left-hand side")
|
||||
}
|
||||
|
||||
if fieldRef.Name == "tag" {
|
||||
return r.renderTagInList(cond.Values)
|
||||
}
|
||||
|
||||
field, ok := r.schema.Field(fieldRef.Name)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("unknown field %q", fieldRef.Name)
|
||||
}
|
||||
|
||||
if field.Kind != FieldKindScalar {
|
||||
return renderResult{}, errors.Errorf("field %q does not support IN()", fieldRef.Name)
|
||||
}
|
||||
|
||||
return r.renderScalarInCondition(field, cond.Values)
|
||||
}
|
||||
|
||||
func (r *renderer) renderTagInList(values []ValueExpr) (renderResult, error) {
|
||||
field, ok := r.schema.ResolveAlias("tag")
|
||||
if !ok {
|
||||
return renderResult{}, errors.New("tag attribute is not configured")
|
||||
}
|
||||
|
||||
conditions := make([]string, 0, len(values))
|
||||
for _, v := range values {
|
||||
lit, err := expectLiteral(v)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
str, ok := lit.(string)
|
||||
if !ok {
|
||||
return renderResult{}, errors.New("tags must be compared with string literals")
|
||||
}
|
||||
|
||||
switch r.dialect {
|
||||
case DialectSQLite:
|
||||
// Support hierarchical tags: match exact tag OR tags with this prefix (e.g., "book" matches "book" and "book/something")
|
||||
exactMatch := fmt.Sprintf("%s LIKE %s", jsonArrayExpr(r.dialect, field), r.addArg(fmt.Sprintf(`%%"%s"%%`, str)))
|
||||
prefixMatch := fmt.Sprintf("%s LIKE %s", jsonArrayExpr(r.dialect, field), r.addArg(fmt.Sprintf(`%%"%s/%%`, str)))
|
||||
expr := fmt.Sprintf("(%s OR %s)", exactMatch, prefixMatch)
|
||||
conditions = append(conditions, expr)
|
||||
case DialectMySQL:
|
||||
// Support hierarchical tags: match exact tag OR tags with this prefix
|
||||
exactMatch := fmt.Sprintf("JSON_CONTAINS(%s, %s)", jsonArrayExpr(r.dialect, field), r.addArg(fmt.Sprintf(`"%s"`, str)))
|
||||
prefixMatch := fmt.Sprintf("%s LIKE %s", jsonArrayExpr(r.dialect, field), r.addArg(fmt.Sprintf(`%%"%s/%%`, str)))
|
||||
expr := fmt.Sprintf("(%s OR %s)", exactMatch, prefixMatch)
|
||||
conditions = append(conditions, expr)
|
||||
case DialectPostgres:
|
||||
// Support hierarchical tags: match exact tag OR tags with this prefix
|
||||
exactMatch := fmt.Sprintf("%s @> jsonb_build_array(%s::json)", jsonArrayExpr(r.dialect, field), r.addArg(fmt.Sprintf(`"%s"`, str)))
|
||||
prefixMatch := fmt.Sprintf("(%s)::text LIKE %s", jsonArrayExpr(r.dialect, field), r.addArg(fmt.Sprintf(`%%"%s/%%`, str)))
|
||||
expr := fmt.Sprintf("(%s OR %s)", exactMatch, prefixMatch)
|
||||
conditions = append(conditions, expr)
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported dialect %s", r.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
if len(conditions) == 1 {
|
||||
return renderResult{sql: conditions[0]}, nil
|
||||
}
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("(%s)", strings.Join(conditions, " OR ")),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *renderer) renderElementInCondition(cond *ElementInCondition) (renderResult, error) {
|
||||
field, ok := r.schema.Field(cond.Field)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("unknown field %q", cond.Field)
|
||||
}
|
||||
if field.Kind != FieldKindJSONList {
|
||||
return renderResult{}, errors.Errorf("field %q is not a tag list", cond.Field)
|
||||
}
|
||||
|
||||
lit, err := expectLiteral(cond.Element)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
str, ok := lit.(string)
|
||||
if !ok {
|
||||
return renderResult{}, errors.New("tags membership requires string literal")
|
||||
}
|
||||
|
||||
switch r.dialect {
|
||||
case DialectSQLite:
|
||||
sql := fmt.Sprintf("%s LIKE %s", jsonArrayExpr(r.dialect, field), r.addArg(fmt.Sprintf(`%%"%s"%%`, str)))
|
||||
return renderResult{sql: sql}, nil
|
||||
case DialectMySQL:
|
||||
sql := fmt.Sprintf("JSON_CONTAINS(%s, %s)", jsonArrayExpr(r.dialect, field), r.addArg(fmt.Sprintf(`"%s"`, str)))
|
||||
return renderResult{sql: sql}, nil
|
||||
case DialectPostgres:
|
||||
sql := fmt.Sprintf("%s @> jsonb_build_array(%s::json)", jsonArrayExpr(r.dialect, field), r.addArg(fmt.Sprintf(`"%s"`, str)))
|
||||
return renderResult{sql: sql}, nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported dialect %s", r.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *renderer) renderScalarInCondition(field Field, values []ValueExpr) (renderResult, error) {
|
||||
placeholders := make([]string, 0, len(values))
|
||||
|
||||
for _, v := range values {
|
||||
lit, err := expectLiteral(v)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
switch field.Type {
|
||||
case FieldTypeString:
|
||||
str, ok := lit.(string)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("field %q expects string values", field.Name)
|
||||
}
|
||||
placeholders = append(placeholders, r.addArg(str))
|
||||
case FieldTypeInt:
|
||||
num, err := toInt64(lit)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
placeholders = append(placeholders, r.addArg(num))
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("field %q does not support IN() comparisons", field.Name)
|
||||
}
|
||||
}
|
||||
|
||||
column := field.columnExpr(r.dialect)
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("%s IN (%s)", column, strings.Join(placeholders, ",")),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *renderer) renderTextMatch(cond *TextMatchCondition) (renderResult, error) {
|
||||
field, ok := r.schema.Field(cond.Field)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("unknown field %q", cond.Field)
|
||||
}
|
||||
column := field.columnExpr(r.dialect)
|
||||
pattern := likePattern(cond.Mode, cond.Value)
|
||||
return renderResult{sql: r.foldedLike(column, pattern)}, nil
|
||||
}
|
||||
|
||||
func (r *renderer) renderRegex(cond *RegexCondition) (renderResult, error) {
|
||||
field, ok := r.schema.Field(cond.Field)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("unknown field %q", cond.Field)
|
||||
}
|
||||
column := field.columnExpr(r.dialect)
|
||||
switch r.dialect {
|
||||
case DialectPostgres:
|
||||
// POSIX regex match operator.
|
||||
return renderResult{sql: fmt.Sprintf("%s ~ %s", column, r.addArg(cond.Pattern))}, nil
|
||||
case DialectMySQL, DialectSQLite:
|
||||
// MySQL has a native REGEXP operator; SQLite uses the registered regexp() function.
|
||||
return renderResult{sql: fmt.Sprintf("%s REGEXP %s", column, r.addArg(cond.Pattern))}, nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported dialect %s", r.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
// foldedLike renders a case-insensitive LIKE comparison of colExpr against a
|
||||
// (already metacharacter-escaped) pattern, using each dialect's case-folding.
|
||||
func (r *renderer) foldedLike(colExpr, pattern string) string {
|
||||
switch r.dialect {
|
||||
case DialectSQLite:
|
||||
// memos_unicode_lower gives Unicode-aware folding; ESCAPE '\' is required
|
||||
// because SQLite has no default LIKE escape character.
|
||||
return fmt.Sprintf(`memos_unicode_lower(%s) LIKE memos_unicode_lower(%s) ESCAPE '\'`, colExpr, r.addArg(pattern))
|
||||
case DialectPostgres:
|
||||
// ILIKE is case-insensitive; backslash is the default escape character.
|
||||
return fmt.Sprintf("%s ILIKE %s", colExpr, r.addArg(pattern))
|
||||
default: // MySQL: default collation is case-insensitive; backslash is the default escape.
|
||||
return fmt.Sprintf("%s LIKE %s", colExpr, r.addArg(pattern))
|
||||
}
|
||||
}
|
||||
|
||||
// likePattern escapes LIKE metacharacters in value and wraps it for the mode.
|
||||
func likePattern(mode TextMatchMode, value string) string {
|
||||
escaped := escapeLikeLiteral(value)
|
||||
switch mode {
|
||||
case TextMatchPrefix:
|
||||
return escaped + "%"
|
||||
case TextMatchSuffix:
|
||||
return "%" + escaped
|
||||
default:
|
||||
return "%" + escaped + "%"
|
||||
}
|
||||
}
|
||||
|
||||
// escapeLikeLiteral escapes the LIKE metacharacters \, %, and _ so user input
|
||||
// is matched literally. Backslash is the escape character on all three dialects.
|
||||
func escapeLikeLiteral(s string) string {
|
||||
return strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`).Replace(s)
|
||||
}
|
||||
|
||||
func (r *renderer) renderListComprehension(cond *ListComprehensionCondition) (renderResult, error) {
|
||||
field, ok := r.schema.Field(cond.Field)
|
||||
if !ok {
|
||||
return renderResult{}, errors.Errorf("unknown field %q", cond.Field)
|
||||
}
|
||||
|
||||
if field.Kind != FieldKindJSONList {
|
||||
return renderResult{}, errors.Errorf("field %q is not a JSON list", cond.Field)
|
||||
}
|
||||
|
||||
if cond.Kind == ComprehensionAll {
|
||||
return r.renderTagAll(field, cond.Predicate)
|
||||
}
|
||||
|
||||
if cond.Kind == ComprehensionExistsOne {
|
||||
return r.renderTagExistsOne(field, cond.Predicate)
|
||||
}
|
||||
|
||||
// Render based on predicate type
|
||||
switch pred := cond.Predicate.(type) {
|
||||
case *EqualsPredicate:
|
||||
return r.renderTagEquals(field, pred.Value, cond.Kind)
|
||||
case *StartsWithPredicate:
|
||||
return r.renderTagStartsWith(field, pred.Prefix, cond.Kind)
|
||||
case *EndsWithPredicate:
|
||||
return r.renderTagEndsWith(field, pred.Suffix, cond.Kind)
|
||||
case *ContainsPredicate:
|
||||
return r.renderTagContains(field, pred.Substring, cond.Kind)
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported predicate type %T in comprehension", pred)
|
||||
}
|
||||
}
|
||||
|
||||
// renderTagAll renders tags.all(t, <pred>): the array is non-empty AND no element
|
||||
// fails the predicate. Element predicates use plain CEL semantics (case-insensitive
|
||||
// for startsWith/endsWith/contains, case-sensitive for ==), evaluated per element.
|
||||
func (r *renderer) renderTagAll(field Field, pred PredicateExpr) (renderResult, error) {
|
||||
arrayExpr := jsonArrayExpr(r.dialect, field)
|
||||
elemCond, err := r.elementPredicateSQL(pred)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
switch r.dialect {
|
||||
case DialectSQLite:
|
||||
nonEmpty := fmt.Sprintf("%s IS NOT NULL AND %s != '[]'", arrayExpr, arrayExpr)
|
||||
sub := fmt.Sprintf("NOT EXISTS (SELECT 1 FROM json_each(%s) WHERE NOT (%s))", arrayExpr, elemCond)
|
||||
return renderResult{sql: fmt.Sprintf("(%s AND %s)", nonEmpty, sub)}, nil
|
||||
case DialectMySQL:
|
||||
nonEmpty := fmt.Sprintf("%s IS NOT NULL AND JSON_LENGTH(%s) > 0", arrayExpr, arrayExpr)
|
||||
sub := fmt.Sprintf("NOT EXISTS (SELECT 1 FROM JSON_TABLE(%s, '$[*]' COLUMNS (value VARCHAR(512) PATH '$')) AS elem WHERE NOT (%s))", arrayExpr, elemCond)
|
||||
return renderResult{sql: fmt.Sprintf("(%s AND %s)", nonEmpty, sub)}, nil
|
||||
case DialectPostgres:
|
||||
nonEmpty := fmt.Sprintf("%s IS NOT NULL AND jsonb_array_length(%s) > 0", arrayExpr, arrayExpr)
|
||||
sub := fmt.Sprintf("NOT EXISTS (SELECT 1 FROM jsonb_array_elements_text(%s) AS elem(value) WHERE NOT (%s))", arrayExpr, elemCond)
|
||||
return renderResult{sql: fmt.Sprintf("(%s AND %s)", nonEmpty, sub)}, nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported dialect %s", r.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
// renderTagExistsOne renders tags.exists_one(t, <pred>): exactly one element
|
||||
// satisfies the predicate, via a COUNT(...) = 1 subquery. A null or empty array
|
||||
// yields COUNT 0, which is correctly not equal to 1.
|
||||
func (r *renderer) renderTagExistsOne(field Field, pred PredicateExpr) (renderResult, error) {
|
||||
arrayExpr := jsonArrayExpr(r.dialect, field)
|
||||
elemCond, err := r.elementPredicateSQL(pred)
|
||||
if err != nil {
|
||||
return renderResult{}, err
|
||||
}
|
||||
switch r.dialect {
|
||||
case DialectSQLite:
|
||||
return renderResult{sql: fmt.Sprintf("(SELECT COUNT(*) FROM json_each(%s) WHERE %s) = 1", arrayExpr, elemCond)}, nil
|
||||
case DialectMySQL:
|
||||
return renderResult{sql: fmt.Sprintf("(SELECT COUNT(*) FROM JSON_TABLE(%s, '$[*]' COLUMNS (value VARCHAR(512) PATH '$')) AS elem WHERE %s) = 1", arrayExpr, elemCond)}, nil
|
||||
case DialectPostgres:
|
||||
return renderResult{sql: fmt.Sprintf("(SELECT COUNT(*) FROM jsonb_array_elements_text(%s) AS elem(value) WHERE %s) = 1", arrayExpr, elemCond)}, nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported dialect %s", r.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
// elementPredicateSQL builds the per-element SQL condition for an all() predicate.
|
||||
// The iterated element is exposed as the unqualified column `value` on all dialects
|
||||
// (json_each.value / JSON_TABLE column / elem(value)).
|
||||
func (r *renderer) elementPredicateSQL(pred PredicateExpr) (string, error) {
|
||||
switch p := pred.(type) {
|
||||
case *EqualsPredicate:
|
||||
return fmt.Sprintf("value = %s", r.addArg(p.Value)), nil
|
||||
case *StartsWithPredicate:
|
||||
return r.foldedLike("value", likePattern(TextMatchPrefix, p.Prefix)), nil
|
||||
case *EndsWithPredicate:
|
||||
return r.foldedLike("value", likePattern(TextMatchSuffix, p.Suffix)), nil
|
||||
case *ContainsPredicate:
|
||||
return r.foldedLike("value", likePattern(TextMatchContains, p.Substring)), nil
|
||||
default:
|
||||
return "", errors.Errorf("unsupported predicate %T in all()", pred)
|
||||
}
|
||||
}
|
||||
|
||||
// renderTagEquals generates SQL for tags.exists(t, t == "value").
|
||||
func (r *renderer) renderTagEquals(field Field, value string, _ ComprehensionKind) (renderResult, error) {
|
||||
arrayExpr := jsonArrayExpr(r.dialect, field)
|
||||
|
||||
switch r.dialect {
|
||||
case DialectSQLite, DialectMySQL:
|
||||
exactMatch := r.buildJSONArrayLike(arrayExpr, fmt.Sprintf(`%%"%s"%%`, value))
|
||||
return renderResult{sql: r.wrapWithNullCheck(arrayExpr, exactMatch)}, nil
|
||||
case DialectPostgres:
|
||||
exactMatch := fmt.Sprintf("%s @> jsonb_build_array(%s::json)", arrayExpr, r.addArg(fmt.Sprintf(`"%s"`, value)))
|
||||
return renderResult{sql: r.wrapWithNullCheck(arrayExpr, exactMatch)}, nil
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported dialect %s", r.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
// renderTagStartsWith generates SQL for tags.exists(t, t.startsWith("prefix")).
|
||||
func (r *renderer) renderTagStartsWith(field Field, prefix string, _ ComprehensionKind) (renderResult, error) {
|
||||
arrayExpr := jsonArrayExpr(r.dialect, field)
|
||||
|
||||
switch r.dialect {
|
||||
case DialectSQLite, DialectMySQL:
|
||||
// Match exact tag or tags with this prefix (hierarchical support)
|
||||
exactMatch := r.buildJSONArrayLike(arrayExpr, fmt.Sprintf(`%%"%s"%%`, prefix))
|
||||
prefixMatch := r.buildJSONArrayLike(arrayExpr, fmt.Sprintf(`%%"%s%%`, prefix))
|
||||
condition := fmt.Sprintf("(%s OR %s)", exactMatch, prefixMatch)
|
||||
return renderResult{sql: r.wrapWithNullCheck(arrayExpr, condition)}, nil
|
||||
|
||||
case DialectPostgres:
|
||||
// Use PostgreSQL's powerful JSON operators
|
||||
exactMatch := fmt.Sprintf("%s @> jsonb_build_array(%s::json)", arrayExpr, r.addArg(fmt.Sprintf(`"%s"`, prefix)))
|
||||
prefixMatch := fmt.Sprintf("(%s)::text LIKE %s", arrayExpr, r.addArg(fmt.Sprintf(`%%"%s%%`, prefix)))
|
||||
condition := fmt.Sprintf("(%s OR %s)", exactMatch, prefixMatch)
|
||||
return renderResult{sql: r.wrapWithNullCheck(arrayExpr, condition)}, nil
|
||||
|
||||
default:
|
||||
return renderResult{}, errors.Errorf("unsupported dialect %s", r.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
// renderTagEndsWith generates SQL for tags.exists(t, t.endsWith("suffix")).
|
||||
func (r *renderer) renderTagEndsWith(field Field, suffix string, _ ComprehensionKind) (renderResult, error) {
|
||||
arrayExpr := jsonArrayExpr(r.dialect, field)
|
||||
pattern := fmt.Sprintf(`%%%s"%%`, suffix)
|
||||
|
||||
likeExpr := r.buildJSONArrayLike(arrayExpr, pattern)
|
||||
return renderResult{sql: r.wrapWithNullCheck(arrayExpr, likeExpr)}, nil
|
||||
}
|
||||
|
||||
// renderTagContains generates SQL for tags.exists(t, t.contains("substring")).
|
||||
func (r *renderer) renderTagContains(field Field, substring string, _ ComprehensionKind) (renderResult, error) {
|
||||
arrayExpr := jsonArrayExpr(r.dialect, field)
|
||||
pattern := fmt.Sprintf(`%%%s%%`, substring)
|
||||
|
||||
likeExpr := r.buildJSONArrayLike(arrayExpr, pattern)
|
||||
return renderResult{sql: r.wrapWithNullCheck(arrayExpr, likeExpr)}, nil
|
||||
}
|
||||
|
||||
// buildJSONArrayLike builds a LIKE expression for matching within a JSON array.
|
||||
// Returns the LIKE clause without NULL/empty checks.
|
||||
func (r *renderer) buildJSONArrayLike(arrayExpr, pattern string) string {
|
||||
switch r.dialect {
|
||||
case DialectSQLite, DialectMySQL:
|
||||
return fmt.Sprintf("%s LIKE %s", arrayExpr, r.addArg(pattern))
|
||||
case DialectPostgres:
|
||||
return fmt.Sprintf("(%s)::text LIKE %s", arrayExpr, r.addArg(pattern))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// wrapWithNullCheck wraps a condition with NULL and empty array checks.
|
||||
// This ensures we don't match against NULL or empty JSON arrays.
|
||||
func (r *renderer) wrapWithNullCheck(arrayExpr, condition string) string {
|
||||
var nullCheck string
|
||||
switch r.dialect {
|
||||
case DialectSQLite:
|
||||
nullCheck = fmt.Sprintf("%s IS NOT NULL AND %s != '[]'", arrayExpr, arrayExpr)
|
||||
case DialectMySQL:
|
||||
nullCheck = fmt.Sprintf("%s IS NOT NULL AND JSON_LENGTH(%s) > 0", arrayExpr, arrayExpr)
|
||||
case DialectPostgres:
|
||||
nullCheck = fmt.Sprintf("%s IS NOT NULL AND jsonb_array_length(%s) > 0", arrayExpr, arrayExpr)
|
||||
default:
|
||||
return condition
|
||||
}
|
||||
return fmt.Sprintf("(%s AND %s)", condition, nullCheck)
|
||||
}
|
||||
|
||||
func (r *renderer) jsonBoolPredicate(field Field) (string, error) {
|
||||
expr := jsonExtractExpr(r.dialect, field)
|
||||
switch r.dialect {
|
||||
case DialectSQLite:
|
||||
return fmt.Sprintf("%s IS TRUE", expr), nil
|
||||
case DialectMySQL:
|
||||
return fmt.Sprintf("COALESCE(%s, CAST('false' AS JSON)) = CAST('true' AS JSON)", expr), nil
|
||||
case DialectPostgres:
|
||||
return fmt.Sprintf("(%s)::boolean IS TRUE", expr), nil
|
||||
default:
|
||||
return "", errors.Errorf("unsupported dialect %s", r.dialect)
|
||||
}
|
||||
}
|
||||
|
||||
func combineAnd(left, right renderResult) renderResult {
|
||||
if left.unsatisfiable || right.unsatisfiable {
|
||||
return renderResult{sql: "1 = 0", unsatisfiable: true}
|
||||
}
|
||||
if left.trivial {
|
||||
return right
|
||||
}
|
||||
if right.trivial {
|
||||
return left
|
||||
}
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("(%s AND %s)", left.sql, right.sql),
|
||||
}
|
||||
}
|
||||
|
||||
func combineOr(left, right renderResult) renderResult {
|
||||
if left.trivial || right.trivial {
|
||||
return renderResult{trivial: true}
|
||||
}
|
||||
if left.unsatisfiable {
|
||||
return right
|
||||
}
|
||||
if right.unsatisfiable {
|
||||
return left
|
||||
}
|
||||
return renderResult{
|
||||
sql: fmt.Sprintf("(%s OR %s)", left.sql, right.sql),
|
||||
}
|
||||
}
|
||||
|
||||
func (r *renderer) addArg(value any) string {
|
||||
r.placeholderCounter++
|
||||
r.args = append(r.args, value)
|
||||
if r.dialect == DialectPostgres {
|
||||
return fmt.Sprintf("$%d", r.placeholderOffset+r.placeholderCounter)
|
||||
}
|
||||
return "?"
|
||||
}
|
||||
|
||||
func (r *renderer) addBoolArg(value bool) string {
|
||||
var v any
|
||||
switch r.dialect {
|
||||
case DialectSQLite:
|
||||
if value {
|
||||
v = 1
|
||||
} else {
|
||||
v = 0
|
||||
}
|
||||
default:
|
||||
v = value
|
||||
}
|
||||
return r.addArg(v)
|
||||
}
|
||||
|
||||
func expectLiteral(expr ValueExpr) (any, error) {
|
||||
lit, ok := expr.(*LiteralValue)
|
||||
if !ok {
|
||||
return nil, errors.New("expression must be a literal")
|
||||
}
|
||||
return lit.Value, nil
|
||||
}
|
||||
|
||||
func expectBool(expr ValueExpr) (bool, error) {
|
||||
lit, err := expectLiteral(expr)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
value, ok := lit.(bool)
|
||||
if !ok {
|
||||
return false, errors.New("boolean literal required")
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func expectNumericLiteral(expr ValueExpr) (int64, error) {
|
||||
lit, err := expectLiteral(expr)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return toInt64(lit)
|
||||
}
|
||||
|
||||
func toInt64(value any) (int64, error) {
|
||||
switch v := value.(type) {
|
||||
case int:
|
||||
return int64(v), nil
|
||||
case int32:
|
||||
return int64(v), nil
|
||||
case int64:
|
||||
return v, nil
|
||||
case uint32:
|
||||
return int64(v), nil
|
||||
case uint64:
|
||||
return int64(v), nil
|
||||
case float32:
|
||||
return int64(v), nil
|
||||
case float64:
|
||||
return int64(v), nil
|
||||
default:
|
||||
return 0, errors.Errorf("cannot convert %T to int64", value)
|
||||
}
|
||||
}
|
||||
|
||||
func sqlOperator(op ComparisonOperator) string {
|
||||
return string(op)
|
||||
}
|
||||
|
||||
func qualifyColumn(d DialectName, col Column) string {
|
||||
switch d {
|
||||
case DialectPostgres:
|
||||
return fmt.Sprintf("%s.%s", col.Table, col.Name)
|
||||
default:
|
||||
return fmt.Sprintf("`%s`.`%s`", col.Table, col.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func jsonPath(field Field) string {
|
||||
return "$." + strings.Join(field.JSONPath, ".")
|
||||
}
|
||||
|
||||
func jsonExtractExpr(d DialectName, field Field) string {
|
||||
column := qualifyColumn(d, field.Column)
|
||||
switch d {
|
||||
case DialectSQLite, DialectMySQL:
|
||||
return fmt.Sprintf("JSON_EXTRACT(%s, '%s')", column, jsonPath(field))
|
||||
case DialectPostgres:
|
||||
return buildPostgresJSONAccessor(column, field.JSONPath, true)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func jsonArrayExpr(d DialectName, field Field) string {
|
||||
column := qualifyColumn(d, field.Column)
|
||||
switch d {
|
||||
case DialectSQLite, DialectMySQL:
|
||||
return fmt.Sprintf("JSON_EXTRACT(%s, '%s')", column, jsonPath(field))
|
||||
case DialectPostgres:
|
||||
return buildPostgresJSONAccessor(column, field.JSONPath, false)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func jsonArrayLengthExpr(d DialectName, field Field) string {
|
||||
arrayExpr := jsonArrayExpr(d, field)
|
||||
switch d {
|
||||
case DialectSQLite:
|
||||
return fmt.Sprintf("JSON_ARRAY_LENGTH(COALESCE(%s, JSON_ARRAY()))", arrayExpr)
|
||||
case DialectMySQL:
|
||||
return fmt.Sprintf("JSON_LENGTH(COALESCE(%s, JSON_ARRAY()))", arrayExpr)
|
||||
case DialectPostgres:
|
||||
return fmt.Sprintf("jsonb_array_length(COALESCE(%s, '[]'::jsonb))", arrayExpr)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func buildPostgresJSONAccessor(base string, path []string, terminalText bool) string {
|
||||
expr := base
|
||||
for idx, part := range path {
|
||||
if idx == len(path)-1 && terminalText {
|
||||
expr = fmt.Sprintf("%s->>'%s'", expr, part)
|
||||
} else {
|
||||
expr = fmt.Sprintf("%s->'%s'", expr, part)
|
||||
}
|
||||
}
|
||||
return expr
|
||||
}
|
||||
@@ -0,0 +1,327 @@
|
||||
package filter
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/google/cel-go/cel"
|
||||
"github.com/google/cel-go/ext"
|
||||
)
|
||||
|
||||
// DialectName enumerates supported SQL dialects.
|
||||
type DialectName string
|
||||
|
||||
const (
|
||||
DialectSQLite DialectName = "sqlite"
|
||||
DialectMySQL DialectName = "mysql"
|
||||
DialectPostgres DialectName = "postgres"
|
||||
)
|
||||
|
||||
// FieldType represents the logical type of a field.
|
||||
type FieldType string
|
||||
|
||||
const (
|
||||
FieldTypeString FieldType = "string"
|
||||
FieldTypeInt FieldType = "int"
|
||||
FieldTypeBool FieldType = "bool"
|
||||
FieldTypeTimestamp FieldType = "timestamp"
|
||||
)
|
||||
|
||||
// FieldKind describes how a field is stored.
|
||||
type FieldKind string
|
||||
|
||||
const (
|
||||
FieldKindScalar FieldKind = "scalar"
|
||||
FieldKindBoolColumn FieldKind = "bool_column"
|
||||
FieldKindJSONBool FieldKind = "json_bool"
|
||||
FieldKindJSONList FieldKind = "json_list"
|
||||
FieldKindVirtualAlias FieldKind = "virtual_alias"
|
||||
)
|
||||
|
||||
// Column identifies the backing table column.
|
||||
type Column struct {
|
||||
Table string
|
||||
Name string
|
||||
}
|
||||
|
||||
// Field captures the schema metadata for an exposed CEL identifier.
|
||||
type Field struct {
|
||||
Name string
|
||||
Kind FieldKind
|
||||
Type FieldType
|
||||
Column Column
|
||||
JSONPath []string
|
||||
AliasFor string
|
||||
SupportsContains bool
|
||||
Expressions map[DialectName]string
|
||||
AllowedComparisonOps map[ComparisonOperator]bool
|
||||
}
|
||||
|
||||
// Schema collects CEL environment options and field metadata.
|
||||
type Schema struct {
|
||||
Name string
|
||||
Fields map[string]Field
|
||||
EnvOptions []cel.EnvOption
|
||||
}
|
||||
|
||||
// Field returns the field metadata if present.
|
||||
func (s Schema) Field(name string) (Field, bool) {
|
||||
f, ok := s.Fields[name]
|
||||
return f, ok
|
||||
}
|
||||
|
||||
// ResolveAlias resolves a virtual alias to its target field.
|
||||
func (s Schema) ResolveAlias(name string) (Field, bool) {
|
||||
field, ok := s.Fields[name]
|
||||
if !ok {
|
||||
return Field{}, false
|
||||
}
|
||||
if field.Kind == FieldKindVirtualAlias {
|
||||
target, ok := s.Fields[field.AliasFor]
|
||||
if !ok {
|
||||
return Field{}, false
|
||||
}
|
||||
return target, true
|
||||
}
|
||||
return field, true
|
||||
}
|
||||
|
||||
// NewSchema constructs the memo filter schema and CEL environment.
|
||||
func NewSchema() Schema {
|
||||
fields := map[string]Field{
|
||||
"content": {
|
||||
Name: "content",
|
||||
Kind: FieldKindScalar,
|
||||
Type: FieldTypeString,
|
||||
Column: Column{Table: "memo", Name: "content"},
|
||||
SupportsContains: true,
|
||||
Expressions: map[DialectName]string{},
|
||||
},
|
||||
"creator": {
|
||||
Name: "creator",
|
||||
Kind: FieldKindScalar,
|
||||
Type: FieldTypeString,
|
||||
Column: Column{Table: "memo_creator", Name: "username"},
|
||||
Expressions: map[DialectName]string{
|
||||
DialectSQLite: "('users/' || %s)",
|
||||
DialectMySQL: "CONCAT('users/', %s)",
|
||||
DialectPostgres: "('users/' || %s)",
|
||||
},
|
||||
AllowedComparisonOps: map[ComparisonOperator]bool{
|
||||
CompareEq: true,
|
||||
CompareNeq: true,
|
||||
},
|
||||
},
|
||||
"creator_id": {
|
||||
Name: "creator_id",
|
||||
Kind: FieldKindScalar,
|
||||
Type: FieldTypeInt,
|
||||
Column: Column{Table: "memo", Name: "creator_id"},
|
||||
Expressions: map[DialectName]string{},
|
||||
AllowedComparisonOps: map[ComparisonOperator]bool{
|
||||
CompareEq: true,
|
||||
CompareNeq: true,
|
||||
},
|
||||
},
|
||||
"created_ts": {
|
||||
Name: "created_ts",
|
||||
Kind: FieldKindScalar,
|
||||
Type: FieldTypeTimestamp,
|
||||
Column: Column{Table: "memo", Name: "created_ts"},
|
||||
Expressions: map[DialectName]string{
|
||||
// MySQL stores created_ts as TIMESTAMP, needs conversion to epoch
|
||||
DialectMySQL: "UNIX_TIMESTAMP(%s)",
|
||||
// PostgreSQL and SQLite store created_ts as BIGINT (epoch), no conversion needed
|
||||
DialectPostgres: "%s",
|
||||
DialectSQLite: "%s",
|
||||
},
|
||||
},
|
||||
"updated_ts": {
|
||||
Name: "updated_ts",
|
||||
Kind: FieldKindScalar,
|
||||
Type: FieldTypeTimestamp,
|
||||
Column: Column{Table: "memo", Name: "updated_ts"},
|
||||
Expressions: map[DialectName]string{
|
||||
// MySQL stores updated_ts as TIMESTAMP, needs conversion to epoch
|
||||
DialectMySQL: "UNIX_TIMESTAMP(%s)",
|
||||
// PostgreSQL and SQLite store updated_ts as BIGINT (epoch), no conversion needed
|
||||
DialectPostgres: "%s",
|
||||
DialectSQLite: "%s",
|
||||
},
|
||||
},
|
||||
"pinned": {
|
||||
Name: "pinned",
|
||||
Kind: FieldKindBoolColumn,
|
||||
Type: FieldTypeBool,
|
||||
Column: Column{Table: "memo", Name: "pinned"},
|
||||
Expressions: map[DialectName]string{},
|
||||
AllowedComparisonOps: map[ComparisonOperator]bool{
|
||||
CompareEq: true,
|
||||
CompareNeq: true,
|
||||
},
|
||||
},
|
||||
"visibility": {
|
||||
Name: "visibility",
|
||||
Kind: FieldKindScalar,
|
||||
Type: FieldTypeString,
|
||||
Column: Column{Table: "memo", Name: "visibility"},
|
||||
Expressions: map[DialectName]string{},
|
||||
AllowedComparisonOps: map[ComparisonOperator]bool{
|
||||
CompareEq: true,
|
||||
CompareNeq: true,
|
||||
},
|
||||
},
|
||||
"tags": {
|
||||
Name: "tags",
|
||||
Kind: FieldKindJSONList,
|
||||
Type: FieldTypeString,
|
||||
Column: Column{Table: "memo", Name: "payload"},
|
||||
JSONPath: []string{"tags"},
|
||||
},
|
||||
"tag": {
|
||||
Name: "tag",
|
||||
Kind: FieldKindVirtualAlias,
|
||||
Type: FieldTypeString,
|
||||
AliasFor: "tags",
|
||||
},
|
||||
"has_task_list": {
|
||||
Name: "has_task_list",
|
||||
Kind: FieldKindJSONBool,
|
||||
Type: FieldTypeBool,
|
||||
Column: Column{Table: "memo", Name: "payload"},
|
||||
JSONPath: []string{"property", "hasTaskList"},
|
||||
AllowedComparisonOps: map[ComparisonOperator]bool{
|
||||
CompareEq: true,
|
||||
CompareNeq: true,
|
||||
},
|
||||
},
|
||||
"has_link": {
|
||||
Name: "has_link",
|
||||
Kind: FieldKindJSONBool,
|
||||
Type: FieldTypeBool,
|
||||
Column: Column{Table: "memo", Name: "payload"},
|
||||
JSONPath: []string{"property", "hasLink"},
|
||||
AllowedComparisonOps: map[ComparisonOperator]bool{
|
||||
CompareEq: true,
|
||||
CompareNeq: true,
|
||||
},
|
||||
},
|
||||
"has_code": {
|
||||
Name: "has_code",
|
||||
Kind: FieldKindJSONBool,
|
||||
Type: FieldTypeBool,
|
||||
Column: Column{Table: "memo", Name: "payload"},
|
||||
JSONPath: []string{"property", "hasCode"},
|
||||
AllowedComparisonOps: map[ComparisonOperator]bool{
|
||||
CompareEq: true,
|
||||
CompareNeq: true,
|
||||
},
|
||||
},
|
||||
"has_incomplete_tasks": {
|
||||
Name: "has_incomplete_tasks",
|
||||
Kind: FieldKindJSONBool,
|
||||
Type: FieldTypeBool,
|
||||
Column: Column{Table: "memo", Name: "payload"},
|
||||
JSONPath: []string{"property", "hasIncompleteTasks"},
|
||||
AllowedComparisonOps: map[ComparisonOperator]bool{
|
||||
CompareEq: true,
|
||||
CompareNeq: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
envOptions := []cel.EnvOption{
|
||||
cel.Variable("content", cel.StringType),
|
||||
cel.Variable("creator", cel.StringType),
|
||||
cel.Variable("creator_id", cel.IntType),
|
||||
cel.Variable("created_ts", cel.TimestampType),
|
||||
cel.Variable("updated_ts", cel.TimestampType),
|
||||
cel.Variable("pinned", cel.BoolType),
|
||||
cel.Variable("tag", cel.StringType),
|
||||
cel.Variable("tags", cel.ListType(cel.StringType)),
|
||||
cel.Variable("visibility", cel.StringType),
|
||||
cel.Variable("has_task_list", cel.BoolType),
|
||||
cel.Variable("has_link", cel.BoolType),
|
||||
cel.Variable("has_code", cel.BoolType),
|
||||
cel.Variable("has_incomplete_tasks", cel.BoolType),
|
||||
cel.Variable("now", cel.TimestampType),
|
||||
ext.Sets(),
|
||||
cel.ASTValidators(cel.ValidateRegexLiterals()),
|
||||
}
|
||||
|
||||
return Schema{
|
||||
Name: "memo",
|
||||
Fields: fields,
|
||||
EnvOptions: envOptions,
|
||||
}
|
||||
}
|
||||
|
||||
// NewAttachmentSchema constructs the attachment filter schema and CEL environment.
|
||||
func NewAttachmentSchema() Schema {
|
||||
fields := map[string]Field{
|
||||
"filename": {
|
||||
Name: "filename",
|
||||
Kind: FieldKindScalar,
|
||||
Type: FieldTypeString,
|
||||
Column: Column{Table: "attachment", Name: "filename"},
|
||||
SupportsContains: true,
|
||||
Expressions: map[DialectName]string{},
|
||||
},
|
||||
"mime_type": {
|
||||
Name: "mime_type",
|
||||
Kind: FieldKindScalar,
|
||||
Type: FieldTypeString,
|
||||
Column: Column{Table: "attachment", Name: "type"},
|
||||
SupportsContains: true,
|
||||
Expressions: map[DialectName]string{},
|
||||
},
|
||||
"create_time": {
|
||||
Name: "create_time",
|
||||
Kind: FieldKindScalar,
|
||||
Type: FieldTypeTimestamp,
|
||||
Column: Column{Table: "attachment", Name: "created_ts"},
|
||||
Expressions: map[DialectName]string{
|
||||
// MySQL stores created_ts as TIMESTAMP, needs conversion to epoch
|
||||
DialectMySQL: "UNIX_TIMESTAMP(%s)",
|
||||
// PostgreSQL and SQLite store created_ts as BIGINT (epoch), no conversion needed
|
||||
DialectPostgres: "%s",
|
||||
DialectSQLite: "%s",
|
||||
},
|
||||
},
|
||||
"memo_id": {
|
||||
Name: "memo_id",
|
||||
Kind: FieldKindScalar,
|
||||
Type: FieldTypeInt,
|
||||
Column: Column{Table: "attachment", Name: "memo_id"},
|
||||
Expressions: map[DialectName]string{},
|
||||
AllowedComparisonOps: map[ComparisonOperator]bool{
|
||||
CompareEq: true,
|
||||
CompareNeq: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
envOptions := []cel.EnvOption{
|
||||
cel.Variable("filename", cel.StringType),
|
||||
cel.Variable("mime_type", cel.StringType),
|
||||
cel.Variable("create_time", cel.TimestampType),
|
||||
cel.Variable("memo_id", cel.AnyType),
|
||||
cel.Variable("now", cel.TimestampType),
|
||||
cel.ASTValidators(cel.ValidateRegexLiterals()),
|
||||
}
|
||||
|
||||
return Schema{
|
||||
Name: "attachment",
|
||||
Fields: fields,
|
||||
EnvOptions: envOptions,
|
||||
}
|
||||
}
|
||||
|
||||
// columnExpr returns the field expression for the given dialect, applying
|
||||
// any schema-specific overrides (e.g. UNIX timestamp conversions).
|
||||
func (f Field) columnExpr(d DialectName) string {
|
||||
base := qualifyColumn(d, f.Column)
|
||||
if expr, ok := f.Expressions[d]; ok && expr != "" {
|
||||
return fmt.Sprintf(expr, base)
|
||||
}
|
||||
return base
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
package filter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// fixedClock returns a deterministic clock for asserting folded `now` values.
|
||||
func fixedClock(epoch int64) func() time.Time {
|
||||
return func() time.Time { return time.Unix(epoch, 0) }
|
||||
}
|
||||
|
||||
func memoEngineAt(t *testing.T, epoch int64) *Engine {
|
||||
t.Helper()
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
engine.nowFunc = fixedClock(epoch)
|
||||
return engine
|
||||
}
|
||||
|
||||
func TestCompileNowVariableFoldsToInjectedClock(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine := memoEngineAt(t, 1750000000)
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `created_ts >= now`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []any{int64(1750000000)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileNowFunctionIsRemoved(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewSchema())
|
||||
require.NoError(t, err)
|
||||
|
||||
// now() was the legacy custom function; it is replaced by the `now` variable.
|
||||
_, err = engine.Compile(context.Background(), `created_ts >= now()`)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestCompileNowMinusDurationFolds(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine := memoEngineAt(t, 1750000000)
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `created_ts >= now - duration("1h")`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []any{int64(1750000000 - 3600)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileNowPlusDurationFolds(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine := memoEngineAt(t, 1750000000)
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `updated_ts < now + duration("24h")`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []any{int64(1750000000 + 86400)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileAbsoluteTimestampStringFolds(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine := memoEngineAt(t, 1750000000)
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `created_ts >= timestamp("2025-01-01T00:00:00Z")`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []any{int64(1735689600)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileTimestampFromEpochIntFolds(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// This is the shape the frontend date-range filter emits.
|
||||
engine := memoEngineAt(t, 1750000000)
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `created_ts >= timestamp(1730000000)`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []any{int64(1730000000)}, stmt.Args)
|
||||
}
|
||||
|
||||
func TestCompileInvalidDurationLiteralErrors(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine := memoEngineAt(t, 1750000000)
|
||||
_, err := engine.Compile(context.Background(), `created_ts >= now - duration("garbage")`)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "duration")
|
||||
}
|
||||
|
||||
func TestCompileAttachmentCreateTimeUsesNow(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
engine, err := NewEngine(NewAttachmentSchema())
|
||||
require.NoError(t, err)
|
||||
engine.nowFunc = fixedClock(1750000000)
|
||||
|
||||
stmt, err := engine.CompileToStatement(context.Background(), `create_time >= now - duration("24h")`, RenderOptions{Dialect: DialectSQLite})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []any{int64(1750000000 - 86400)}, stmt.Args)
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
package httpgetter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
"golang.org/x/net/html"
|
||||
"golang.org/x/net/html/atom"
|
||||
)
|
||||
|
||||
var ErrInternalIP = errors.New("internal IP addresses are not allowed")
|
||||
|
||||
const maxHTMLMetaBytes = 512 * 1024
|
||||
|
||||
var (
|
||||
lookupIPAddr = net.DefaultResolver.LookupIPAddr
|
||||
dialContext = (&net.Dialer{
|
||||
Timeout: 30 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
}).DialContext
|
||||
httpClient = newHTTPClient()
|
||||
)
|
||||
|
||||
func newHTTPClient() *http.Client {
|
||||
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||
transport.Proxy = nil
|
||||
transport.DialContext = secureDialContext
|
||||
|
||||
return &http.Client{
|
||||
Transport: transport,
|
||||
Timeout: 5 * time.Second,
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if err := validateURL(req.URL.String()); err != nil {
|
||||
return errors.Wrap(err, "redirect to internal IP")
|
||||
}
|
||||
if len(via) >= 10 {
|
||||
return errors.New("too many redirects")
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func secureDialContext(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "invalid address")
|
||||
}
|
||||
|
||||
ips, err := resolveAllowedIPs(ctx, host)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var dialErr error
|
||||
for _, ip := range ips {
|
||||
conn, err := dialContext(ctx, network, net.JoinHostPort(ip.String(), port))
|
||||
if err == nil {
|
||||
return conn, nil
|
||||
}
|
||||
dialErr = err
|
||||
}
|
||||
return nil, dialErr
|
||||
}
|
||||
|
||||
func resolveAllowedIPs(ctx context.Context, host string) ([]net.IP, error) {
|
||||
if ip := net.ParseIP(host); ip != nil {
|
||||
if isInternalIP(ip) {
|
||||
return nil, errors.Wrap(ErrInternalIP, ip.String())
|
||||
}
|
||||
return []net.IP{ip}, nil
|
||||
}
|
||||
|
||||
addrs, err := lookupIPAddr(ctx, host)
|
||||
if err != nil {
|
||||
return nil, errors.Errorf("failed to resolve hostname: %v", err)
|
||||
}
|
||||
|
||||
ips := make([]net.IP, 0, len(addrs))
|
||||
for _, addr := range addrs {
|
||||
ip := addr.IP
|
||||
if ip == nil {
|
||||
continue
|
||||
}
|
||||
if isInternalIP(ip) {
|
||||
return nil, errors.Wrapf(ErrInternalIP, "host=%s, ip=%s", host, ip.String())
|
||||
}
|
||||
ips = append(ips, ip)
|
||||
}
|
||||
if len(ips) == 0 {
|
||||
return nil, errors.New("hostname resolved to no addresses")
|
||||
}
|
||||
|
||||
return ips, nil
|
||||
}
|
||||
|
||||
func isInternalIP(ip net.IP) bool {
|
||||
return ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsUnspecified()
|
||||
}
|
||||
|
||||
func validateURL(urlStr string) error {
|
||||
u, err := url.Parse(urlStr)
|
||||
if err != nil {
|
||||
return errors.New("invalid URL format")
|
||||
}
|
||||
|
||||
if u.Scheme != "http" && u.Scheme != "https" {
|
||||
return errors.New("only http/https protocols are allowed")
|
||||
}
|
||||
|
||||
host := u.Hostname()
|
||||
if host == "" {
|
||||
return errors.New("empty hostname")
|
||||
}
|
||||
|
||||
if ip := net.ParseIP(host); ip != nil && isInternalIP(ip) {
|
||||
return errors.Wrap(ErrInternalIP, ip.String())
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type HTMLMeta struct {
|
||||
Title string `json:"title"`
|
||||
Description string `json:"description"`
|
||||
Image string `json:"image"`
|
||||
}
|
||||
|
||||
func GetHTMLMeta(urlStr string) (*HTMLMeta, error) {
|
||||
if err := validateURL(urlStr); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
response, err := httpClient.Get(urlStr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
|
||||
mediatype, err := getMediatype(response)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if mediatype != "text/html" {
|
||||
return nil, errors.New("not a HTML page")
|
||||
}
|
||||
|
||||
htmlMeta := extractHTMLMeta(io.LimitReader(response.Body, maxHTMLMetaBytes))
|
||||
enrichSiteMeta(response.Request.URL, htmlMeta)
|
||||
return htmlMeta, nil
|
||||
}
|
||||
|
||||
func extractHTMLMeta(resp io.Reader) *HTMLMeta {
|
||||
tokenizer := html.NewTokenizer(resp)
|
||||
htmlMeta := new(HTMLMeta)
|
||||
|
||||
for {
|
||||
tokenType := tokenizer.Next()
|
||||
if tokenType == html.ErrorToken {
|
||||
break
|
||||
} else if tokenType == html.StartTagToken || tokenType == html.SelfClosingTagToken {
|
||||
token := tokenizer.Token()
|
||||
if token.DataAtom == atom.Body {
|
||||
break
|
||||
}
|
||||
|
||||
if token.DataAtom == atom.Title {
|
||||
tokenizer.Next()
|
||||
token := tokenizer.Token()
|
||||
htmlMeta.Title = token.Data
|
||||
} else if token.DataAtom == atom.Meta {
|
||||
ogTitle, ok := extractMetaProperty(token, "og:title")
|
||||
if ok {
|
||||
htmlMeta.Title = ogTitle
|
||||
}
|
||||
|
||||
ogDescription, ok := extractMetaProperty(token, "og:description")
|
||||
if ok {
|
||||
htmlMeta.Description = ogDescription
|
||||
}
|
||||
|
||||
ogImage, ok := extractMetaProperty(token, "og:image")
|
||||
if ok {
|
||||
htmlMeta.Image = ogImage
|
||||
}
|
||||
|
||||
description, ok := extractMetaProperty(token, "description")
|
||||
if ok && htmlMeta.Description == "" {
|
||||
htmlMeta.Description = description
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return htmlMeta
|
||||
}
|
||||
|
||||
func extractMetaProperty(token html.Token, prop string) (content string, ok bool) {
|
||||
content, ok = "", false
|
||||
for _, attr := range token.Attr {
|
||||
if (attr.Key == "property" || attr.Key == "name") && strings.EqualFold(attr.Val, prop) {
|
||||
ok = true
|
||||
}
|
||||
if attr.Key == "content" {
|
||||
content = attr.Val
|
||||
}
|
||||
}
|
||||
return content, ok
|
||||
}
|
||||
|
||||
func enrichSiteMeta(url *url.URL, meta *HTMLMeta) {
|
||||
if url.Hostname() == "www.youtube.com" {
|
||||
if url.Path == "/watch" {
|
||||
vid := url.Query().Get("v")
|
||||
if vid != "" {
|
||||
meta.Image = fmt.Sprintf("https://img.youtube.com/vi/%s/mqdefault.jpg", vid)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package httpgetter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return f(req)
|
||||
}
|
||||
|
||||
func TestGetHTMLMeta(t *testing.T) {
|
||||
originalHTTPClient := httpClient
|
||||
t.Cleanup(func() {
|
||||
httpClient = originalHTTPClient
|
||||
})
|
||||
|
||||
httpClient = &http.Client{
|
||||
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
require.Equal(t, "http://93.184.216.34/article", req.URL.String())
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/html; charset=utf-8"}},
|
||||
Body: io.NopCloser(strings.NewReader(`<!doctype html>
|
||||
<html>
|
||||
<head>
|
||||
<title>Fallback title</title>
|
||||
<meta name="description" content="Fallback description">
|
||||
<meta property="og:title" content="Open Graph title">
|
||||
<meta property="og:description" content="Open Graph description">
|
||||
<meta property="og:image" content="https://example.com/cover.png">
|
||||
</head>
|
||||
<body>ignored</body>
|
||||
</html>`)),
|
||||
Request: req,
|
||||
}, nil
|
||||
}),
|
||||
}
|
||||
|
||||
metadata, err := GetHTMLMeta("http://93.184.216.34/article")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, HTMLMeta{
|
||||
Title: "Open Graph title",
|
||||
Description: "Open Graph description",
|
||||
Image: "https://example.com/cover.png",
|
||||
}, *metadata)
|
||||
}
|
||||
|
||||
func TestGetHTMLMetaWithNameOnly(t *testing.T) {
|
||||
originalHTTPClient := httpClient
|
||||
t.Cleanup(func() {
|
||||
httpClient = originalHTTPClient
|
||||
})
|
||||
|
||||
httpClient = &http.Client{
|
||||
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
require.Equal(t, "http://93.184.216.34/blog", req.URL.String())
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/html; charset=utf-8"}},
|
||||
Body: io.NopCloser(strings.NewReader(`<!doctype html>
|
||||
<html>
|
||||
<head>
|
||||
<title>Sample Page</title>
|
||||
<meta name="description" content="This description should appear in the link preview.">
|
||||
</head>
|
||||
<body>Hello</body>
|
||||
</html>`)),
|
||||
Request: req,
|
||||
}, nil
|
||||
}),
|
||||
}
|
||||
|
||||
metadata, err := GetHTMLMeta("http://93.184.216.34/blog")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, HTMLMeta{
|
||||
Title: "Sample Page",
|
||||
Description: "This description should appear in the link preview.",
|
||||
Image: "",
|
||||
}, *metadata)
|
||||
}
|
||||
|
||||
func TestGetHTMLMetaWithNameCaseInsensitive(t *testing.T) {
|
||||
originalHTTPClient := httpClient
|
||||
t.Cleanup(func() {
|
||||
httpClient = originalHTTPClient
|
||||
})
|
||||
|
||||
httpClient = &http.Client{
|
||||
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
require.Equal(t, "http://93.184.216.34/blog", req.URL.String())
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/html; charset=utf-8"}},
|
||||
Body: io.NopCloser(strings.NewReader(`<!doctype html>
|
||||
<html>
|
||||
<head>
|
||||
<title>Sample Page</title>
|
||||
<meta name="Description" content="Case insensitive description match.">
|
||||
</head>
|
||||
<body>Hello</body>
|
||||
</html>`)),
|
||||
Request: req,
|
||||
}, nil
|
||||
}),
|
||||
}
|
||||
|
||||
metadata, err := GetHTMLMeta("http://93.184.216.34/blog")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, HTMLMeta{
|
||||
Title: "Sample Page",
|
||||
Description: "Case insensitive description match.",
|
||||
Image: "",
|
||||
}, *metadata)
|
||||
}
|
||||
|
||||
func TestGetHTMLMetaForInternal(t *testing.T) {
|
||||
// test for internal IP
|
||||
if _, err := GetHTMLMeta("http://192.168.0.1"); !errors.Is(err, ErrInternalIP) {
|
||||
t.Errorf("Expected error for internal IP, got %v", err)
|
||||
}
|
||||
|
||||
// test for resolved internal IP
|
||||
if _, err := GetHTMLMeta("http://localhost"); !errors.Is(err, ErrInternalIP) {
|
||||
t.Errorf("Expected error for resolved internal IP, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPClientHasTimeout(t *testing.T) {
|
||||
require.NotZero(t, httpClient.Timeout)
|
||||
}
|
||||
|
||||
func TestSecureDialContextRejectsResolvedInternalIP(t *testing.T) {
|
||||
originalLookupIPAddr := lookupIPAddr
|
||||
originalDialContext := dialContext
|
||||
t.Cleanup(func() {
|
||||
lookupIPAddr = originalLookupIPAddr
|
||||
dialContext = originalDialContext
|
||||
})
|
||||
|
||||
lookupIPAddr = func(context.Context, string) ([]net.IPAddr, error) {
|
||||
return []net.IPAddr{{IP: net.ParseIP("127.0.0.1")}}, nil
|
||||
}
|
||||
dialContext = func(context.Context, string, string) (net.Conn, error) {
|
||||
t.Fatal("internal IP should be rejected before dialing")
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
_, err := secureDialContext(context.Background(), "tcp", "rebind.example:80")
|
||||
require.ErrorIs(t, err, ErrInternalIP)
|
||||
}
|
||||
|
||||
func TestSecureDialContextDialsResolvedIP(t *testing.T) {
|
||||
originalLookupIPAddr := lookupIPAddr
|
||||
originalDialContext := dialContext
|
||||
t.Cleanup(func() {
|
||||
lookupIPAddr = originalLookupIPAddr
|
||||
dialContext = originalDialContext
|
||||
})
|
||||
|
||||
lookupIPAddr = func(context.Context, string) ([]net.IPAddr, error) {
|
||||
return []net.IPAddr{{IP: net.ParseIP("93.184.216.34")}}, nil
|
||||
}
|
||||
|
||||
var dialedAddress string
|
||||
dialContext = func(_ context.Context, _ string, address string) (net.Conn, error) {
|
||||
dialedAddress = address
|
||||
clientConn, serverConn := net.Pipe()
|
||||
t.Cleanup(func() {
|
||||
clientConn.Close()
|
||||
serverConn.Close()
|
||||
})
|
||||
return clientConn, nil
|
||||
}
|
||||
|
||||
conn, err := secureDialContext(context.Background(), "tcp", "rebind.example:80")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, conn)
|
||||
require.Equal(t, "93.184.216.34:80", dialedAddress)
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
package httpgetter
|
||||
@@ -0,0 +1,45 @@
|
||||
package httpgetter
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type Image struct {
|
||||
Blob []byte
|
||||
Mediatype string
|
||||
}
|
||||
|
||||
func GetImage(urlStr string) (*Image, error) {
|
||||
if _, err := url.Parse(urlStr); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
response, err := http.Get(urlStr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
|
||||
mediatype, err := getMediatype(response)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !strings.HasPrefix(mediatype, "image/") {
|
||||
return nil, errors.New("wrong image mediatype")
|
||||
}
|
||||
|
||||
bodyBytes, err := io.ReadAll(response.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
image := &Image{
|
||||
Blob: bodyBytes,
|
||||
Mediatype: mediatype,
|
||||
}
|
||||
return image, nil
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package httpgetter
|
||||
|
||||
import (
|
||||
"mime"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
func getMediatype(response *http.Response) (string, error) {
|
||||
contentType := response.Header.Get("content-type")
|
||||
mediatype, _, err := mime.ParseMediaType(contentType)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return mediatype, nil
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
package idp
|
||||
|
||||
type IdentityProviderUserInfo struct {
|
||||
Identifier string
|
||||
DisplayName string
|
||||
Email string
|
||||
AvatarURL string
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
// Package oauth2 implements the OAuth2 identity provider integration.
|
||||
package oauth2
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
"golang.org/x/oauth2"
|
||||
|
||||
"github.com/usememos/memos/internal/idp"
|
||||
storepb "github.com/usememos/memos/proto/gen/store"
|
||||
)
|
||||
|
||||
// IdentityProvider represents an OAuth2 Identity Provider.
|
||||
type IdentityProvider struct {
|
||||
config *storepb.OAuth2Config
|
||||
}
|
||||
|
||||
const userInfoRequestTimeout = 10 * time.Second
|
||||
|
||||
// NewIdentityProvider initializes a new OAuth2 Identity Provider with the given configuration.
|
||||
func NewIdentityProvider(config *storepb.OAuth2Config) (*IdentityProvider, error) {
|
||||
for v, field := range map[string]string{
|
||||
config.ClientId: "clientId",
|
||||
config.ClientSecret: "clientSecret",
|
||||
config.TokenUrl: "tokenUrl",
|
||||
config.UserInfoUrl: "userInfoUrl",
|
||||
config.FieldMapping.Identifier: "fieldMapping.identifier",
|
||||
} {
|
||||
if v == "" {
|
||||
return nil, errors.Errorf(`the field "%s" is empty but required`, field)
|
||||
}
|
||||
}
|
||||
|
||||
return &IdentityProvider{
|
||||
config: config,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ExchangeToken returns the exchanged OAuth2 token using the given authorization code.
|
||||
// If codeVerifier is provided, it will be used for PKCE (Proof Key for Code Exchange) validation.
|
||||
func (p *IdentityProvider) ExchangeToken(ctx context.Context, redirectURL, code, codeVerifier string) (string, error) {
|
||||
conf := &oauth2.Config{
|
||||
ClientID: p.config.ClientId,
|
||||
ClientSecret: p.config.ClientSecret,
|
||||
RedirectURL: redirectURL,
|
||||
Scopes: p.config.Scopes,
|
||||
Endpoint: oauth2.Endpoint{
|
||||
AuthURL: p.config.AuthUrl,
|
||||
TokenURL: p.config.TokenUrl,
|
||||
AuthStyle: oauth2.AuthStyleAutoDetect,
|
||||
},
|
||||
}
|
||||
|
||||
// Prepare token exchange options
|
||||
opts := []oauth2.AuthCodeOption{}
|
||||
|
||||
// Add PKCE code_verifier if provided
|
||||
if codeVerifier != "" {
|
||||
opts = append(opts, oauth2.SetAuthURLParam("code_verifier", codeVerifier))
|
||||
}
|
||||
|
||||
token, err := conf.Exchange(ctx, code, opts...)
|
||||
if err != nil {
|
||||
return "", errors.Wrap(err, "failed to exchange access token")
|
||||
}
|
||||
|
||||
// Use the standard AccessToken field instead of Extra()
|
||||
// This is more reliable across different OAuth providers
|
||||
if token.AccessToken == "" {
|
||||
return "", errors.New("missing access token from authorization response")
|
||||
}
|
||||
|
||||
return token.AccessToken, nil
|
||||
}
|
||||
|
||||
// UserInfo returns the parsed user information using the given OAuth2 token.
|
||||
func (p *IdentityProvider) UserInfo(ctx context.Context, token string) (*idp.IdentityProviderUserInfo, error) {
|
||||
client := &http.Client{Timeout: userInfoRequestTimeout}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, p.config.UserInfoUrl, nil)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to create http request")
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to get user information")
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
body, readErr := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
if readErr != nil {
|
||||
return nil, errors.Wrap(readErr, "failed to read error response body")
|
||||
}
|
||||
return nil, errors.Errorf("userinfo request failed with status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to read response body")
|
||||
}
|
||||
|
||||
var claims map[string]any
|
||||
if err := json.Unmarshal(body, &claims); err != nil {
|
||||
return nil, errors.Wrap(err, "failed to unmarshal response body")
|
||||
}
|
||||
userInfo := &idp.IdentityProviderUserInfo{}
|
||||
if v, ok := claims[p.config.FieldMapping.Identifier].(string); ok {
|
||||
userInfo.Identifier = v
|
||||
}
|
||||
if userInfo.Identifier == "" {
|
||||
return nil, errors.Errorf("the field %q is not found in claims or has empty value", p.config.FieldMapping.Identifier)
|
||||
}
|
||||
|
||||
// Best effort to map optional fields
|
||||
if p.config.FieldMapping.DisplayName != "" {
|
||||
if v, ok := claims[p.config.FieldMapping.DisplayName].(string); ok {
|
||||
userInfo.DisplayName = v
|
||||
}
|
||||
}
|
||||
if userInfo.DisplayName == "" {
|
||||
userInfo.DisplayName = userInfo.Identifier
|
||||
}
|
||||
if p.config.FieldMapping.Email != "" {
|
||||
if v, ok := claims[p.config.FieldMapping.Email].(string); ok {
|
||||
userInfo.Email = v
|
||||
}
|
||||
}
|
||||
if p.config.FieldMapping.AvatarUrl != "" {
|
||||
if v, ok := claims[p.config.FieldMapping.AvatarUrl].(string); ok {
|
||||
userInfo.AvatarURL = v
|
||||
}
|
||||
}
|
||||
return userInfo, nil
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
package oauth2
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/usememos/memos/internal/idp"
|
||||
storepb "github.com/usememos/memos/proto/gen/store"
|
||||
)
|
||||
|
||||
func TestNewIdentityProvider(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
config *storepb.OAuth2Config
|
||||
containsErr string
|
||||
}{
|
||||
{
|
||||
name: "no tokenUrl",
|
||||
config: &storepb.OAuth2Config{
|
||||
ClientId: "test-client-id",
|
||||
ClientSecret: "test-client-secret",
|
||||
AuthUrl: "",
|
||||
TokenUrl: "",
|
||||
UserInfoUrl: "https://example.com/api/user",
|
||||
FieldMapping: &storepb.FieldMapping{
|
||||
Identifier: "login",
|
||||
},
|
||||
},
|
||||
containsErr: `the field "tokenUrl" is empty but required`,
|
||||
},
|
||||
{
|
||||
name: "no userInfoUrl",
|
||||
config: &storepb.OAuth2Config{
|
||||
ClientId: "test-client-id",
|
||||
ClientSecret: "test-client-secret",
|
||||
AuthUrl: "",
|
||||
TokenUrl: "https://example.com/token",
|
||||
UserInfoUrl: "",
|
||||
FieldMapping: &storepb.FieldMapping{
|
||||
Identifier: "login",
|
||||
},
|
||||
},
|
||||
containsErr: `the field "userInfoUrl" is empty but required`,
|
||||
},
|
||||
{
|
||||
name: "no field mapping identifier",
|
||||
config: &storepb.OAuth2Config{
|
||||
ClientId: "test-client-id",
|
||||
ClientSecret: "test-client-secret",
|
||||
AuthUrl: "",
|
||||
TokenUrl: "https://example.com/token",
|
||||
UserInfoUrl: "https://example.com/api/user",
|
||||
FieldMapping: &storepb.FieldMapping{
|
||||
Identifier: "",
|
||||
},
|
||||
},
|
||||
containsErr: `the field "fieldMapping.identifier" is empty but required`,
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(*testing.T) {
|
||||
_, err := NewIdentityProvider(test.config)
|
||||
assert.ErrorContains(t, err, test.containsErr)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func newMockServer(t *testing.T, code, accessToken string, userinfo []byte) *httptest.Server {
|
||||
mux := http.NewServeMux()
|
||||
|
||||
var rawIDToken string
|
||||
mux.HandleFunc("/oauth2/token", func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, http.MethodPost, r.Method)
|
||||
|
||||
body, err := io.ReadAll(r.Body)
|
||||
require.NoError(t, err)
|
||||
vals, err := url.ParseQuery(string(body))
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Equal(t, code, vals.Get("code"))
|
||||
require.Equal(t, "authorization_code", vals.Get("grant_type"))
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
err = json.NewEncoder(w).Encode(map[string]any{
|
||||
"access_token": accessToken,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
"id_token": rawIDToken,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
})
|
||||
mux.HandleFunc("/oauth2/userinfo", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, err := w.Write(userinfo)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
|
||||
s := httptest.NewServer(mux)
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
func TestIdentityProvider(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
const (
|
||||
testClientID = "test-client-id"
|
||||
testCode = "test-code"
|
||||
testAccessToken = "test-access-token"
|
||||
testSubject = "123456789"
|
||||
testName = "John Doe"
|
||||
testEmail = "john.doe@example.com"
|
||||
)
|
||||
userInfo, err := json.Marshal(
|
||||
map[string]any{
|
||||
"sub": testSubject,
|
||||
"name": testName,
|
||||
"email": testEmail,
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
s := newMockServer(t, testCode, testAccessToken, userInfo)
|
||||
|
||||
oauth2, err := NewIdentityProvider(
|
||||
&storepb.OAuth2Config{
|
||||
ClientId: testClientID,
|
||||
ClientSecret: "test-client-secret",
|
||||
TokenUrl: fmt.Sprintf("%s/oauth2/token", s.URL),
|
||||
UserInfoUrl: fmt.Sprintf("%s/oauth2/userinfo", s.URL),
|
||||
FieldMapping: &storepb.FieldMapping{
|
||||
Identifier: "sub",
|
||||
DisplayName: "name",
|
||||
Email: "email",
|
||||
},
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
redirectURL := "https://example.com/oauth/callback"
|
||||
// Test without PKCE (backward compatibility)
|
||||
oauthToken, err := oauth2.ExchangeToken(ctx, redirectURL, testCode, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, testAccessToken, oauthToken)
|
||||
|
||||
userInfoResult, err := oauth2.UserInfo(ctx, oauthToken)
|
||||
require.NoError(t, err)
|
||||
|
||||
wantUserInfo := &idp.IdentityProviderUserInfo{
|
||||
Identifier: testSubject,
|
||||
DisplayName: testName,
|
||||
Email: testEmail,
|
||||
}
|
||||
assert.Equal(t, wantUserInfo, userInfoResult)
|
||||
}
|
||||
|
||||
func TestIdentityProviderExchangeTokenClientAuthentication(t *testing.T) {
|
||||
const (
|
||||
clientID = "test-client-id"
|
||||
clientSecret = "test-client-secret"
|
||||
code = "test-code"
|
||||
accessToken = "test-access-token"
|
||||
codeVerifier = "test-code-verifier"
|
||||
)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
acceptBasicAuth bool
|
||||
expectedRequests int
|
||||
}{
|
||||
{
|
||||
name: "client secret basic",
|
||||
acceptBasicAuth: true,
|
||||
expectedRequests: 1,
|
||||
},
|
||||
{
|
||||
name: "client secret post fallback",
|
||||
acceptBasicAuth: false,
|
||||
expectedRequests: 2,
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
requestCount := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requestCount++
|
||||
require.NoError(t, r.ParseForm())
|
||||
require.Equal(t, code, r.Form.Get("code"))
|
||||
require.Equal(t, codeVerifier, r.Form.Get("code_verifier"))
|
||||
|
||||
username, password, hasBasicAuth := r.BasicAuth()
|
||||
if test.acceptBasicAuth {
|
||||
require.True(t, hasBasicAuth)
|
||||
require.Equal(t, clientID, username)
|
||||
require.Equal(t, clientSecret, password)
|
||||
require.Empty(t, r.Form.Get("client_id"))
|
||||
require.Empty(t, r.Form.Get("client_secret"))
|
||||
} else if hasBasicAuth {
|
||||
http.Error(w, `{"error":"invalid_client"}`, http.StatusUnauthorized)
|
||||
return
|
||||
} else {
|
||||
require.Equal(t, clientID, r.Form.Get("client_id"))
|
||||
require.Equal(t, clientSecret, r.Form.Get("client_secret"))
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
require.NoError(t, json.NewEncoder(w).Encode(map[string]any{
|
||||
"access_token": accessToken,
|
||||
"token_type": "Bearer",
|
||||
}))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
provider, err := NewIdentityProvider(&storepb.OAuth2Config{
|
||||
ClientId: clientID,
|
||||
ClientSecret: clientSecret,
|
||||
TokenUrl: server.URL,
|
||||
UserInfoUrl: "https://example.com/oauth2/userinfo",
|
||||
FieldMapping: &storepb.FieldMapping{Identifier: "sub"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
token, err := provider.ExchangeToken(context.Background(), "https://example.com/auth/callback", code, codeVerifier)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, accessToken, token)
|
||||
assert.Equal(t, test.expectedRequests, requestCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdentityProviderUserInfoUsesContext(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer s.Close()
|
||||
|
||||
oauth2, err := NewIdentityProvider(
|
||||
&storepb.OAuth2Config{
|
||||
ClientId: "test-client-id",
|
||||
ClientSecret: "test-client-secret",
|
||||
TokenUrl: "https://example.com/oauth2/token",
|
||||
UserInfoUrl: s.URL,
|
||||
FieldMapping: &storepb.FieldMapping{
|
||||
Identifier: "sub",
|
||||
},
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = oauth2.UserInfo(ctx, "test-access-token")
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "failed to get user information")
|
||||
}
|
||||
|
||||
func TestIdentityProviderUserInfoRejectsNon2xx(t *testing.T) {
|
||||
s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
http.Error(w, "upstream failure", http.StatusBadGateway)
|
||||
}))
|
||||
defer s.Close()
|
||||
|
||||
oauth2, err := NewIdentityProvider(
|
||||
&storepb.OAuth2Config{
|
||||
ClientId: "test-client-id",
|
||||
ClientSecret: "test-client-secret",
|
||||
TokenUrl: "https://example.com/oauth2/token",
|
||||
UserInfoUrl: s.URL,
|
||||
FieldMapping: &storepb.FieldMapping{
|
||||
Identifier: "sub",
|
||||
},
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = oauth2.UserInfo(context.Background(), "test-access-token")
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "userinfo request failed with status 502")
|
||||
assert.ErrorContains(t, err, "upstream failure")
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package ast
|
||||
|
||||
import (
|
||||
gast "github.com/yuin/goldmark/ast"
|
||||
)
|
||||
|
||||
// MentionNode represents an @mention in the markdown AST.
|
||||
type MentionNode struct {
|
||||
gast.BaseInline
|
||||
|
||||
// Username without the @ prefix.
|
||||
Username []byte
|
||||
}
|
||||
|
||||
// KindMention is the NodeKind for MentionNode.
|
||||
var KindMention = gast.NewNodeKind("Mention")
|
||||
|
||||
// Kind returns KindMention.
|
||||
func (*MentionNode) Kind() gast.NodeKind {
|
||||
return KindMention
|
||||
}
|
||||
|
||||
// Dump implements Node.Dump for debugging.
|
||||
func (n *MentionNode) Dump(source []byte, level int) {
|
||||
gast.DumpHelper(n, source, level, map[string]string{
|
||||
"Username": string(n.Username),
|
||||
}, nil)
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package ast
|
||||
|
||||
import (
|
||||
gast "github.com/yuin/goldmark/ast"
|
||||
)
|
||||
|
||||
// TagNode represents a #tag in the markdown AST.
|
||||
type TagNode struct {
|
||||
gast.BaseInline
|
||||
|
||||
// Tag name without the # prefix
|
||||
Tag []byte
|
||||
}
|
||||
|
||||
// KindTag is the NodeKind for TagNode.
|
||||
var KindTag = gast.NewNodeKind("Tag")
|
||||
|
||||
// Kind returns KindTag.
|
||||
func (*TagNode) Kind() gast.NodeKind {
|
||||
return KindTag
|
||||
}
|
||||
|
||||
// Dump implements Node.Dump for debugging.
|
||||
func (n *TagNode) Dump(source []byte, level int) {
|
||||
gast.DumpHelper(n, source, level, map[string]string{
|
||||
"Tag": string(n.Tag),
|
||||
}, nil)
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package extensions
|
||||
|
||||
import (
|
||||
"github.com/yuin/goldmark"
|
||||
"github.com/yuin/goldmark/parser"
|
||||
"github.com/yuin/goldmark/util"
|
||||
|
||||
mparser "github.com/usememos/memos/internal/markdown/parser"
|
||||
)
|
||||
|
||||
type mentionExtension struct{}
|
||||
|
||||
// MentionExtension is a goldmark extension for @mention syntax.
|
||||
var MentionExtension = &mentionExtension{}
|
||||
|
||||
// Extend extends the goldmark parser with mention support.
|
||||
func (*mentionExtension) Extend(m goldmark.Markdown) {
|
||||
m.Parser().AddOptions(
|
||||
parser.WithInlineParsers(
|
||||
// Priority 200 - run before standard link parser (500).
|
||||
util.Prioritized(mparser.NewMentionParser(), 200),
|
||||
),
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package extensions
|
||||
|
||||
import (
|
||||
"github.com/yuin/goldmark"
|
||||
"github.com/yuin/goldmark/parser"
|
||||
"github.com/yuin/goldmark/util"
|
||||
|
||||
mparser "github.com/usememos/memos/internal/markdown/parser"
|
||||
)
|
||||
|
||||
type tagExtension struct{}
|
||||
|
||||
// TagExtension is a goldmark extension for #tag syntax.
|
||||
var TagExtension = &tagExtension{}
|
||||
|
||||
// Extend extends the goldmark parser with tag support.
|
||||
func (*tagExtension) Extend(m goldmark.Markdown) {
|
||||
m.Parser().AddOptions(
|
||||
parser.WithInlineParsers(
|
||||
// Priority 200 - run before standard link parser (500)
|
||||
util.Prioritized(mparser.NewTagParser(), 200),
|
||||
),
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,489 @@
|
||||
package markdown
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
|
||||
"github.com/yuin/goldmark"
|
||||
gast "github.com/yuin/goldmark/ast"
|
||||
"github.com/yuin/goldmark/extension"
|
||||
east "github.com/yuin/goldmark/extension/ast"
|
||||
"github.com/yuin/goldmark/parser"
|
||||
"github.com/yuin/goldmark/text"
|
||||
|
||||
mast "github.com/usememos/memos/internal/markdown/ast"
|
||||
"github.com/usememos/memos/internal/markdown/extensions"
|
||||
"github.com/usememos/memos/internal/markdown/renderer"
|
||||
storepb "github.com/usememos/memos/proto/gen/store"
|
||||
)
|
||||
|
||||
// ExtractedData contains all metadata extracted from markdown in a single pass.
|
||||
type ExtractedData struct {
|
||||
Tags []string
|
||||
Mentions []string
|
||||
Property *storepb.MemoPayload_Property
|
||||
}
|
||||
|
||||
// Service handles markdown metadata extraction.
|
||||
// It uses goldmark to parse markdown and extract tags, properties, and snippets.
|
||||
// HTML rendering is primarily done on frontend using markdown-it, but backend provides
|
||||
// RenderHTML for RSS feeds and other server-side rendering needs.
|
||||
type Service interface {
|
||||
// ExtractAll extracts tags, properties, and references in a single parse (most efficient)
|
||||
ExtractAll(content []byte) (*ExtractedData, error)
|
||||
|
||||
// ExtractTags returns all #tags found in content
|
||||
ExtractTags(content []byte) ([]string, error)
|
||||
|
||||
// ExtractProperties computes boolean properties
|
||||
ExtractProperties(content []byte) (*storepb.MemoPayload_Property, error)
|
||||
|
||||
// RenderMarkdown renders goldmark AST back to markdown text
|
||||
RenderMarkdown(content []byte) (string, error)
|
||||
|
||||
// RenderHTML renders markdown content to HTML
|
||||
RenderHTML(content []byte) (string, error)
|
||||
|
||||
// GenerateSnippet creates plain text summary
|
||||
GenerateSnippet(content []byte, maxLength int) (string, error)
|
||||
|
||||
// ValidateContent checks for syntax errors
|
||||
ValidateContent(content []byte) error
|
||||
|
||||
// RenameTag renames all occurrences of oldTag to newTag in content
|
||||
RenameTag(content []byte, oldTag, newTag string) (string, error)
|
||||
}
|
||||
|
||||
// service implements the Service interface.
|
||||
type service struct {
|
||||
md goldmark.Markdown
|
||||
}
|
||||
|
||||
// Option configures the markdown service.
|
||||
type Option func(*config)
|
||||
|
||||
type config struct {
|
||||
enableTags bool
|
||||
enableMentions bool
|
||||
}
|
||||
|
||||
// WithTagExtension enables #tag parsing.
|
||||
func WithTagExtension() Option {
|
||||
return func(c *config) {
|
||||
c.enableTags = true
|
||||
}
|
||||
}
|
||||
|
||||
// WithMentionExtension enables @mention parsing.
|
||||
func WithMentionExtension() Option {
|
||||
return func(c *config) {
|
||||
c.enableMentions = true
|
||||
}
|
||||
}
|
||||
|
||||
// NewService creates a new markdown service with the given options.
|
||||
func NewService(opts ...Option) Service {
|
||||
cfg := &config{}
|
||||
for _, opt := range opts {
|
||||
opt(cfg)
|
||||
}
|
||||
|
||||
exts := []goldmark.Extender{
|
||||
extension.GFM, // GitHub Flavored Markdown (tables, strikethrough, task lists, autolinks)
|
||||
}
|
||||
|
||||
// Add custom extensions based on config
|
||||
if cfg.enableTags {
|
||||
exts = append(exts, extensions.TagExtension)
|
||||
}
|
||||
if cfg.enableMentions {
|
||||
exts = append(exts, extensions.MentionExtension)
|
||||
}
|
||||
|
||||
md := goldmark.New(
|
||||
goldmark.WithExtensions(exts...),
|
||||
goldmark.WithParserOptions(
|
||||
parser.WithAutoHeadingID(), // Generate heading IDs
|
||||
),
|
||||
)
|
||||
|
||||
return &service{
|
||||
md: md,
|
||||
}
|
||||
}
|
||||
|
||||
// parse is an internal helper to parse content into AST.
|
||||
func (s *service) parse(content []byte) (gast.Node, error) {
|
||||
reader := text.NewReader(content)
|
||||
doc := s.md.Parser().Parse(reader)
|
||||
return doc, nil
|
||||
}
|
||||
|
||||
func isTagNodeInLinkOrImage(n gast.Node) bool {
|
||||
for parent := n.Parent(); parent != nil; parent = parent.Parent() {
|
||||
switch parent.Kind() {
|
||||
case gast.KindLink, gast.KindImage:
|
||||
return true
|
||||
default:
|
||||
// Keep walking ancestors.
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func asMemoTagNode(n gast.Node) (*mast.TagNode, bool) {
|
||||
tagNode, ok := n.(*mast.TagNode)
|
||||
if !ok || isTagNodeInLinkOrImage(n) {
|
||||
return nil, false
|
||||
}
|
||||
return tagNode, true
|
||||
}
|
||||
|
||||
// ExtractTags returns all #tags found in content.
|
||||
func (s *service) ExtractTags(content []byte) ([]string, error) {
|
||||
root, err := s.parse(content)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var tags []string
|
||||
|
||||
// Walk the AST to find tag nodes
|
||||
err = gast.Walk(root, func(n gast.Node, entering bool) (gast.WalkStatus, error) {
|
||||
if !entering {
|
||||
return gast.WalkContinue, nil
|
||||
}
|
||||
|
||||
if tagNode, ok := asMemoTagNode(n); ok {
|
||||
tags = append(tags, string(tagNode.Tag))
|
||||
}
|
||||
|
||||
return gast.WalkContinue, nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Deduplicate tags while preserving original case
|
||||
return uniquePreserveCase(tags), nil
|
||||
}
|
||||
|
||||
// extractHeadingText extracts plain text content from a heading node.
|
||||
func extractHeadingText(n gast.Node, source []byte) string {
|
||||
var buf strings.Builder
|
||||
for child := n.FirstChild(); child != nil; child = child.NextSibling() {
|
||||
extractTextFromNode(child, source, &buf)
|
||||
}
|
||||
return buf.String()
|
||||
}
|
||||
|
||||
// extractTextFromNode recursively extracts plain text from a node and its children.
|
||||
func extractTextFromNode(n gast.Node, source []byte, buf *strings.Builder) {
|
||||
if textNode, ok := n.(*gast.Text); ok {
|
||||
buf.Write(textNode.Segment.Value(source))
|
||||
return
|
||||
}
|
||||
for child := n.FirstChild(); child != nil; child = child.NextSibling() {
|
||||
extractTextFromNode(child, source, buf)
|
||||
}
|
||||
}
|
||||
|
||||
// ExtractProperties computes boolean properties about the content.
|
||||
func (s *service) ExtractProperties(content []byte) (*storepb.MemoPayload_Property, error) {
|
||||
root, err := s.parse(content)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
prop := &storepb.MemoPayload_Property{}
|
||||
firstBlockChecked := false
|
||||
|
||||
err = gast.Walk(root, func(n gast.Node, entering bool) (gast.WalkStatus, error) {
|
||||
if !entering {
|
||||
return gast.WalkContinue, nil
|
||||
}
|
||||
|
||||
// Check if the first block-level child of the document is an H1 heading.
|
||||
if !firstBlockChecked && n.Parent() != nil && n.Parent().Kind() == gast.KindDocument {
|
||||
firstBlockChecked = true
|
||||
if heading, ok := n.(*gast.Heading); ok && heading.Level == 1 {
|
||||
prop.Title = extractHeadingText(n, content)
|
||||
}
|
||||
}
|
||||
|
||||
switch n.Kind() {
|
||||
case gast.KindLink:
|
||||
prop.HasLink = true
|
||||
|
||||
case gast.KindCodeBlock, gast.KindFencedCodeBlock, gast.KindCodeSpan:
|
||||
prop.HasCode = true
|
||||
|
||||
case east.KindTaskCheckBox:
|
||||
prop.HasTaskList = true
|
||||
if checkBox, ok := n.(*east.TaskCheckBox); ok {
|
||||
if !checkBox.IsChecked {
|
||||
prop.HasIncompleteTasks = true
|
||||
}
|
||||
}
|
||||
default:
|
||||
// No special handling for other node types
|
||||
}
|
||||
|
||||
return gast.WalkContinue, nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return prop, nil
|
||||
}
|
||||
|
||||
// RenderMarkdown renders goldmark AST back to markdown text.
|
||||
func (s *service) RenderMarkdown(content []byte) (string, error) {
|
||||
root, err := s.parse(content)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
mdRenderer := renderer.NewMarkdownRenderer()
|
||||
return mdRenderer.Render(root, content), nil
|
||||
}
|
||||
|
||||
// RenderHTML renders markdown content to HTML using goldmark's built-in HTML renderer.
|
||||
func (s *service) RenderHTML(content []byte) (string, error) {
|
||||
var buf bytes.Buffer
|
||||
if err := s.md.Convert(content, &buf); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return buf.String(), nil
|
||||
}
|
||||
|
||||
// GenerateSnippet creates a plain text summary from markdown content.
|
||||
func (s *service) GenerateSnippet(content []byte, maxLength int) (string, error) {
|
||||
root, err := s.parse(content)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
var buf strings.Builder
|
||||
var lastNodeWasBlock bool
|
||||
|
||||
err = gast.Walk(root, func(n gast.Node, entering bool) (gast.WalkStatus, error) {
|
||||
if entering {
|
||||
// Skip code blocks entirely (but keep inline code spans for snippet text)
|
||||
switch n.Kind() {
|
||||
case gast.KindCodeBlock, gast.KindFencedCodeBlock:
|
||||
return gast.WalkSkipChildren, nil
|
||||
default:
|
||||
// Continue walking for other node types
|
||||
}
|
||||
|
||||
// Add space before block elements (except first)
|
||||
switch n.Kind() {
|
||||
case gast.KindParagraph, gast.KindHeading, gast.KindListItem, east.KindTableCell, east.KindTableRow, east.KindTableHeader:
|
||||
if buf.Len() > 0 && lastNodeWasBlock {
|
||||
buf.WriteByte(' ')
|
||||
}
|
||||
default:
|
||||
// No space needed for other node types
|
||||
}
|
||||
}
|
||||
|
||||
if !entering {
|
||||
// Mark that we just exited a block element
|
||||
switch n.Kind() {
|
||||
case gast.KindParagraph, gast.KindHeading, gast.KindListItem, east.KindTableCell, east.KindTableRow, east.KindTableHeader:
|
||||
lastNodeWasBlock = true
|
||||
default:
|
||||
// Not a block element
|
||||
}
|
||||
return gast.WalkContinue, nil
|
||||
}
|
||||
|
||||
lastNodeWasBlock = false
|
||||
|
||||
// Extract text from various node types
|
||||
switch node := n.(type) {
|
||||
case *gast.Text:
|
||||
segment := node.Segment
|
||||
buf.Write(segment.Value(content))
|
||||
if node.SoftLineBreak() {
|
||||
buf.WriteByte(' ')
|
||||
}
|
||||
case *gast.AutoLink:
|
||||
buf.Write(node.URL(content))
|
||||
return gast.WalkSkipChildren, nil
|
||||
case *mast.TagNode:
|
||||
buf.WriteByte('#')
|
||||
buf.Write(node.Tag)
|
||||
default:
|
||||
// Ignore other node types.
|
||||
}
|
||||
|
||||
// Stop walking if we've exceeded double the max length
|
||||
// (we'll truncate precisely later)
|
||||
if buf.Len() > maxLength*2 {
|
||||
return gast.WalkStop, nil
|
||||
}
|
||||
|
||||
return gast.WalkContinue, nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
snippet := buf.String()
|
||||
|
||||
// Truncate at word boundary if needed
|
||||
if len(snippet) > maxLength {
|
||||
snippet = truncateAtWord(snippet, maxLength)
|
||||
}
|
||||
|
||||
return strings.TrimSpace(snippet), nil
|
||||
}
|
||||
|
||||
// ValidateContent checks if the markdown content is valid.
|
||||
func (s *service) ValidateContent(content []byte) error {
|
||||
// Try to parse the content
|
||||
_, err := s.parse(content)
|
||||
return err
|
||||
}
|
||||
|
||||
// ExtractAll extracts tags, properties, and references in a single parse for efficiency.
|
||||
func (s *service) ExtractAll(content []byte) (*ExtractedData, error) {
|
||||
root, err := s.parse(content)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
data := &ExtractedData{
|
||||
Tags: []string{},
|
||||
Mentions: []string{},
|
||||
Property: &storepb.MemoPayload_Property{},
|
||||
}
|
||||
|
||||
firstBlockChecked := false
|
||||
|
||||
// Single walk to collect all data
|
||||
err = gast.Walk(root, func(n gast.Node, entering bool) (gast.WalkStatus, error) {
|
||||
if !entering {
|
||||
return gast.WalkContinue, nil
|
||||
}
|
||||
|
||||
if tagNode, ok := asMemoTagNode(n); ok {
|
||||
data.Tags = append(data.Tags, string(tagNode.Tag))
|
||||
}
|
||||
if mentionNode, ok := n.(*mast.MentionNode); ok {
|
||||
data.Mentions = append(data.Mentions, strings.ToLower(string(mentionNode.Username)))
|
||||
}
|
||||
|
||||
// Check if the first block-level child of the document is an H1 heading.
|
||||
if !firstBlockChecked && n.Parent() != nil && n.Parent().Kind() == gast.KindDocument {
|
||||
firstBlockChecked = true
|
||||
if heading, ok := n.(*gast.Heading); ok && heading.Level == 1 {
|
||||
data.Property.Title = extractHeadingText(n, content)
|
||||
}
|
||||
}
|
||||
|
||||
// Extract properties based on node kind
|
||||
switch n.Kind() {
|
||||
case gast.KindLink:
|
||||
data.Property.HasLink = true
|
||||
|
||||
case gast.KindCodeBlock, gast.KindFencedCodeBlock, gast.KindCodeSpan:
|
||||
data.Property.HasCode = true
|
||||
|
||||
case east.KindTaskCheckBox:
|
||||
data.Property.HasTaskList = true
|
||||
if checkBox, ok := n.(*east.TaskCheckBox); ok {
|
||||
if !checkBox.IsChecked {
|
||||
data.Property.HasIncompleteTasks = true
|
||||
}
|
||||
}
|
||||
default:
|
||||
// No special handling for other node types
|
||||
}
|
||||
|
||||
return gast.WalkContinue, nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Deduplicate tags while preserving original case
|
||||
data.Tags = uniquePreserveCase(data.Tags)
|
||||
data.Mentions = uniquePreserveCase(data.Mentions)
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// RenameTag renames all occurrences of oldTag to newTag in content.
|
||||
func (s *service) RenameTag(content []byte, oldTag, newTag string) (string, error) {
|
||||
root, err := s.parse(content)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Walk the AST to find and rename tag nodes
|
||||
err = gast.Walk(root, func(n gast.Node, entering bool) (gast.WalkStatus, error) {
|
||||
if !entering {
|
||||
return gast.WalkContinue, nil
|
||||
}
|
||||
|
||||
if tagNode, ok := asMemoTagNode(n); ok {
|
||||
if string(tagNode.Tag) == oldTag {
|
||||
tagNode.Tag = []byte(newTag)
|
||||
}
|
||||
}
|
||||
|
||||
return gast.WalkContinue, nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Render back to markdown using the already-parsed AST
|
||||
mdRenderer := renderer.NewMarkdownRenderer()
|
||||
return mdRenderer.Render(root, content), nil
|
||||
}
|
||||
|
||||
// uniquePreserveCase returns unique strings from input while preserving case.
|
||||
func uniquePreserveCase(strs []string) []string {
|
||||
seen := make(map[string]struct{})
|
||||
var result []string
|
||||
|
||||
for _, s := range strs {
|
||||
if _, exists := seen[s]; !exists {
|
||||
seen[s] = struct{}{}
|
||||
result = append(result, s)
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// truncateAtWord truncates a string at the last word boundary before maxLength.
|
||||
// maxLength is treated as a rune (character) count to properly handle UTF-8 multi-byte characters.
|
||||
func truncateAtWord(s string, maxLength int) string {
|
||||
// Convert to runes to properly handle multi-byte UTF-8 characters
|
||||
runes := []rune(s)
|
||||
if len(runes) <= maxLength {
|
||||
return s
|
||||
}
|
||||
|
||||
// Truncate to max length (by character count, not byte count)
|
||||
truncated := string(runes[:maxLength])
|
||||
|
||||
// Find last space to avoid cutting in the middle of a word
|
||||
lastSpace := strings.LastIndexAny(truncated, " \t\n\r")
|
||||
if lastSpace > 0 {
|
||||
truncated = truncated[:lastSpace]
|
||||
}
|
||||
|
||||
return truncated + " ..."
|
||||
}
|
||||
@@ -0,0 +1,654 @@
|
||||
package markdown
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNewService(t *testing.T) {
|
||||
svc := NewService()
|
||||
assert.NotNil(t, svc)
|
||||
}
|
||||
|
||||
func TestValidateContent(t *testing.T) {
|
||||
svc := NewService()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
content string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid markdown",
|
||||
content: "# Hello\n\nThis is **bold** text.",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "empty content",
|
||||
content: "",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "complex markdown",
|
||||
content: "# Title\n\n- List item 1\n- List item 2\n\n```go\ncode block\n```",
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := svc.ValidateContent([]byte(tt.content))
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateSnippet(t *testing.T) {
|
||||
svc := NewService()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
content string
|
||||
maxLength int
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "simple text",
|
||||
content: "Hello world",
|
||||
maxLength: 100,
|
||||
expected: "Hello world",
|
||||
},
|
||||
{
|
||||
name: "text with formatting",
|
||||
content: "This is **bold** and *italic* text.",
|
||||
maxLength: 100,
|
||||
expected: "This is bold and italic text.",
|
||||
},
|
||||
{
|
||||
name: "truncate long text",
|
||||
content: "This is a very long piece of text that should be truncated at a word boundary.",
|
||||
maxLength: 30,
|
||||
expected: "This is a very long piece of ...",
|
||||
},
|
||||
{
|
||||
name: "heading and paragraph",
|
||||
content: "# My Title\n\nThis is the first paragraph.",
|
||||
maxLength: 100,
|
||||
expected: "My Title This is the first paragraph.",
|
||||
},
|
||||
{
|
||||
name: "code block removed",
|
||||
content: "Text before\n\n```go\ncode\n```\n\nText after",
|
||||
maxLength: 100,
|
||||
expected: "Text before Text after",
|
||||
},
|
||||
{
|
||||
name: "list items",
|
||||
content: "- Item 1\n- Item 2\n- Item 3",
|
||||
maxLength: 100,
|
||||
expected: "Item 1 Item 2 Item 3",
|
||||
},
|
||||
{
|
||||
name: "inline code preserved",
|
||||
content: "`console.log('hello')`",
|
||||
maxLength: 100,
|
||||
expected: "console.log('hello')",
|
||||
},
|
||||
{
|
||||
name: "text with inline code",
|
||||
content: "Use `fmt.Println` to print output.",
|
||||
maxLength: 100,
|
||||
expected: "Use fmt.Println to print output.",
|
||||
},
|
||||
{
|
||||
name: "image alt text",
|
||||
content: "",
|
||||
maxLength: 100,
|
||||
expected: "alt text",
|
||||
},
|
||||
{
|
||||
name: "strikethrough text",
|
||||
content: "~~deleted text~~",
|
||||
maxLength: 100,
|
||||
expected: "deleted text",
|
||||
},
|
||||
{
|
||||
name: "blockquote",
|
||||
content: "> quoted text",
|
||||
maxLength: 100,
|
||||
expected: "quoted text",
|
||||
},
|
||||
{
|
||||
name: "table cells spaced",
|
||||
content: "| a | b |\n|---|---|\n| 1 | 2 |",
|
||||
maxLength: 100,
|
||||
expected: "a b 1 2",
|
||||
},
|
||||
{
|
||||
name: "plain URL autolink",
|
||||
content: "https://usememos.com",
|
||||
maxLength: 100,
|
||||
expected: "https://usememos.com",
|
||||
},
|
||||
{
|
||||
name: "text with plain URL",
|
||||
content: "Check out https://usememos.com for more info.",
|
||||
maxLength: 100,
|
||||
expected: "Check out https://usememos.com for more info.",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
snippet, err := svc.GenerateSnippet([]byte(tt.content), tt.maxLength)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tt.expected, snippet)
|
||||
})
|
||||
}
|
||||
|
||||
// Test with tag extension enabled (matches production config).
|
||||
svcWithTags := NewService(WithTagExtension())
|
||||
tagTests := []struct {
|
||||
name string
|
||||
content string
|
||||
maxLength int
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "tag only",
|
||||
content: "#todo",
|
||||
maxLength: 100,
|
||||
expected: "#todo",
|
||||
},
|
||||
{
|
||||
name: "text with tags",
|
||||
content: "Remember to #review the #code",
|
||||
maxLength: 100,
|
||||
expected: "Remember to #review the #code",
|
||||
},
|
||||
}
|
||||
for _, tt := range tagTests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
snippet, err := svcWithTags.GenerateSnippet([]byte(tt.content), tt.maxLength)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tt.expected, snippet)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractProperties(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
content string
|
||||
hasLink bool
|
||||
hasCode bool
|
||||
hasTasks bool
|
||||
hasInc bool
|
||||
title string
|
||||
}{
|
||||
{
|
||||
name: "plain text",
|
||||
content: "Just plain text",
|
||||
hasLink: false,
|
||||
hasCode: false,
|
||||
hasTasks: false,
|
||||
hasInc: false,
|
||||
title: "",
|
||||
},
|
||||
{
|
||||
name: "with link",
|
||||
content: "Check out [this link](https://example.com)",
|
||||
hasLink: true,
|
||||
hasCode: false,
|
||||
hasTasks: false,
|
||||
hasInc: false,
|
||||
title: "",
|
||||
},
|
||||
{
|
||||
name: "with inline code",
|
||||
content: "Use `console.log()` to debug",
|
||||
hasLink: false,
|
||||
hasCode: true,
|
||||
hasTasks: false,
|
||||
hasInc: false,
|
||||
title: "",
|
||||
},
|
||||
{
|
||||
name: "with code block",
|
||||
content: "```go\nfunc main() {}\n```",
|
||||
hasLink: false,
|
||||
hasCode: true,
|
||||
hasTasks: false,
|
||||
hasInc: false,
|
||||
title: "",
|
||||
},
|
||||
{
|
||||
name: "with completed task",
|
||||
content: "- [x] Completed task",
|
||||
hasLink: false,
|
||||
hasCode: false,
|
||||
hasTasks: true,
|
||||
hasInc: false,
|
||||
title: "",
|
||||
},
|
||||
{
|
||||
name: "with incomplete task",
|
||||
content: "- [ ] Todo item",
|
||||
hasLink: false,
|
||||
hasCode: false,
|
||||
hasTasks: true,
|
||||
hasInc: true,
|
||||
title: "",
|
||||
},
|
||||
{
|
||||
name: "mixed tasks",
|
||||
content: "- [x] Done\n- [ ] Not done",
|
||||
hasLink: false,
|
||||
hasCode: false,
|
||||
hasTasks: true,
|
||||
hasInc: true,
|
||||
title: "",
|
||||
},
|
||||
{
|
||||
name: "everything",
|
||||
content: "# Title\n\n[Link](url)\n\n`code`\n\n- [ ] Task",
|
||||
hasLink: true,
|
||||
hasCode: true,
|
||||
hasTasks: true,
|
||||
hasInc: true,
|
||||
title: "Title",
|
||||
},
|
||||
{
|
||||
name: "h1 as first node extracts title",
|
||||
content: "# My Article Title\n\nBody text here.",
|
||||
title: "My Article Title",
|
||||
},
|
||||
{
|
||||
name: "h2 as first node does not extract title",
|
||||
content: "## Sub Heading\n\nBody text.",
|
||||
title: "",
|
||||
},
|
||||
{
|
||||
name: "h1 not first node does not extract title",
|
||||
content: "Some text\n\n# Heading Later",
|
||||
title: "",
|
||||
},
|
||||
{
|
||||
name: "h1 with inline formatting extracts plain text",
|
||||
content: "# Title with **bold** and *italic*\n\nBody.",
|
||||
title: "Title with bold and italic",
|
||||
},
|
||||
{
|
||||
name: "empty content has no title",
|
||||
content: "",
|
||||
title: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
svc := NewService()
|
||||
|
||||
props, err := svc.ExtractProperties([]byte(tt.content))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tt.hasLink, props.HasLink, "HasLink")
|
||||
assert.Equal(t, tt.hasCode, props.HasCode, "HasCode")
|
||||
assert.Equal(t, tt.hasTasks, props.HasTaskList, "HasTaskList")
|
||||
assert.Equal(t, tt.hasInc, props.HasIncompleteTasks, "HasIncompleteTasks")
|
||||
assert.Equal(t, tt.title, props.Title, "Title")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractAllTitle(t *testing.T) {
|
||||
svc := NewService(WithTagExtension())
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
content string
|
||||
title string
|
||||
}{
|
||||
{
|
||||
name: "h1 first node",
|
||||
content: "# Article Title\n\nContent with #tag",
|
||||
title: "Article Title",
|
||||
},
|
||||
{
|
||||
name: "no h1",
|
||||
content: "Just text with #tag",
|
||||
title: "",
|
||||
},
|
||||
{
|
||||
name: "h1 not first",
|
||||
content: "Intro\n\n# Late Heading",
|
||||
title: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
data, err := svc.ExtractAll([]byte(tt.content))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tt.title, data.Property.Title, "Title")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractAllMentions(t *testing.T) {
|
||||
svc := NewService(WithTagExtension(), WithMentionExtension())
|
||||
|
||||
data, err := svc.ExtractAll([]byte("Hi @Alice and @bob. Email support@example.com should stay plain. #tag"))
|
||||
require.NoError(t, err)
|
||||
assert.ElementsMatch(t, []string{"alice", "bob"}, data.Mentions)
|
||||
assert.ElementsMatch(t, []string{"tag"}, data.Tags)
|
||||
}
|
||||
|
||||
func TestExtractAllSkipsTagsInsideLinks(t *testing.T) {
|
||||
svc := NewService(WithTagExtension())
|
||||
|
||||
data, err := svc.ExtractAll([]byte(
|
||||
"[release #notes](https://example.com/releases#release-notes)\n\n" +
|
||||
"\n\n" +
|
||||
"Outside #memo-tag",
|
||||
))
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.ElementsMatch(t, []string{"memo-tag"}, data.Tags)
|
||||
assert.True(t, data.Property.HasLink)
|
||||
}
|
||||
|
||||
func TestExtractTags(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
content string
|
||||
withExt bool
|
||||
expected []string
|
||||
}{
|
||||
{
|
||||
name: "no tags",
|
||||
content: "Just plain text",
|
||||
withExt: false,
|
||||
expected: []string{},
|
||||
},
|
||||
{
|
||||
name: "single tag",
|
||||
content: "Text with #tag",
|
||||
withExt: true,
|
||||
expected: []string{"tag"},
|
||||
},
|
||||
{
|
||||
name: "multiple tags",
|
||||
content: "Text with #tag1 and #tag2",
|
||||
withExt: true,
|
||||
expected: []string{"tag1", "tag2"},
|
||||
},
|
||||
{
|
||||
name: "duplicate tags",
|
||||
content: "#work is important. #Work #WORK",
|
||||
withExt: true,
|
||||
expected: []string{"work", "Work", "WORK"},
|
||||
},
|
||||
{
|
||||
name: "tags with hyphens and underscores",
|
||||
content: "Tags: #work-notes #2024_plans",
|
||||
withExt: true,
|
||||
expected: []string{"work-notes", "2024_plans"},
|
||||
},
|
||||
{
|
||||
name: "tags at end of sentence",
|
||||
content: "This is important #urgent.",
|
||||
withExt: true,
|
||||
expected: []string{"urgent"},
|
||||
},
|
||||
{
|
||||
name: "headings not tags",
|
||||
content: "## Heading\n\n# Title\n\nText with #realtag",
|
||||
withExt: true,
|
||||
expected: []string{"realtag"},
|
||||
},
|
||||
{
|
||||
name: "numeric tag",
|
||||
content: "Issue #123",
|
||||
withExt: true,
|
||||
expected: []string{"123"},
|
||||
},
|
||||
{
|
||||
name: "tag in list",
|
||||
content: "- Item 1 #todo\n- Item 2 #done",
|
||||
withExt: true,
|
||||
expected: []string{"todo", "done"},
|
||||
},
|
||||
{
|
||||
name: "autolink URL fragment not tag",
|
||||
content: "https://github.com/dmtrKovalenko/fff#pi-agent-extension\n\nProject #memo-tag",
|
||||
withExt: true,
|
||||
expected: []string{"memo-tag"},
|
||||
},
|
||||
{
|
||||
name: "markdown link text and fragment not tags",
|
||||
content: "[release #notes](https://example.com/releases#release-notes) Outside #memo-tag",
|
||||
withExt: true,
|
||||
expected: []string{"memo-tag"},
|
||||
},
|
||||
{
|
||||
name: "reference link text and fragment not tags",
|
||||
content: "[reference #anchor][docs]\n\n[docs]: https://example.com/docs#reference-anchor\n\nOutside #memo-tag",
|
||||
withExt: true,
|
||||
expected: []string{"memo-tag"},
|
||||
},
|
||||
{
|
||||
name: "image alt text and fragment not tags",
|
||||
content: "\n\nOutside #memo-tag",
|
||||
withExt: true,
|
||||
expected: []string{"memo-tag"},
|
||||
},
|
||||
{
|
||||
name: "no extension enabled",
|
||||
content: "Text with #tag",
|
||||
withExt: false,
|
||||
expected: []string{},
|
||||
},
|
||||
{
|
||||
name: "Chinese tag",
|
||||
content: "Text with #测试",
|
||||
withExt: true,
|
||||
expected: []string{"测试"},
|
||||
},
|
||||
{
|
||||
name: "Chinese tag followed by punctuation",
|
||||
content: "Text #测试。 More text",
|
||||
withExt: true,
|
||||
expected: []string{"测试"},
|
||||
},
|
||||
{
|
||||
name: "mixed Chinese and ASCII tag",
|
||||
content: "#测试test123 content",
|
||||
withExt: true,
|
||||
expected: []string{"测试test123"},
|
||||
},
|
||||
{
|
||||
name: "Japanese tag",
|
||||
content: "#日本語 content",
|
||||
withExt: true,
|
||||
expected: []string{"日本語"},
|
||||
},
|
||||
{
|
||||
name: "Korean tag",
|
||||
content: "#한국어 content",
|
||||
withExt: true,
|
||||
expected: []string{"한국어"},
|
||||
},
|
||||
{
|
||||
name: "hierarchical tag with Chinese",
|
||||
content: "#work/测试/项目",
|
||||
withExt: true,
|
||||
expected: []string{"work/测试/项目"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
var svc Service
|
||||
if tt.withExt {
|
||||
svc = NewService(WithTagExtension())
|
||||
} else {
|
||||
svc = NewService()
|
||||
}
|
||||
|
||||
tags, err := svc.ExtractTags([]byte(tt.content))
|
||||
require.NoError(t, err)
|
||||
assert.ElementsMatch(t, tt.expected, tags)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenameTagSkipsTagsInsideLinks(t *testing.T) {
|
||||
svc := NewService(WithTagExtension())
|
||||
|
||||
result, err := svc.RenameTag(
|
||||
[]byte("[release #notes](https://example.com/releases#release-notes)\n\nOutside #notes"),
|
||||
"notes",
|
||||
"done",
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "[release #notes](https://example.com/releases#release-notes)\n\nOutside #done", result)
|
||||
}
|
||||
|
||||
func TestUniquePreserveCase(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input []string
|
||||
expected []string
|
||||
}{
|
||||
{
|
||||
name: "empty",
|
||||
input: []string{},
|
||||
expected: []string{},
|
||||
},
|
||||
{
|
||||
name: "unique items",
|
||||
input: []string{"tag1", "tag2", "tag3"},
|
||||
expected: []string{"tag1", "tag2", "tag3"},
|
||||
},
|
||||
{
|
||||
name: "duplicates",
|
||||
input: []string{"tag", "TAG", "Tag"},
|
||||
expected: []string{"tag", "TAG", "Tag"},
|
||||
},
|
||||
{
|
||||
name: "mixed",
|
||||
input: []string{"Work", "work", "Important", "work"},
|
||||
expected: []string{"Work", "work", "Important"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := uniquePreserveCase(tt.input)
|
||||
assert.ElementsMatch(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTruncateAtWord(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
maxLength int
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "no truncation needed",
|
||||
input: "short",
|
||||
maxLength: 10,
|
||||
expected: "short",
|
||||
},
|
||||
{
|
||||
name: "exact length",
|
||||
input: "exactly ten",
|
||||
maxLength: 11,
|
||||
expected: "exactly ten",
|
||||
},
|
||||
{
|
||||
name: "truncate at word",
|
||||
input: "this is a long sentence",
|
||||
maxLength: 10,
|
||||
expected: "this is a ...",
|
||||
},
|
||||
{
|
||||
name: "truncate very long word",
|
||||
input: "supercalifragilisticexpialidocious",
|
||||
maxLength: 10,
|
||||
expected: "supercalif ...",
|
||||
},
|
||||
{
|
||||
name: "CJK characters without spaces",
|
||||
input: "这是一个很长的中文句子没有空格的情况下也要正确处理",
|
||||
maxLength: 15,
|
||||
expected: "这是一个很长的中文句子没有空格 ...",
|
||||
},
|
||||
{
|
||||
name: "mixed CJK and Latin",
|
||||
input: "这是中文mixed with English文字",
|
||||
maxLength: 10,
|
||||
expected: "这是中文mixed ...",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := truncateAtWord(tt.input, tt.maxLength)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Benchmark tests.
|
||||
func BenchmarkGenerateSnippet(b *testing.B) {
|
||||
svc := NewService()
|
||||
content := []byte(`# Large Document
|
||||
|
||||
This is a large document with multiple paragraphs and formatting.
|
||||
|
||||
## Section 1
|
||||
|
||||
Here is some **bold** text and *italic* text with [links](https://example.com).
|
||||
|
||||
- List item 1
|
||||
- List item 2
|
||||
- List item 3
|
||||
|
||||
## Section 2
|
||||
|
||||
More content here with ` + "`inline code`" + ` and other elements.
|
||||
|
||||
` + "```go\nfunc example() {\n return true\n}\n```")
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := svc.GenerateSnippet(content, 200)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkExtractProperties(b *testing.B) {
|
||||
svc := NewService()
|
||||
content := []byte("# Title\n\n[Link](url)\n\n`code`\n\n- [ ] Task\n- [x] Done")
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err := svc.ExtractProperties(content)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
|
||||
gast "github.com/yuin/goldmark/ast"
|
||||
"github.com/yuin/goldmark/parser"
|
||||
"github.com/yuin/goldmark/text"
|
||||
|
||||
mast "github.com/usememos/memos/internal/markdown/ast"
|
||||
)
|
||||
|
||||
const (
|
||||
// MaxMentionLength matches the username token length accepted by the API.
|
||||
MaxMentionLength = 32
|
||||
)
|
||||
|
||||
type mentionParser struct{}
|
||||
|
||||
// NewMentionParser creates a new inline parser for @mention syntax.
|
||||
func NewMentionParser() parser.InlineParser {
|
||||
return &mentionParser{}
|
||||
}
|
||||
|
||||
// Trigger returns the characters that trigger this parser.
|
||||
func (*mentionParser) Trigger() []byte {
|
||||
return []byte{'@'}
|
||||
}
|
||||
|
||||
func isValidMentionRune(r rune) bool {
|
||||
return unicode.IsLetter(r) || unicode.IsNumber(r) || r == '-'
|
||||
}
|
||||
|
||||
func isMentionBoundary(r rune) bool {
|
||||
return unicode.IsSpace(r) || unicode.IsPunct(r) || unicode.IsSymbol(r)
|
||||
}
|
||||
|
||||
// Parse parses @mention syntax while avoiding email-address matches.
|
||||
func (*mentionParser) Parse(_ gast.Node, block text.Reader, _ parser.Context) gast.Node {
|
||||
line, _ := block.PeekLine()
|
||||
if len(line) == 0 || line[0] != '@' {
|
||||
return nil
|
||||
}
|
||||
|
||||
prev := block.PrecendingCharacter()
|
||||
if prev != '\n' && !isMentionBoundary(prev) {
|
||||
return nil
|
||||
}
|
||||
|
||||
start := 1
|
||||
pos := start
|
||||
runeCount := 0
|
||||
hasLetterOrNumber := false
|
||||
|
||||
for pos < len(line) {
|
||||
r, size := utf8.DecodeRune(line[pos:])
|
||||
if r == utf8.RuneError && size == 1 {
|
||||
break
|
||||
}
|
||||
if !isValidMentionRune(r) {
|
||||
break
|
||||
}
|
||||
if unicode.IsLetter(r) || unicode.IsNumber(r) {
|
||||
hasLetterOrNumber = true
|
||||
}
|
||||
runeCount++
|
||||
if runeCount > MaxMentionLength {
|
||||
break
|
||||
}
|
||||
pos += size
|
||||
}
|
||||
|
||||
if pos <= start || !hasLetterOrNumber {
|
||||
return nil
|
||||
}
|
||||
|
||||
username := line[start:pos]
|
||||
usernameCopy := make([]byte, len(username))
|
||||
copy(usernameCopy, username)
|
||||
|
||||
block.Advance(pos)
|
||||
|
||||
return &mast.MentionNode{
|
||||
Username: usernameCopy,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
|
||||
gast "github.com/yuin/goldmark/ast"
|
||||
"github.com/yuin/goldmark/parser"
|
||||
"github.com/yuin/goldmark/text"
|
||||
|
||||
mast "github.com/usememos/memos/internal/markdown/ast"
|
||||
)
|
||||
|
||||
const (
|
||||
// MaxTagLength defines the maximum number of runes allowed in a tag.
|
||||
MaxTagLength = 100
|
||||
)
|
||||
|
||||
type tagParser struct{}
|
||||
|
||||
// NewTagParser creates a new inline parser for #tag syntax.
|
||||
func NewTagParser() parser.InlineParser {
|
||||
return &tagParser{}
|
||||
}
|
||||
|
||||
// Trigger returns the characters that trigger this parser.
|
||||
func (*tagParser) Trigger() []byte {
|
||||
return []byte{'#'}
|
||||
}
|
||||
|
||||
// isValidTagRune checks if a Unicode rune is valid in a tag.
|
||||
// Uses Unicode categories for proper international character support.
|
||||
func isValidTagRune(r rune) bool {
|
||||
// Allow Unicode letters (any script: Latin, CJK, Arabic, Cyrillic, etc.)
|
||||
if unicode.IsLetter(r) {
|
||||
return true
|
||||
}
|
||||
|
||||
// Allow Unicode digits
|
||||
if unicode.IsNumber(r) {
|
||||
return true
|
||||
}
|
||||
|
||||
// Allow emoji and symbols (So category: Symbol, Other)
|
||||
// This includes emoji, which are essential for social media-style tagging
|
||||
if unicode.IsSymbol(r) {
|
||||
return true
|
||||
}
|
||||
|
||||
// Allow marks (non-spacing, spacing combining, enclosing)
|
||||
// This covers variation selectors (e.g. VS16 \uFE0F) and combining marks (e.g. Keycap \u20E3, accents)
|
||||
if unicode.IsMark(r) {
|
||||
return true
|
||||
}
|
||||
|
||||
// Allow Zero Width Joiner (ZWJ) for emoji sequences
|
||||
if r == '\u200D' {
|
||||
return true
|
||||
}
|
||||
|
||||
// Allow specific ASCII symbols for tag structure
|
||||
// Underscore: word separation (snake_case)
|
||||
// Hyphen: word separation (kebab-case)
|
||||
// Forward slash: hierarchical tags (category/subcategory)
|
||||
// Ampersand: compound tags (science&tech)
|
||||
if r == '_' || r == '-' || r == '/' || r == '&' {
|
||||
return true
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// Parse parses #tag syntax using Unicode-aware validation.
|
||||
// Tags support international characters and follow these rules:
|
||||
// - Must start with # followed by valid tag characters
|
||||
// - Valid characters: Unicode letters, Unicode digits, underscore (_), hyphen (-), forward slash (/)
|
||||
// - Maximum length: 100 runes (Unicode characters)
|
||||
// - Stops at: whitespace, punctuation, or other invalid characters
|
||||
func (*tagParser) Parse(_ gast.Node, block text.Reader, _ parser.Context) gast.Node {
|
||||
line, _ := block.PeekLine()
|
||||
|
||||
// Must start with #
|
||||
if len(line) == 0 || line[0] != '#' {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if it's a heading (## or space after #)
|
||||
if len(line) > 1 {
|
||||
if line[1] == '#' {
|
||||
// It's a heading (##), not a tag
|
||||
return nil
|
||||
}
|
||||
if line[1] == ' ' {
|
||||
// Space after # - heading or just a hash
|
||||
return nil
|
||||
}
|
||||
} else {
|
||||
// Just a lone #
|
||||
return nil
|
||||
}
|
||||
|
||||
// Parse tag using UTF-8 aware rune iteration
|
||||
tagStart := 1
|
||||
pos := tagStart
|
||||
runeCount := 0
|
||||
|
||||
for pos < len(line) {
|
||||
r, size := utf8.DecodeRune(line[pos:])
|
||||
|
||||
// Stop at invalid UTF-8
|
||||
if r == utf8.RuneError && size == 1 {
|
||||
break
|
||||
}
|
||||
|
||||
// Validate character using Unicode categories
|
||||
if !isValidTagRune(r) {
|
||||
break
|
||||
}
|
||||
|
||||
// Enforce max length (by rune count, not byte count)
|
||||
runeCount++
|
||||
if runeCount > MaxTagLength {
|
||||
break
|
||||
}
|
||||
|
||||
pos += size
|
||||
}
|
||||
|
||||
// Must have at least one character after #
|
||||
if pos <= tagStart {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Extract tag (without #)
|
||||
tagName := line[tagStart:pos]
|
||||
|
||||
// Make a copy of the tag name
|
||||
tagCopy := make([]byte, len(tagName))
|
||||
copy(tagCopy, tagName)
|
||||
|
||||
// Advance reader
|
||||
block.Advance(pos)
|
||||
|
||||
// Create node
|
||||
node := &mast.TagNode{
|
||||
Tag: tagCopy,
|
||||
}
|
||||
|
||||
return node
|
||||
}
|
||||
@@ -0,0 +1,263 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/yuin/goldmark/parser"
|
||||
"github.com/yuin/goldmark/text"
|
||||
|
||||
mast "github.com/usememos/memos/internal/markdown/ast"
|
||||
)
|
||||
|
||||
func TestTagParser(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expectedTag string
|
||||
shouldParse bool
|
||||
}{
|
||||
{
|
||||
name: "basic tag",
|
||||
input: "#tag",
|
||||
expectedTag: "tag",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "tag with hyphen",
|
||||
input: "#work-notes",
|
||||
expectedTag: "work-notes",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "tag with ampersand",
|
||||
input: "#science&tech",
|
||||
expectedTag: "science&tech",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "tag with underscore",
|
||||
input: "#2024_plans",
|
||||
expectedTag: "2024_plans",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "numeric tag",
|
||||
input: "#123",
|
||||
expectedTag: "123",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "tag followed by space",
|
||||
input: "#tag ",
|
||||
expectedTag: "tag",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "tag followed by punctuation",
|
||||
input: "#tag.",
|
||||
expectedTag: "tag",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "tag in sentence",
|
||||
input: "#important task",
|
||||
expectedTag: "important",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "heading (##)",
|
||||
input: "## Heading",
|
||||
expectedTag: "",
|
||||
shouldParse: false,
|
||||
},
|
||||
{
|
||||
name: "space after hash",
|
||||
input: "# heading",
|
||||
expectedTag: "",
|
||||
shouldParse: false,
|
||||
},
|
||||
{
|
||||
name: "lone hash",
|
||||
input: "#",
|
||||
expectedTag: "",
|
||||
shouldParse: false,
|
||||
},
|
||||
{
|
||||
name: "hash with space",
|
||||
input: "# ",
|
||||
expectedTag: "",
|
||||
shouldParse: false,
|
||||
},
|
||||
{
|
||||
name: "special characters",
|
||||
input: "#tag@special",
|
||||
expectedTag: "tag",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "mixed case",
|
||||
input: "#WorkNotes",
|
||||
expectedTag: "WorkNotes",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "hierarchical tag with slash",
|
||||
input: "#tag1/subtag",
|
||||
expectedTag: "tag1/subtag",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "hierarchical tag with multiple levels",
|
||||
input: "#tag1/subtag/subtag2",
|
||||
expectedTag: "tag1/subtag/subtag2",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "hierarchical tag followed by space",
|
||||
input: "#work/notes ",
|
||||
expectedTag: "work/notes",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "hierarchical tag followed by punctuation",
|
||||
input: "#project/2024.",
|
||||
expectedTag: "project/2024",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "hierarchical tag with numbers and dashes",
|
||||
input: "#work-log/2024/q1",
|
||||
expectedTag: "work-log/2024/q1",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "Chinese characters",
|
||||
input: "#测试",
|
||||
expectedTag: "测试",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "Chinese tag followed by space",
|
||||
input: "#测试 some text",
|
||||
expectedTag: "测试",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "Chinese tag followed by punctuation",
|
||||
input: "#测试。",
|
||||
expectedTag: "测试",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "mixed Chinese and ASCII",
|
||||
input: "#测试test123",
|
||||
expectedTag: "测试test123",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "Japanese characters",
|
||||
input: "#テスト",
|
||||
expectedTag: "テスト",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "Korean characters",
|
||||
input: "#테스트",
|
||||
expectedTag: "테스트",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "emoji",
|
||||
input: "#test🚀",
|
||||
expectedTag: "test🚀",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "emoji with VS16",
|
||||
input: "#test👁️", // Eye + VS16
|
||||
expectedTag: "test👁️",
|
||||
shouldParse: true,
|
||||
},
|
||||
{
|
||||
name: "emoji with ZWJ sequence",
|
||||
input: "#family👨👩👧👦", // Family ZWJ sequence
|
||||
expectedTag: "family👨👩👧👦",
|
||||
shouldParse: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
p := NewTagParser()
|
||||
reader := text.NewReader([]byte(tt.input))
|
||||
ctx := parser.NewContext()
|
||||
|
||||
node := p.Parse(nil, reader, ctx)
|
||||
|
||||
if tt.shouldParse {
|
||||
require.NotNil(t, node, "Expected tag to be parsed")
|
||||
require.IsType(t, &mast.TagNode{}, node)
|
||||
|
||||
tagNode, ok := node.(*mast.TagNode)
|
||||
require.True(t, ok, "Expected node to be *mast.TagNode")
|
||||
assert.Equal(t, tt.expectedTag, string(tagNode.Tag))
|
||||
} else {
|
||||
assert.Nil(t, node, "Expected tag NOT to be parsed")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTagParser_Trigger(t *testing.T) {
|
||||
p := NewTagParser()
|
||||
triggers := p.Trigger()
|
||||
|
||||
assert.Equal(t, []byte{'#'}, triggers)
|
||||
}
|
||||
|
||||
func TestTagParser_MultipleTags(t *testing.T) {
|
||||
// Test that parser correctly handles multiple tags in sequence
|
||||
input := "#tag1 #tag2"
|
||||
|
||||
p := NewTagParser()
|
||||
reader := text.NewReader([]byte(input))
|
||||
ctx := parser.NewContext()
|
||||
|
||||
// Parse first tag
|
||||
node1 := p.Parse(nil, reader, ctx)
|
||||
require.NotNil(t, node1)
|
||||
tagNode1, ok := node1.(*mast.TagNode)
|
||||
require.True(t, ok, "Expected node1 to be *mast.TagNode")
|
||||
assert.Equal(t, "tag1", string(tagNode1.Tag))
|
||||
|
||||
// Advance past the space
|
||||
reader.Advance(1)
|
||||
|
||||
// Parse second tag
|
||||
node2 := p.Parse(nil, reader, ctx)
|
||||
require.NotNil(t, node2)
|
||||
tagNode2, ok := node2.(*mast.TagNode)
|
||||
require.True(t, ok, "Expected node2 to be *mast.TagNode")
|
||||
assert.Equal(t, "tag2", string(tagNode2.Tag))
|
||||
}
|
||||
|
||||
func TestTagNode_Kind(t *testing.T) {
|
||||
node := &mast.TagNode{
|
||||
Tag: []byte("test"),
|
||||
}
|
||||
|
||||
assert.Equal(t, mast.KindTag, node.Kind())
|
||||
}
|
||||
|
||||
func TestTagNode_Dump(t *testing.T) {
|
||||
node := &mast.TagNode{
|
||||
Tag: []byte("test"),
|
||||
}
|
||||
|
||||
// Should not panic
|
||||
assert.NotPanics(t, func() {
|
||||
node.Dump([]byte("#test"), 0)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,270 @@
|
||||
package renderer
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
gast "github.com/yuin/goldmark/ast"
|
||||
east "github.com/yuin/goldmark/extension/ast"
|
||||
|
||||
mast "github.com/usememos/memos/internal/markdown/ast"
|
||||
)
|
||||
|
||||
// MarkdownRenderer renders goldmark AST back to markdown text.
|
||||
type MarkdownRenderer struct {
|
||||
buf *bytes.Buffer
|
||||
}
|
||||
|
||||
// NewMarkdownRenderer creates a new markdown renderer.
|
||||
func NewMarkdownRenderer() *MarkdownRenderer {
|
||||
return &MarkdownRenderer{
|
||||
buf: &bytes.Buffer{},
|
||||
}
|
||||
}
|
||||
|
||||
// Render renders the AST node to markdown and returns the result.
|
||||
func (r *MarkdownRenderer) Render(node gast.Node, source []byte) string {
|
||||
r.buf.Reset()
|
||||
r.renderNode(node, source, 0)
|
||||
return r.buf.String()
|
||||
}
|
||||
|
||||
// renderNode renders a single node and its children.
|
||||
func (r *MarkdownRenderer) renderNode(node gast.Node, source []byte, depth int) {
|
||||
switch n := node.(type) {
|
||||
case *gast.Document:
|
||||
r.renderChildren(n, source, depth)
|
||||
|
||||
case *gast.Paragraph:
|
||||
r.renderChildren(n, source, depth)
|
||||
if node.NextSibling() != nil {
|
||||
r.buf.WriteString("\n\n")
|
||||
}
|
||||
|
||||
case *gast.Text:
|
||||
// Text nodes store their content as segments in the source
|
||||
segment := n.Segment
|
||||
r.buf.Write(segment.Value(source))
|
||||
if n.SoftLineBreak() {
|
||||
r.buf.WriteByte('\n')
|
||||
} else if n.HardLineBreak() {
|
||||
r.buf.WriteString(" \n")
|
||||
}
|
||||
|
||||
case *gast.CodeSpan:
|
||||
r.buf.WriteByte('`')
|
||||
r.renderChildren(n, source, depth)
|
||||
r.buf.WriteByte('`')
|
||||
|
||||
case *gast.Emphasis:
|
||||
symbol := "*"
|
||||
if n.Level == 2 {
|
||||
symbol = "**"
|
||||
}
|
||||
r.buf.WriteString(symbol)
|
||||
r.renderChildren(n, source, depth)
|
||||
r.buf.WriteString(symbol)
|
||||
|
||||
case *gast.Link:
|
||||
r.buf.WriteString("[")
|
||||
r.renderChildren(n, source, depth)
|
||||
r.buf.WriteString("](")
|
||||
r.buf.Write(n.Destination)
|
||||
if len(n.Title) > 0 {
|
||||
r.buf.WriteString(` "`)
|
||||
r.buf.Write(n.Title)
|
||||
r.buf.WriteString(`"`)
|
||||
}
|
||||
r.buf.WriteString(")")
|
||||
|
||||
case *gast.AutoLink:
|
||||
url := n.URL(source)
|
||||
if n.AutoLinkType == gast.AutoLinkEmail {
|
||||
r.buf.WriteString("<")
|
||||
r.buf.Write(url)
|
||||
r.buf.WriteString(">")
|
||||
} else {
|
||||
r.buf.Write(url)
|
||||
}
|
||||
|
||||
case *gast.Image:
|
||||
r.buf.WriteString("
|
||||
r.buf.Write(n.Destination)
|
||||
if len(n.Title) > 0 {
|
||||
r.buf.WriteString(` "`)
|
||||
r.buf.Write(n.Title)
|
||||
r.buf.WriteString(`"`)
|
||||
}
|
||||
r.buf.WriteString(")")
|
||||
|
||||
case *gast.Heading:
|
||||
r.buf.WriteString(strings.Repeat("#", n.Level))
|
||||
r.buf.WriteByte(' ')
|
||||
r.renderChildren(n, source, depth)
|
||||
if node.NextSibling() != nil {
|
||||
r.buf.WriteString("\n\n")
|
||||
}
|
||||
|
||||
case *gast.CodeBlock, *gast.FencedCodeBlock:
|
||||
r.renderCodeBlock(n, source)
|
||||
|
||||
case *gast.Blockquote:
|
||||
// Render each child line with "> " prefix
|
||||
r.renderBlockquote(n, source, depth)
|
||||
if node.NextSibling() != nil {
|
||||
r.buf.WriteString("\n\n")
|
||||
}
|
||||
|
||||
case *gast.List:
|
||||
r.renderChildren(n, source, depth)
|
||||
if node.NextSibling() != nil {
|
||||
r.buf.WriteString("\n\n")
|
||||
}
|
||||
|
||||
case *gast.ListItem:
|
||||
r.renderListItem(n, source, depth)
|
||||
|
||||
case *gast.ThematicBreak:
|
||||
r.buf.WriteString("---")
|
||||
if node.NextSibling() != nil {
|
||||
r.buf.WriteString("\n\n")
|
||||
}
|
||||
|
||||
case *east.Strikethrough:
|
||||
r.buf.WriteString("~~")
|
||||
r.renderChildren(n, source, depth)
|
||||
r.buf.WriteString("~~")
|
||||
|
||||
case *east.TaskCheckBox:
|
||||
if n.IsChecked {
|
||||
r.buf.WriteString("[x] ")
|
||||
} else {
|
||||
r.buf.WriteString("[ ] ")
|
||||
}
|
||||
|
||||
case *east.Table:
|
||||
r.renderTable(n, source)
|
||||
if node.NextSibling() != nil {
|
||||
r.buf.WriteString("\n\n")
|
||||
}
|
||||
|
||||
// Custom Memos nodes
|
||||
case *mast.TagNode:
|
||||
r.buf.WriteByte('#')
|
||||
r.buf.Write(n.Tag)
|
||||
|
||||
case *mast.MentionNode:
|
||||
r.buf.WriteByte('@')
|
||||
r.buf.Write(n.Username)
|
||||
|
||||
default:
|
||||
// For unknown nodes, try to render children
|
||||
r.renderChildren(n, source, depth)
|
||||
}
|
||||
}
|
||||
|
||||
// renderChildren renders all children of a node.
|
||||
func (r *MarkdownRenderer) renderChildren(node gast.Node, source []byte, depth int) {
|
||||
child := node.FirstChild()
|
||||
for child != nil {
|
||||
r.renderNode(child, source, depth+1)
|
||||
child = child.NextSibling()
|
||||
}
|
||||
}
|
||||
|
||||
// renderCodeBlock renders a code block.
|
||||
func (r *MarkdownRenderer) renderCodeBlock(node gast.Node, source []byte) {
|
||||
if fenced, ok := node.(*gast.FencedCodeBlock); ok {
|
||||
// Fenced code block with language
|
||||
r.buf.WriteString("```")
|
||||
if lang := fenced.Language(source); len(lang) > 0 {
|
||||
r.buf.Write(lang)
|
||||
}
|
||||
r.buf.WriteByte('\n')
|
||||
|
||||
// Write all lines
|
||||
lines := fenced.Lines()
|
||||
for i := 0; i < lines.Len(); i++ {
|
||||
line := lines.At(i)
|
||||
r.buf.Write(line.Value(source))
|
||||
}
|
||||
|
||||
r.buf.WriteString("```")
|
||||
if node.NextSibling() != nil {
|
||||
r.buf.WriteString("\n\n")
|
||||
}
|
||||
} else if codeBlock, ok := node.(*gast.CodeBlock); ok {
|
||||
// Indented code block
|
||||
lines := codeBlock.Lines()
|
||||
for i := 0; i < lines.Len(); i++ {
|
||||
line := lines.At(i)
|
||||
r.buf.WriteString(" ")
|
||||
r.buf.Write(line.Value(source))
|
||||
}
|
||||
if node.NextSibling() != nil {
|
||||
r.buf.WriteString("\n\n")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// renderBlockquote renders a blockquote with "> " prefix.
|
||||
func (r *MarkdownRenderer) renderBlockquote(node *gast.Blockquote, source []byte, depth int) {
|
||||
// Create a temporary buffer for the blockquote content
|
||||
tempBuf := &bytes.Buffer{}
|
||||
tempRenderer := &MarkdownRenderer{buf: tempBuf}
|
||||
tempRenderer.renderChildren(node, source, depth)
|
||||
|
||||
// Add "> " prefix to each line
|
||||
content := tempBuf.String()
|
||||
lines := strings.Split(strings.TrimRight(content, "\n"), "\n")
|
||||
for i, line := range lines {
|
||||
r.buf.WriteString("> ")
|
||||
r.buf.WriteString(line)
|
||||
if i < len(lines)-1 {
|
||||
r.buf.WriteByte('\n')
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// renderListItem renders a list item with proper indentation and markers.
|
||||
func (r *MarkdownRenderer) renderListItem(node *gast.ListItem, source []byte, depth int) {
|
||||
parent := node.Parent()
|
||||
list, ok := parent.(*gast.List)
|
||||
if !ok {
|
||||
r.renderChildren(node, source, depth)
|
||||
return
|
||||
}
|
||||
|
||||
// Add indentation only for nested lists
|
||||
// Document=0, List=1, ListItem=2 (no indent), nested ListItem=3+ (indent)
|
||||
if depth > 2 {
|
||||
indent := strings.Repeat(" ", depth-2)
|
||||
r.buf.WriteString(indent)
|
||||
}
|
||||
|
||||
// Add list marker
|
||||
if list.IsOrdered() {
|
||||
fmt.Fprintf(r.buf, "%d. ", list.Start)
|
||||
list.Start++ // Increment for next item
|
||||
} else {
|
||||
r.buf.WriteString("- ")
|
||||
}
|
||||
|
||||
// Render content
|
||||
r.renderChildren(node, source, depth)
|
||||
|
||||
// Add newline if there's a next sibling
|
||||
if node.NextSibling() != nil {
|
||||
r.buf.WriteByte('\n')
|
||||
}
|
||||
}
|
||||
|
||||
// renderTable renders a table in markdown format.
|
||||
func (r *MarkdownRenderer) renderTable(table *east.Table, source []byte) {
|
||||
// This is a simplified table renderer
|
||||
// A full implementation would need to handle alignment, etc.
|
||||
r.renderChildren(table, source, 0)
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
package renderer
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/yuin/goldmark"
|
||||
"github.com/yuin/goldmark/extension"
|
||||
"github.com/yuin/goldmark/parser"
|
||||
"github.com/yuin/goldmark/text"
|
||||
|
||||
"github.com/usememos/memos/internal/markdown/extensions"
|
||||
)
|
||||
|
||||
func TestMarkdownRenderer(t *testing.T) {
|
||||
// Create goldmark instance with all extensions
|
||||
md := goldmark.New(
|
||||
goldmark.WithExtensions(
|
||||
extension.GFM,
|
||||
extensions.TagExtension,
|
||||
),
|
||||
goldmark.WithParserOptions(
|
||||
parser.WithAutoHeadingID(),
|
||||
),
|
||||
)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "simple text",
|
||||
input: "Hello world",
|
||||
expected: "Hello world",
|
||||
},
|
||||
{
|
||||
name: "paragraph with newlines",
|
||||
input: "First paragraph\n\nSecond paragraph",
|
||||
expected: "First paragraph\n\nSecond paragraph",
|
||||
},
|
||||
{
|
||||
name: "emphasis",
|
||||
input: "This is *italic* and **bold** text",
|
||||
expected: "This is *italic* and **bold** text",
|
||||
},
|
||||
{
|
||||
name: "headings",
|
||||
input: "# Heading 1\n\n## Heading 2\n\n### Heading 3",
|
||||
expected: "# Heading 1\n\n## Heading 2\n\n### Heading 3",
|
||||
},
|
||||
{
|
||||
name: "link",
|
||||
input: "Check [this link](https://example.com)",
|
||||
expected: "Check [this link](https://example.com)",
|
||||
},
|
||||
{
|
||||
name: "image",
|
||||
input: "",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "code inline",
|
||||
input: "This is `inline code` here",
|
||||
expected: "This is `inline code` here",
|
||||
},
|
||||
{
|
||||
name: "code block fenced",
|
||||
input: "```go\nfunc main() {\n}\n```",
|
||||
expected: "```go\nfunc main() {\n}\n```",
|
||||
},
|
||||
{
|
||||
name: "unordered list",
|
||||
input: "- Item 1\n- Item 2\n- Item 3",
|
||||
expected: "- Item 1\n- Item 2\n- Item 3",
|
||||
},
|
||||
{
|
||||
name: "ordered list",
|
||||
input: "1. First\n2. Second\n3. Third",
|
||||
expected: "1. First\n2. Second\n3. Third",
|
||||
},
|
||||
{
|
||||
name: "blockquote",
|
||||
input: "> This is a quote\n> Second line",
|
||||
expected: "> This is a quote\n> Second line",
|
||||
},
|
||||
{
|
||||
name: "horizontal rule",
|
||||
input: "Text before\n\n---\n\nText after",
|
||||
expected: "Text before\n\n---\n\nText after",
|
||||
},
|
||||
{
|
||||
name: "strikethrough",
|
||||
input: "This is ~~deleted~~ text",
|
||||
expected: "This is ~~deleted~~ text",
|
||||
},
|
||||
{
|
||||
name: "task list",
|
||||
input: "- [x] Completed task\n- [ ] Incomplete task",
|
||||
expected: "- [x] Completed task\n- [ ] Incomplete task",
|
||||
},
|
||||
{
|
||||
name: "tag",
|
||||
input: "This has #tag in it",
|
||||
expected: "This has #tag in it",
|
||||
},
|
||||
{
|
||||
name: "multiple tags",
|
||||
input: "#work #important meeting notes",
|
||||
expected: "#work #important meeting notes",
|
||||
},
|
||||
{
|
||||
name: "complex mixed content",
|
||||
input: "# Meeting Notes\n\n**Date**: 2024-01-01\n\n## Attendees\n- Alice\n- Bob\n\n## Discussion\n\nWe discussed #project status.\n\n```python\nprint('hello')\n```",
|
||||
expected: "# Meeting Notes\n\n**Date**: 2024-01-01\n\n## Attendees\n\n- Alice\n- Bob\n\n## Discussion\n\nWe discussed #project status.\n\n```python\nprint('hello')\n```",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Parse the input
|
||||
source := []byte(tt.input)
|
||||
reader := text.NewReader(source)
|
||||
doc := md.Parser().Parse(reader)
|
||||
require.NotNil(t, doc)
|
||||
|
||||
// Render back to markdown
|
||||
renderer := NewMarkdownRenderer()
|
||||
result := renderer.Render(doc, source)
|
||||
|
||||
// For debugging
|
||||
if result != tt.expected {
|
||||
t.Logf("Input: %q", tt.input)
|
||||
t.Logf("Expected: %q", tt.expected)
|
||||
t.Logf("Got: %q", result)
|
||||
}
|
||||
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkdownRendererPreservesStructure(t *testing.T) {
|
||||
// Test that parsing and rendering preserves structure
|
||||
md := goldmark.New(
|
||||
goldmark.WithExtensions(
|
||||
extension.GFM,
|
||||
extensions.TagExtension,
|
||||
),
|
||||
)
|
||||
|
||||
inputs := []string{
|
||||
"# Title\n\nParagraph",
|
||||
"**Bold** and *italic*",
|
||||
"- List\n- Items",
|
||||
"#tag #another",
|
||||
"> Quote",
|
||||
}
|
||||
|
||||
renderer := NewMarkdownRenderer()
|
||||
|
||||
for _, input := range inputs {
|
||||
t.Run(input, func(t *testing.T) {
|
||||
source := []byte(input)
|
||||
reader := text.NewReader(source)
|
||||
doc := md.Parser().Parse(reader)
|
||||
|
||||
result := renderer.Render(doc, source)
|
||||
|
||||
// The result should be structurally similar
|
||||
// (may have minor formatting differences)
|
||||
assert.NotEmpty(t, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package motionphoto
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"regexp"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
type Detection struct {
|
||||
VideoStart int
|
||||
PresentationTimestampUs int64
|
||||
}
|
||||
|
||||
var (
|
||||
motionPhotoMarkerRegex = regexp.MustCompile(`(?i)(?:Camera:MotionPhoto|GCamera:MotionPhoto|MicroVideo)["'=:\s>]+1`)
|
||||
presentationRegex = regexp.MustCompile(`(?i)(?:Camera:MotionPhotoPresentationTimestampUs|GCamera:MotionPhotoPresentationTimestampUs)["'=:\s>]+(-?\d+)`)
|
||||
microVideoOffsetRegex = regexp.MustCompile(`(?i)(?:Camera:MicroVideoOffset|GCamera:MicroVideoOffset)["'=:\s>]+(\d+)`)
|
||||
)
|
||||
|
||||
const maxMetadataScanBytes = 256 * 1024
|
||||
|
||||
func DetectJPEG(blob []byte) *Detection {
|
||||
if len(blob) < 16 || !bytes.HasPrefix(blob, []byte{0xFF, 0xD8}) {
|
||||
return nil
|
||||
}
|
||||
|
||||
text := string(blob[:min(len(blob), maxMetadataScanBytes)])
|
||||
if !motionPhotoMarkerRegex.MatchString(text) {
|
||||
return nil
|
||||
}
|
||||
|
||||
videoStart := detectVideoStart(blob, text)
|
||||
if videoStart < 0 || videoStart >= len(blob) {
|
||||
return nil
|
||||
}
|
||||
|
||||
return &Detection{
|
||||
VideoStart: videoStart,
|
||||
PresentationTimestampUs: parsePresentationTimestampUs(text),
|
||||
}
|
||||
}
|
||||
|
||||
func ExtractVideo(blob []byte) ([]byte, *Detection) {
|
||||
detection := DetectJPEG(blob)
|
||||
if detection == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
videoBlob := blob[detection.VideoStart:]
|
||||
if !looksLikeMP4(videoBlob) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
return videoBlob, detection
|
||||
}
|
||||
|
||||
func detectVideoStart(blob []byte, text string) int {
|
||||
if matches := microVideoOffsetRegex.FindStringSubmatch(text); len(matches) == 2 {
|
||||
if offset, err := strconv.Atoi(matches[1]); err == nil && offset > 0 && offset < len(blob) {
|
||||
start := len(blob) - offset
|
||||
if looksLikeMP4(blob[start:]) {
|
||||
return start
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return findEmbeddedMP4Start(blob)
|
||||
}
|
||||
|
||||
func parsePresentationTimestampUs(text string) int64 {
|
||||
matches := presentationRegex.FindStringSubmatch(text)
|
||||
if len(matches) != 2 {
|
||||
return 0
|
||||
}
|
||||
|
||||
value, err := strconv.ParseInt(matches[1], 10, 64)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func findEmbeddedMP4Start(blob []byte) int {
|
||||
searchFrom := len(blob)
|
||||
for searchFrom > 8 {
|
||||
index := bytes.LastIndex(blob[:searchFrom], []byte("ftyp"))
|
||||
if index < 4 {
|
||||
return -1
|
||||
}
|
||||
|
||||
start := index - 4
|
||||
if looksLikeMP4(blob[start:]) {
|
||||
return start
|
||||
}
|
||||
|
||||
searchFrom = index - 1
|
||||
}
|
||||
|
||||
return -1
|
||||
}
|
||||
|
||||
func looksLikeMP4(blob []byte) bool {
|
||||
if len(blob) < 12 || !bytes.Equal(blob[4:8], []byte("ftyp")) {
|
||||
return false
|
||||
}
|
||||
|
||||
size := binary.BigEndian.Uint32(blob[:4])
|
||||
return size == 1 || size >= 8
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
package motionphoto
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/usememos/memos/internal/testutil"
|
||||
)
|
||||
|
||||
func TestDetectJPEG(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
blob := testutil.BuildMotionPhotoJPEG()
|
||||
detection := DetectJPEG(blob)
|
||||
require.NotNil(t, detection)
|
||||
require.Positive(t, detection.VideoStart)
|
||||
require.EqualValues(t, 123456, detection.PresentationTimestampUs)
|
||||
|
||||
videoBlob, extracted := ExtractVideo(blob)
|
||||
require.NotNil(t, extracted)
|
||||
require.True(t, bytes.Equal(videoBlob[:4], []byte{0x00, 0x00, 0x00, 0x10}))
|
||||
require.Equal(t, []byte("ftyp"), videoBlob[4:8])
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package profile
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
// Profile is the configuration to start main server.
|
||||
type Profile struct {
|
||||
// Demo indicates if the server is in demo mode
|
||||
Demo bool
|
||||
// Addr is the binding address for server
|
||||
Addr string
|
||||
// Port is the binding port for server
|
||||
Port int
|
||||
// UNIXSock is the IPC binding path. Overrides Addr and Port
|
||||
UNIXSock string
|
||||
// Data is the data directory
|
||||
Data string
|
||||
// DSN points to where memos stores its own data
|
||||
DSN string
|
||||
// Driver is the database driver
|
||||
// sqlite, mysql
|
||||
Driver string
|
||||
// Version is the current version of server
|
||||
Version string
|
||||
// Commit is the current build commit of server
|
||||
Commit string
|
||||
// InstanceURL is the url of your memos instance.
|
||||
InstanceURL string
|
||||
}
|
||||
|
||||
// AllowAnonymous reports whether unauthenticated visitors may access the instance.
|
||||
//
|
||||
// Anonymous access is enabled only when an InstanceURL is configured. An instance
|
||||
// with no InstanceURL set is treated as private: anonymous callers are limited to
|
||||
// the auth-bootstrap endpoints (sign-in, share links, etc.) and the web UI redirects
|
||||
// them to the sign-in page instead of the public Explore view. Authenticated callers
|
||||
// (session, access token, or personal access token) are never affected.
|
||||
func (p *Profile) AllowAnonymous() bool {
|
||||
return strings.TrimSpace(p.InstanceURL) != ""
|
||||
}
|
||||
|
||||
func checkDataDir(dataDir string) (string, error) {
|
||||
// Convert to absolute path if relative path is supplied.
|
||||
if !filepath.IsAbs(dataDir) {
|
||||
// Use current working directory, not the binary's directory
|
||||
// This ensures we use the actual working directory where the process runs
|
||||
absDir, err := filepath.Abs(dataDir)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
dataDir = absDir
|
||||
}
|
||||
|
||||
// Trim trailing \ or / in case user supplies
|
||||
dataDir = strings.TrimRight(dataDir, "\\/")
|
||||
if _, err := os.Stat(dataDir); err != nil {
|
||||
return "", errors.Wrapf(err, "unable to access data folder %s", dataDir)
|
||||
}
|
||||
return dataDir, nil
|
||||
}
|
||||
|
||||
func (p *Profile) Validate() error {
|
||||
// Set default data directory if not specified
|
||||
if p.Data == "" {
|
||||
if runtime.GOOS == "windows" {
|
||||
p.Data = filepath.Join(os.Getenv("ProgramData"), "memos")
|
||||
} else {
|
||||
// On Linux/macOS, check if /var/opt/memos exists and is writable (Docker scenario)
|
||||
if info, err := os.Stat("/var/opt/memos"); err == nil && info.IsDir() {
|
||||
// Check if we can write to this directory
|
||||
testFile := filepath.Join("/var/opt/memos", ".write-test")
|
||||
if err := os.WriteFile(testFile, []byte("test"), 0600); err == nil {
|
||||
os.Remove(testFile)
|
||||
p.Data = "/var/opt/memos"
|
||||
} else {
|
||||
// /var/opt/memos exists but is not writable, use current directory
|
||||
slog.Warn("/var/opt/memos is not writable, using current directory")
|
||||
p.Data = "."
|
||||
}
|
||||
} else {
|
||||
// /var/opt/memos doesn't exist, use current directory (local development)
|
||||
p.Data = "."
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Create data directory if it doesn't exist
|
||||
if _, err := os.Stat(p.Data); os.IsNotExist(err) {
|
||||
if err := os.MkdirAll(p.Data, 0770); err != nil {
|
||||
slog.Error("failed to create data directory", slog.String("data", p.Data), slog.String("error", err.Error()))
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
dataDir, err := checkDataDir(p.Data)
|
||||
if err != nil {
|
||||
slog.Error("failed to check dsn", slog.String("data", dataDir), slog.String("error", err.Error()))
|
||||
return err
|
||||
}
|
||||
|
||||
p.Data = dataDir
|
||||
if p.Driver == "sqlite" && p.DSN == "" {
|
||||
mode := "prod"
|
||||
if p.Demo {
|
||||
mode = "demo"
|
||||
}
|
||||
dbFile := fmt.Sprintf("memos_%s.db", mode)
|
||||
p.DSN = filepath.Join(dataDir, dbFile)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package profile
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestAllowAnonymous(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
url string
|
||||
want bool
|
||||
}{
|
||||
{"empty is private", "", false},
|
||||
{"whitespace only is private", " ", false},
|
||||
{"configured url is public", "https://memos.example.com", true},
|
||||
{"configured url with padding is public", " https://memos.example.com ", true},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
p := &Profile{InstanceURL: c.url}
|
||||
if got := p.AllowAnonymous(); got != c.want {
|
||||
t.Fatalf("AllowAnonymous() with InstanceURL=%q = %v, want %v", c.url, got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,367 @@
|
||||
# Scheduler Plugin
|
||||
|
||||
A production-ready, GitHub Actions-inspired cron job scheduler for Go.
|
||||
|
||||
## Features
|
||||
|
||||
- **Standard Cron Syntax**: Supports both 5-field and 6-field (with seconds) cron expressions
|
||||
- **Timezone-Aware**: Explicit timezone handling to avoid DST surprises
|
||||
- **Middleware Pattern**: Composable job wrappers for logging, metrics, panic recovery, timeouts
|
||||
- **Graceful Shutdown**: Jobs complete cleanly or cancel when context expires
|
||||
- **Zero Dependencies**: Core functionality uses only the standard library
|
||||
- **Type-Safe**: Strong typing with clear error messages
|
||||
- **Well-Tested**: Comprehensive test coverage
|
||||
|
||||
## Installation
|
||||
|
||||
This package is included with Memos. No separate installation required.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```go
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/usememos/memos/internal/scheduler"
|
||||
)
|
||||
|
||||
func main() {
|
||||
s := scheduler.New()
|
||||
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "daily-cleanup",
|
||||
Schedule: "0 2 * * *", // 2 AM daily
|
||||
Handler: func(ctx context.Context) error {
|
||||
fmt.Println("Running cleanup...")
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
s.Start()
|
||||
defer s.Stop(context.Background())
|
||||
|
||||
// Keep running...
|
||||
select {}
|
||||
}
|
||||
```
|
||||
|
||||
## Cron Expression Format
|
||||
|
||||
### 5-Field Format (Standard)
|
||||
```
|
||||
┌───────────── minute (0 - 59)
|
||||
│ ┌───────────── hour (0 - 23)
|
||||
│ │ ┌───────────── day of month (1 - 31)
|
||||
│ │ │ ┌───────────── month (1 - 12)
|
||||
│ │ │ │ ┌───────────── day of week (0 - 7) (Sunday = 0 or 7)
|
||||
│ │ │ │ │
|
||||
* * * * *
|
||||
```
|
||||
|
||||
### 6-Field Format (With Seconds)
|
||||
```
|
||||
┌───────────── second (0 - 59)
|
||||
│ ┌───────────── minute (0 - 59)
|
||||
│ │ ┌───────────── hour (0 - 23)
|
||||
│ │ │ ┌───────────── day of month (1 - 31)
|
||||
│ │ │ │ ┌───────────── month (1 - 12)
|
||||
│ │ │ │ │ ┌───────────── day of week (0 - 7)
|
||||
│ │ │ │ │ │
|
||||
* * * * * *
|
||||
```
|
||||
|
||||
### Special Characters
|
||||
|
||||
- `*` - Any value (every minute, every hour, etc.)
|
||||
- `,` - List of values: `1,15,30` (1st, 15th, and 30th)
|
||||
- `-` - Range: `9-17` (9 AM through 5 PM)
|
||||
- `/` - Step: `*/15` (every 15 units)
|
||||
|
||||
### Common Examples
|
||||
|
||||
| Schedule | Description |
|
||||
|----------|-------------|
|
||||
| `* * * * *` | Every minute |
|
||||
| `0 * * * *` | Every hour |
|
||||
| `0 0 * * *` | Daily at midnight |
|
||||
| `0 9 * * 1-5` | Weekdays at 9 AM |
|
||||
| `*/15 * * * *` | Every 15 minutes |
|
||||
| `0 0 1 * *` | First day of every month |
|
||||
| `0 0 * * 0` | Every Sunday at midnight |
|
||||
| `30 14 * * *` | Every day at 2:30 PM |
|
||||
|
||||
## Timezone Support
|
||||
|
||||
```go
|
||||
// Global timezone for all jobs
|
||||
s := scheduler.New(
|
||||
scheduler.WithTimezone("America/New_York"),
|
||||
)
|
||||
|
||||
// Per-job timezone (overrides global)
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "tokyo-report",
|
||||
Schedule: "0 9 * * *", // 9 AM Tokyo time
|
||||
Timezone: "Asia/Tokyo",
|
||||
Handler: func(ctx context.Context) error {
|
||||
// Runs at 9 AM in Tokyo
|
||||
return nil
|
||||
},
|
||||
})
|
||||
```
|
||||
|
||||
**Important**: Always use IANA timezone names (`America/New_York`, not `EST`).
|
||||
|
||||
## Middleware
|
||||
|
||||
Middleware wraps job handlers to add cross-cutting behavior. Multiple middleware can be chained together.
|
||||
|
||||
### Built-in Middleware
|
||||
|
||||
#### Recovery (Panic Handling)
|
||||
|
||||
```go
|
||||
s := scheduler.New(
|
||||
scheduler.WithMiddleware(
|
||||
scheduler.Recovery(func(jobName string, r interface{}) {
|
||||
log.Printf("Job %s panicked: %v", jobName, r)
|
||||
}),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
#### Logging
|
||||
|
||||
```go
|
||||
type Logger interface {
|
||||
Info(msg string, args ...interface{})
|
||||
Error(msg string, args ...interface{})
|
||||
}
|
||||
|
||||
s := scheduler.New(
|
||||
scheduler.WithMiddleware(
|
||||
scheduler.Logging(myLogger),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
#### Timeout
|
||||
|
||||
```go
|
||||
s := scheduler.New(
|
||||
scheduler.WithMiddleware(
|
||||
scheduler.Timeout(5 * time.Minute),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
### Combining Middleware
|
||||
|
||||
```go
|
||||
s := scheduler.New(
|
||||
scheduler.WithMiddleware(
|
||||
scheduler.Recovery(panicHandler),
|
||||
scheduler.Logging(logger),
|
||||
scheduler.Timeout(10 * time.Minute),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
**Order matters**: Middleware are applied left-to-right. In the example above:
|
||||
1. Recovery (outermost) catches panics from everything
|
||||
2. Logging logs the execution
|
||||
3. Timeout (innermost) wraps the actual handler
|
||||
|
||||
### Custom Middleware
|
||||
|
||||
```go
|
||||
func Metrics(recorder MetricsRecorder) scheduler.Middleware {
|
||||
return func(next scheduler.JobHandler) scheduler.JobHandler {
|
||||
return func(ctx context.Context) error {
|
||||
start := time.Now()
|
||||
err := next(ctx)
|
||||
duration := time.Since(start)
|
||||
|
||||
jobName := scheduler.GetJobName(ctx)
|
||||
recorder.Record(jobName, duration, err)
|
||||
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Graceful Shutdown
|
||||
|
||||
Always use `Stop()` with a context to allow jobs to finish cleanly:
|
||||
|
||||
```go
|
||||
// Give jobs up to 30 seconds to complete
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.Stop(ctx); err != nil {
|
||||
log.Printf("Shutdown error: %v", err)
|
||||
}
|
||||
```
|
||||
|
||||
Jobs should respect context cancellation:
|
||||
|
||||
```go
|
||||
Handler: func(ctx context.Context) error {
|
||||
for i := 0; i < 100; i++ {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err() // Canceled
|
||||
default:
|
||||
// Do work
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
```
|
||||
|
||||
## Best Practices
|
||||
|
||||
### 1. Always Name Your Jobs
|
||||
|
||||
Names are used for logging, metrics, and debugging:
|
||||
|
||||
```go
|
||||
Name: "user-cleanup-job" // Good
|
||||
Name: "job1" // Bad
|
||||
```
|
||||
|
||||
### 2. Add Descriptions and Tags
|
||||
|
||||
```go
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "stale-session-cleanup",
|
||||
Description: "Removes user sessions older than 30 days",
|
||||
Tags: []string{"maintenance", "security"},
|
||||
Schedule: "0 3 * * *",
|
||||
Handler: cleanupSessions,
|
||||
})
|
||||
```
|
||||
|
||||
### 3. Use Appropriate Middleware
|
||||
|
||||
Always include Recovery and Logging in production:
|
||||
|
||||
```go
|
||||
scheduler.New(
|
||||
scheduler.WithMiddleware(
|
||||
scheduler.Recovery(logPanic),
|
||||
scheduler.Logging(logger),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
### 4. Avoid Scheduling Exactly on the Hour
|
||||
|
||||
Many systems schedule jobs at `:00`, causing load spikes. Stagger your jobs:
|
||||
|
||||
```go
|
||||
"5 2 * * *" // 2:05 AM (good)
|
||||
"0 2 * * *" // 2:00 AM (often overloaded)
|
||||
```
|
||||
|
||||
### 5. Make Jobs Idempotent
|
||||
|
||||
Jobs may run multiple times (crash recovery, etc.). Design them to be safely re-runnable:
|
||||
|
||||
```go
|
||||
Handler: func(ctx context.Context) error {
|
||||
// Use unique constraint or check-before-insert
|
||||
db.Exec("INSERT IGNORE INTO processed_items ...")
|
||||
return nil
|
||||
}
|
||||
```
|
||||
|
||||
### 6. Handle Timezones Explicitly
|
||||
|
||||
Always specify timezone for business-hour jobs:
|
||||
|
||||
```go
|
||||
Timezone: "America/New_York" // Good
|
||||
// Timezone: "" // Bad (defaults to UTC)
|
||||
```
|
||||
|
||||
### 7. Test Your Cron Expressions
|
||||
|
||||
Use a cron expression calculator before deploying:
|
||||
- [crontab.guru](https://crontab.guru/)
|
||||
- Write unit tests with the parser
|
||||
|
||||
## Testing Jobs
|
||||
|
||||
Test job handlers independently of the scheduler:
|
||||
|
||||
```go
|
||||
func TestCleanupJob(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
err := cleanupHandler(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("cleanup failed: %v", err)
|
||||
}
|
||||
|
||||
// Verify cleanup occurred
|
||||
}
|
||||
```
|
||||
|
||||
Test schedule parsing:
|
||||
|
||||
```go
|
||||
func TestScheduleParsing(t *testing.T) {
|
||||
job := &scheduler.Job{
|
||||
Name: "test",
|
||||
Schedule: "0 2 * * *",
|
||||
Handler: func(ctx context.Context) error { return nil },
|
||||
}
|
||||
|
||||
if err := job.Validate(); err != nil {
|
||||
t.Fatalf("invalid job: %v", err)
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Comparison to Other Solutions
|
||||
|
||||
| Feature | scheduler | robfig/cron | github.com/go-co-op/gocron |
|
||||
|---------|-----------|-------------|----------------------------|
|
||||
| Standard cron syntax | ✅ | ✅ | ✅ |
|
||||
| Seconds support | ✅ | ✅ | ✅ |
|
||||
| Timezone support | ✅ | ✅ | ✅ |
|
||||
| Middleware pattern | ✅ | ⚠️ (basic) | ❌ |
|
||||
| Graceful shutdown | ✅ | ⚠️ (basic) | ✅ |
|
||||
| Zero dependencies | ✅ | ❌ | ❌ |
|
||||
| Job metadata | ✅ | ❌ | ⚠️ (limited) |
|
||||
|
||||
## API Reference
|
||||
|
||||
See [example_test.go](./example_test.go) for comprehensive examples.
|
||||
|
||||
### Core Types
|
||||
|
||||
- `Scheduler` - Manages scheduled jobs
|
||||
- `Job` - Job definition with schedule and handler
|
||||
- `Middleware` - Function that wraps job handlers
|
||||
|
||||
### Functions
|
||||
|
||||
- `New(opts ...Option) *Scheduler` - Create new scheduler
|
||||
- `WithTimezone(tz string) Option` - Set default timezone
|
||||
- `WithMiddleware(mw ...Middleware) Option` - Add middleware
|
||||
|
||||
### Methods
|
||||
|
||||
- `Register(job *Job) error` - Add job to scheduler
|
||||
- `Start() error` - Begin executing jobs
|
||||
- `Stop(ctx context.Context) error` - Graceful shutdown
|
||||
|
||||
## License
|
||||
|
||||
This package is part of the Memos project and shares its license.
|
||||
@@ -0,0 +1,35 @@
|
||||
// Package scheduler provides a GitHub Actions-inspired cron job scheduler.
|
||||
//
|
||||
// Features:
|
||||
// - Standard cron expression syntax (5-field and 6-field formats)
|
||||
// - Timezone-aware scheduling
|
||||
// - Middleware pattern for cross-cutting concerns (logging, metrics, recovery)
|
||||
// - Graceful shutdown with context cancellation
|
||||
// - Zero external dependencies
|
||||
//
|
||||
// Basic usage:
|
||||
//
|
||||
// s := scheduler.New()
|
||||
//
|
||||
// s.Register(&scheduler.Job{
|
||||
// Name: "daily-cleanup",
|
||||
// Schedule: "0 2 * * *", // 2 AM daily
|
||||
// Handler: func(ctx context.Context) error {
|
||||
// // Your cleanup logic here
|
||||
// return nil
|
||||
// },
|
||||
// })
|
||||
//
|
||||
// s.Start()
|
||||
// defer s.Stop(context.Background())
|
||||
//
|
||||
// With middleware:
|
||||
//
|
||||
// s := scheduler.New(
|
||||
// scheduler.WithTimezone("America/New_York"),
|
||||
// scheduler.WithMiddleware(
|
||||
// scheduler.Recovery(),
|
||||
// scheduler.Logging(),
|
||||
// ),
|
||||
// )
|
||||
package scheduler
|
||||
@@ -0,0 +1,165 @@
|
||||
package scheduler_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/usememos/memos/internal/scheduler"
|
||||
)
|
||||
|
||||
// Example demonstrates basic scheduler usage.
|
||||
func Example_basic() {
|
||||
s := scheduler.New()
|
||||
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "hello",
|
||||
Schedule: "*/5 * * * *", // Every 5 minutes
|
||||
Description: "Say hello",
|
||||
Handler: func(_ context.Context) error {
|
||||
fmt.Println("Hello from scheduler!")
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
s.Start()
|
||||
defer s.Stop(context.Background())
|
||||
|
||||
// Scheduler runs in background
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
|
||||
// Example demonstrates timezone-aware scheduling.
|
||||
func Example_timezone() {
|
||||
s := scheduler.New(
|
||||
scheduler.WithTimezone("America/New_York"),
|
||||
)
|
||||
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "daily-report",
|
||||
Schedule: "0 9 * * *", // 9 AM in New York
|
||||
Handler: func(_ context.Context) error {
|
||||
fmt.Println("Generating daily report...")
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
s.Start()
|
||||
defer s.Stop(context.Background())
|
||||
}
|
||||
|
||||
// Example demonstrates middleware usage.
|
||||
func Example_middleware() {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
|
||||
s := scheduler.New(
|
||||
scheduler.WithMiddleware(
|
||||
scheduler.Recovery(func(jobName string, r interface{}) {
|
||||
logger.Error("Job panicked", "job", jobName, "panic", r)
|
||||
}),
|
||||
scheduler.Logging(&slogAdapter{logger}),
|
||||
scheduler.Timeout(5*time.Minute),
|
||||
),
|
||||
)
|
||||
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "data-sync",
|
||||
Schedule: "0 */2 * * *", // Every 2 hours
|
||||
Handler: func(_ context.Context) error {
|
||||
// Your sync logic here
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
s.Start()
|
||||
defer s.Stop(context.Background())
|
||||
}
|
||||
|
||||
// slogAdapter adapts slog.Logger to scheduler.Logger interface.
|
||||
type slogAdapter struct {
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
func (a *slogAdapter) Info(msg string, args ...interface{}) {
|
||||
a.logger.Info(msg, args...)
|
||||
}
|
||||
|
||||
func (a *slogAdapter) Error(msg string, args ...interface{}) {
|
||||
a.logger.Error(msg, args...)
|
||||
}
|
||||
|
||||
// Example demonstrates multiple jobs with different schedules.
|
||||
func Example_multipleJobs() {
|
||||
s := scheduler.New()
|
||||
|
||||
// Cleanup old data every night at 2 AM
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "cleanup",
|
||||
Schedule: "0 2 * * *",
|
||||
Tags: []string{"maintenance"},
|
||||
Handler: func(_ context.Context) error {
|
||||
fmt.Println("Cleaning up old data...")
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
// Health check every 5 minutes
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "health-check",
|
||||
Schedule: "*/5 * * * *",
|
||||
Tags: []string{"monitoring"},
|
||||
Handler: func(_ context.Context) error {
|
||||
fmt.Println("Running health check...")
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
// Weekly backup on Sundays at 1 AM
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "weekly-backup",
|
||||
Schedule: "0 1 * * 0",
|
||||
Tags: []string{"backup"},
|
||||
Handler: func(_ context.Context) error {
|
||||
fmt.Println("Creating weekly backup...")
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
s.Start()
|
||||
defer s.Stop(context.Background())
|
||||
}
|
||||
|
||||
// Example demonstrates graceful shutdown with timeout.
|
||||
func Example_gracefulShutdown() {
|
||||
s := scheduler.New()
|
||||
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "long-running",
|
||||
Schedule: "* * * * *",
|
||||
Handler: func(ctx context.Context) error {
|
||||
select {
|
||||
case <-time.After(30 * time.Second):
|
||||
fmt.Println("Job completed")
|
||||
case <-ctx.Done():
|
||||
fmt.Println("Job canceled, cleaning up...")
|
||||
return ctx.Err()
|
||||
}
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
s.Start()
|
||||
|
||||
// Simulate shutdown signal
|
||||
time.Sleep(5 * time.Second)
|
||||
|
||||
// Give jobs 10 seconds to finish
|
||||
shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.Stop(shutdownCtx); err != nil {
|
||||
fmt.Printf("Shutdown error: %v\n", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,393 @@
|
||||
package scheduler_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/usememos/memos/internal/scheduler"
|
||||
)
|
||||
|
||||
// TestRealWorldScenario tests a realistic multi-job scenario.
|
||||
func TestRealWorldScenario(t *testing.T) {
|
||||
var (
|
||||
quickJobCount atomic.Int32
|
||||
hourlyJobCount atomic.Int32
|
||||
logEntries []string
|
||||
logMu sync.Mutex
|
||||
)
|
||||
|
||||
logger := &testLogger{
|
||||
onInfo: func(msg string, _ ...interface{}) {
|
||||
logMu.Lock()
|
||||
logEntries = append(logEntries, fmt.Sprintf("INFO: %s", msg))
|
||||
logMu.Unlock()
|
||||
},
|
||||
onError: func(msg string, _ ...interface{}) {
|
||||
logMu.Lock()
|
||||
logEntries = append(logEntries, fmt.Sprintf("ERROR: %s", msg))
|
||||
logMu.Unlock()
|
||||
},
|
||||
}
|
||||
|
||||
s := scheduler.New(
|
||||
scheduler.WithTimezone("UTC"),
|
||||
scheduler.WithMiddleware(
|
||||
scheduler.Recovery(func(jobName string, r interface{}) {
|
||||
t.Logf("Job %s panicked: %v", jobName, r)
|
||||
}),
|
||||
scheduler.Logging(logger),
|
||||
scheduler.Timeout(5*time.Second),
|
||||
),
|
||||
)
|
||||
|
||||
// Quick job (every second)
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "quick-check",
|
||||
Schedule: "* * * * * *",
|
||||
Handler: func(_ context.Context) error {
|
||||
quickJobCount.Add(1)
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
// Slower job (every 2 seconds)
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "slow-process",
|
||||
Schedule: "*/2 * * * * *",
|
||||
Handler: func(_ context.Context) error {
|
||||
hourlyJobCount.Add(1)
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
// Start scheduler
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatalf("failed to start scheduler: %v", err)
|
||||
}
|
||||
|
||||
// Let it run for 5 seconds
|
||||
time.Sleep(5 * time.Second)
|
||||
|
||||
// Graceful shutdown
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.Stop(ctx); err != nil {
|
||||
t.Fatalf("failed to stop scheduler: %v", err)
|
||||
}
|
||||
|
||||
// Verify execution counts
|
||||
quick := quickJobCount.Load()
|
||||
slow := hourlyJobCount.Load()
|
||||
|
||||
t.Logf("Quick job ran %d times", quick)
|
||||
t.Logf("Slow job ran %d times", slow)
|
||||
|
||||
if quick < 4 {
|
||||
t.Errorf("expected quick job to run at least 4 times, ran %d", quick)
|
||||
}
|
||||
|
||||
if slow < 2 {
|
||||
t.Errorf("expected slow job to run at least 2 times, ran %d", slow)
|
||||
}
|
||||
|
||||
// Verify logging
|
||||
logMu.Lock()
|
||||
defer logMu.Unlock()
|
||||
|
||||
hasStartLog := false
|
||||
hasCompleteLog := false
|
||||
for _, entry := range logEntries {
|
||||
if contains(entry, "Job started") {
|
||||
hasStartLog = true
|
||||
}
|
||||
if contains(entry, "Job completed") {
|
||||
hasCompleteLog = true
|
||||
}
|
||||
}
|
||||
|
||||
if !hasStartLog {
|
||||
t.Error("expected job start logs")
|
||||
}
|
||||
if !hasCompleteLog {
|
||||
t.Error("expected job completion logs")
|
||||
}
|
||||
}
|
||||
|
||||
// TestCancellationDuringExecution verifies jobs can be canceled mid-execution.
|
||||
func TestCancellationDuringExecution(t *testing.T) {
|
||||
var canceled atomic.Bool
|
||||
var started atomic.Bool
|
||||
|
||||
s := scheduler.New()
|
||||
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "long-job",
|
||||
Schedule: "* * * * * *",
|
||||
Handler: func(ctx context.Context) error {
|
||||
started.Store(true)
|
||||
// Simulate long-running work
|
||||
for i := 0; i < 100; i++ {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
canceled.Store(true)
|
||||
return ctx.Err()
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
// Keep working
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatalf("failed to start: %v", err)
|
||||
}
|
||||
|
||||
// Wait until job starts
|
||||
for i := 0; i < 30; i++ {
|
||||
if started.Load() {
|
||||
break
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
|
||||
if !started.Load() {
|
||||
t.Fatal("job did not start within timeout")
|
||||
}
|
||||
|
||||
// Stop with reasonable timeout
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.Stop(ctx); err != nil {
|
||||
t.Logf("stop returned error (may be expected): %v", err)
|
||||
}
|
||||
|
||||
if !canceled.Load() {
|
||||
t.Error("expected job to detect cancellation")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTimezoneHandling verifies timezone-aware scheduling.
|
||||
func TestTimezoneHandling(t *testing.T) {
|
||||
// Parse a schedule in a specific timezone
|
||||
schedule, err := scheduler.ParseCronExpression("0 9 * * *") // 9 AM
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse schedule: %v", err)
|
||||
}
|
||||
|
||||
// Test in New York timezone
|
||||
nyc, err := time.LoadLocation("America/New_York")
|
||||
if err != nil {
|
||||
t.Fatalf("failed to load timezone: %v", err)
|
||||
}
|
||||
|
||||
// Current time: 8:30 AM in New York
|
||||
now := time.Date(2025, 1, 15, 8, 30, 0, 0, nyc)
|
||||
|
||||
// Next run should be 9:00 AM same day
|
||||
next := schedule.Next(now)
|
||||
expected := time.Date(2025, 1, 15, 9, 0, 0, 0, nyc)
|
||||
|
||||
if !next.Equal(expected) {
|
||||
t.Errorf("next = %v, expected %v", next, expected)
|
||||
}
|
||||
|
||||
// If it's already past 9 AM
|
||||
now = time.Date(2025, 1, 15, 9, 30, 0, 0, nyc)
|
||||
next = schedule.Next(now)
|
||||
expected = time.Date(2025, 1, 16, 9, 0, 0, 0, nyc)
|
||||
|
||||
if !next.Equal(expected) {
|
||||
t.Errorf("next = %v, expected %v", next, expected)
|
||||
}
|
||||
}
|
||||
|
||||
// TestErrorPropagation verifies error handling.
|
||||
func TestErrorPropagation(t *testing.T) {
|
||||
var errorLogged atomic.Bool
|
||||
|
||||
logger := &testLogger{
|
||||
onError: func(msg string, _ ...interface{}) {
|
||||
if msg == "Job failed" {
|
||||
errorLogged.Store(true)
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
s := scheduler.New(
|
||||
scheduler.WithMiddleware(
|
||||
scheduler.Logging(logger),
|
||||
),
|
||||
)
|
||||
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "failing-job",
|
||||
Schedule: "* * * * * *",
|
||||
Handler: func(_ context.Context) error {
|
||||
return errors.New("intentional error")
|
||||
},
|
||||
})
|
||||
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatalf("failed to start: %v", err)
|
||||
}
|
||||
|
||||
// Let it run once
|
||||
time.Sleep(1500 * time.Millisecond)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.Stop(ctx); err != nil {
|
||||
t.Fatalf("failed to stop: %v", err)
|
||||
}
|
||||
|
||||
if !errorLogged.Load() {
|
||||
t.Error("expected error to be logged")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPanicRecovery verifies panic recovery middleware.
|
||||
func TestPanicRecovery(t *testing.T) {
|
||||
var panicRecovered atomic.Bool
|
||||
|
||||
s := scheduler.New(
|
||||
scheduler.WithMiddleware(
|
||||
scheduler.Recovery(func(jobName string, r interface{}) {
|
||||
panicRecovered.Store(true)
|
||||
t.Logf("Recovered from panic in job %s: %v", jobName, r)
|
||||
}),
|
||||
),
|
||||
)
|
||||
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "panicking-job",
|
||||
Schedule: "* * * * * *",
|
||||
Handler: func(_ context.Context) error {
|
||||
panic("intentional panic for testing")
|
||||
},
|
||||
})
|
||||
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatalf("failed to start: %v", err)
|
||||
}
|
||||
|
||||
// Let it run once
|
||||
time.Sleep(1500 * time.Millisecond)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.Stop(ctx); err != nil {
|
||||
t.Fatalf("failed to stop: %v", err)
|
||||
}
|
||||
|
||||
if !panicRecovered.Load() {
|
||||
t.Error("expected panic to be recovered")
|
||||
}
|
||||
}
|
||||
|
||||
// TestMultipleJobsWithDifferentSchedules verifies concurrent job execution.
|
||||
func TestMultipleJobsWithDifferentSchedules(t *testing.T) {
|
||||
var (
|
||||
job1Count atomic.Int32
|
||||
job2Count atomic.Int32
|
||||
job3Count atomic.Int32
|
||||
)
|
||||
|
||||
s := scheduler.New()
|
||||
|
||||
// Job 1: Every second
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "job-1sec",
|
||||
Schedule: "* * * * * *",
|
||||
Handler: func(_ context.Context) error {
|
||||
job1Count.Add(1)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
// Job 2: Every 2 seconds
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "job-2sec",
|
||||
Schedule: "*/2 * * * * *",
|
||||
Handler: func(_ context.Context) error {
|
||||
job2Count.Add(1)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
// Job 3: Every 3 seconds
|
||||
s.Register(&scheduler.Job{
|
||||
Name: "job-3sec",
|
||||
Schedule: "*/3 * * * * *",
|
||||
Handler: func(_ context.Context) error {
|
||||
job3Count.Add(1)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatalf("failed to start: %v", err)
|
||||
}
|
||||
|
||||
// Let them run for 6 seconds
|
||||
time.Sleep(6 * time.Second)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.Stop(ctx); err != nil {
|
||||
t.Fatalf("failed to stop: %v", err)
|
||||
}
|
||||
|
||||
// Verify counts (allowing for timing variance)
|
||||
c1 := job1Count.Load()
|
||||
c2 := job2Count.Load()
|
||||
c3 := job3Count.Load()
|
||||
|
||||
t.Logf("Job 1 ran %d times, Job 2 ran %d times, Job 3 ran %d times", c1, c2, c3)
|
||||
|
||||
if c1 < 5 {
|
||||
t.Errorf("expected job1 to run at least 5 times, ran %d", c1)
|
||||
}
|
||||
if c2 < 2 {
|
||||
t.Errorf("expected job2 to run at least 2 times, ran %d", c2)
|
||||
}
|
||||
if c3 < 1 {
|
||||
t.Errorf("expected job3 to run at least 1 time, ran %d", c3)
|
||||
}
|
||||
}
|
||||
|
||||
// Helpers
|
||||
|
||||
type testLogger struct {
|
||||
onInfo func(msg string, args ...interface{})
|
||||
onError func(msg string, args ...interface{})
|
||||
}
|
||||
|
||||
func (l *testLogger) Info(msg string, args ...interface{}) {
|
||||
if l.onInfo != nil {
|
||||
l.onInfo(msg, args...)
|
||||
}
|
||||
}
|
||||
|
||||
func (l *testLogger) Error(msg string, args ...interface{}) {
|
||||
if l.onError != nil {
|
||||
l.onError(msg, args...)
|
||||
}
|
||||
}
|
||||
|
||||
func contains(s, substr string) bool {
|
||||
return strings.Contains(s, substr)
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
// JobHandler is the function signature for scheduled job handlers.
|
||||
// The context passed to the handler will be canceled if the scheduler is shutting down.
|
||||
type JobHandler func(ctx context.Context) error
|
||||
|
||||
// Job represents a scheduled task.
|
||||
type Job struct {
|
||||
// Name is a unique identifier for this job (required).
|
||||
// Used for logging and metrics.
|
||||
Name string
|
||||
|
||||
// Schedule is a cron expression defining when this job runs (required).
|
||||
// Supports standard 5-field format: "minute hour day month weekday"
|
||||
// Examples: "0 * * * *" (hourly), "0 0 * * *" (daily at midnight)
|
||||
Schedule string
|
||||
|
||||
// Timezone for schedule evaluation (optional, defaults to UTC).
|
||||
// Use IANA timezone names: "America/New_York", "Europe/London", etc.
|
||||
Timezone string
|
||||
|
||||
// Handler is the function to execute when the job triggers (required).
|
||||
Handler JobHandler
|
||||
|
||||
// Description provides human-readable context about what this job does (optional).
|
||||
Description string
|
||||
|
||||
// Tags allow categorizing jobs for filtering/monitoring (optional).
|
||||
Tags []string
|
||||
}
|
||||
|
||||
// Validate checks if the job definition is valid.
|
||||
func (j *Job) Validate() error {
|
||||
if j.Name == "" {
|
||||
return errors.New("job name is required")
|
||||
}
|
||||
|
||||
if j.Schedule == "" {
|
||||
return errors.New("job schedule is required")
|
||||
}
|
||||
|
||||
// Validate cron expression using parser
|
||||
if _, err := ParseCronExpression(j.Schedule); err != nil {
|
||||
return errors.Wrap(err, "invalid cron expression")
|
||||
}
|
||||
|
||||
if j.Handler == nil {
|
||||
return errors.New("job handler is required")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestJobDefinition(t *testing.T) {
|
||||
callCount := 0
|
||||
job := &Job{
|
||||
Name: "test-job",
|
||||
Handler: func(_ context.Context) error {
|
||||
callCount++
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
if job.Name != "test-job" {
|
||||
t.Errorf("expected name 'test-job', got %s", job.Name)
|
||||
}
|
||||
|
||||
// Test handler execution
|
||||
if err := job.Handler(context.Background()); err != nil {
|
||||
t.Fatalf("handler failed: %v", err)
|
||||
}
|
||||
|
||||
if callCount != 1 {
|
||||
t.Errorf("expected handler to be called once, called %d times", callCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestJobValidation(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
job *Job
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid job",
|
||||
job: &Job{
|
||||
Name: "valid",
|
||||
Schedule: "0 * * * *",
|
||||
Handler: func(_ context.Context) error { return nil },
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "missing name",
|
||||
job: &Job{
|
||||
Schedule: "0 * * * *",
|
||||
Handler: func(_ context.Context) error { return nil },
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "missing schedule",
|
||||
job: &Job{
|
||||
Name: "test",
|
||||
Handler: func(_ context.Context) error { return nil },
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "invalid cron expression",
|
||||
job: &Job{
|
||||
Name: "test",
|
||||
Schedule: "invalid cron",
|
||||
Handler: func(_ context.Context) error { return nil },
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "missing handler",
|
||||
job: &Job{
|
||||
Name: "test",
|
||||
Schedule: "0 * * * *",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := tt.job.Validate()
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("Validate() error = %v, wantErr %v", err, tt.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"slices"
|
||||
"time"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
// Middleware wraps a JobHandler to add cross-cutting behavior.
|
||||
type Middleware func(JobHandler) JobHandler
|
||||
|
||||
// Chain combines multiple middleware into a single middleware.
|
||||
// Middleware are applied in the order they're provided (left to right).
|
||||
func Chain(middlewares ...Middleware) Middleware {
|
||||
return func(handler JobHandler) JobHandler {
|
||||
// Apply middleware in reverse order so first middleware wraps outermost
|
||||
for _, middleware := range slices.Backward(middlewares) {
|
||||
handler = middleware(handler)
|
||||
}
|
||||
return handler
|
||||
}
|
||||
}
|
||||
|
||||
// Recovery recovers from panics in job handlers and converts them to errors.
|
||||
func Recovery(onPanic func(jobName string, recovered interface{})) Middleware {
|
||||
return func(next JobHandler) JobHandler {
|
||||
return func(ctx context.Context) (err error) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
jobName := getJobName(ctx)
|
||||
if onPanic != nil {
|
||||
onPanic(jobName, r)
|
||||
}
|
||||
err = errors.Errorf("job %q panicked: %v", jobName, r)
|
||||
}
|
||||
}()
|
||||
return next(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Logger is a minimal logging interface.
|
||||
type Logger interface {
|
||||
Info(msg string, args ...interface{})
|
||||
Error(msg string, args ...interface{})
|
||||
}
|
||||
|
||||
// Logging adds execution logging to jobs.
|
||||
func Logging(logger Logger) Middleware {
|
||||
return func(next JobHandler) JobHandler {
|
||||
return func(ctx context.Context) error {
|
||||
jobName := getJobName(ctx)
|
||||
start := time.Now()
|
||||
|
||||
logger.Info("Job started", "job", jobName)
|
||||
|
||||
err := next(ctx)
|
||||
duration := time.Since(start)
|
||||
|
||||
if err != nil {
|
||||
logger.Error("Job failed", "job", jobName, "duration", duration, "error", err)
|
||||
} else {
|
||||
logger.Info("Job completed", "job", jobName, "duration", duration)
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Timeout wraps a job handler with a timeout.
|
||||
func Timeout(duration time.Duration) Middleware {
|
||||
return func(next JobHandler) JobHandler {
|
||||
return func(ctx context.Context) error {
|
||||
ctx, cancel := context.WithTimeout(ctx, duration)
|
||||
defer cancel()
|
||||
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- next(ctx)
|
||||
}()
|
||||
|
||||
select {
|
||||
case err := <-done:
|
||||
return err
|
||||
case <-ctx.Done():
|
||||
return errors.Errorf("job %q timed out after %v", getJobName(ctx), duration)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Context keys for job metadata.
|
||||
type contextKey int
|
||||
|
||||
const (
|
||||
jobNameKey contextKey = iota
|
||||
)
|
||||
|
||||
// withJobName adds the job name to the context.
|
||||
func withJobName(ctx context.Context, name string) context.Context {
|
||||
return context.WithValue(ctx, jobNameKey, name)
|
||||
}
|
||||
|
||||
// getJobName retrieves the job name from the context.
|
||||
func getJobName(ctx context.Context) string {
|
||||
if name, ok := ctx.Value(jobNameKey).(string); ok {
|
||||
return name
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
|
||||
// GetJobName retrieves the job name from the context (public API).
|
||||
// Returns empty string if not found.
|
||||
//
|
||||
//nolint:revive // GetJobName is the public API, getJobName is internal
|
||||
func GetJobName(ctx context.Context) string {
|
||||
return getJobName(ctx)
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMiddlewareChaining(t *testing.T) {
|
||||
var order []string
|
||||
|
||||
mw1 := func(next JobHandler) JobHandler {
|
||||
return func(ctx context.Context) error {
|
||||
order = append(order, "before-1")
|
||||
err := next(ctx)
|
||||
order = append(order, "after-1")
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
mw2 := func(next JobHandler) JobHandler {
|
||||
return func(ctx context.Context) error {
|
||||
order = append(order, "before-2")
|
||||
err := next(ctx)
|
||||
order = append(order, "after-2")
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
handler := func(_ context.Context) error {
|
||||
order = append(order, "handler")
|
||||
return nil
|
||||
}
|
||||
|
||||
chain := Chain(mw1, mw2)
|
||||
wrapped := chain(handler)
|
||||
|
||||
if err := wrapped(context.Background()); err != nil {
|
||||
t.Fatalf("wrapped handler failed: %v", err)
|
||||
}
|
||||
|
||||
expected := []string{"before-1", "before-2", "handler", "after-2", "after-1"}
|
||||
if len(order) != len(expected) {
|
||||
t.Fatalf("expected %d calls, got %d", len(expected), len(order))
|
||||
}
|
||||
|
||||
for i, want := range expected {
|
||||
if order[i] != want {
|
||||
t.Errorf("order[%d] = %q, want %q", i, order[i], want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecoveryMiddleware(t *testing.T) {
|
||||
var panicRecovered atomic.Bool
|
||||
|
||||
onPanic := func(_ string, _ interface{}) {
|
||||
panicRecovered.Store(true)
|
||||
}
|
||||
|
||||
handler := func(_ context.Context) error {
|
||||
panic("simulated panic")
|
||||
}
|
||||
|
||||
recovery := Recovery(onPanic)
|
||||
wrapped := recovery(handler)
|
||||
|
||||
// Should not panic, error should be returned
|
||||
err := wrapped(withJobName(context.Background(), "test-job"))
|
||||
if err == nil {
|
||||
t.Error("expected error from recovered panic")
|
||||
}
|
||||
|
||||
if !panicRecovered.Load() {
|
||||
t.Error("panic handler was not called")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoggingMiddleware(t *testing.T) {
|
||||
var loggedStart, loggedEnd atomic.Bool
|
||||
var loggedError atomic.Bool
|
||||
|
||||
logger := &testLogger{
|
||||
onInfo: func(msg string, _ ...interface{}) {
|
||||
if msg == "Job started" {
|
||||
loggedStart.Store(true)
|
||||
} else if msg == "Job completed" {
|
||||
loggedEnd.Store(true)
|
||||
}
|
||||
},
|
||||
onError: func(msg string, _ ...interface{}) {
|
||||
if msg == "Job failed" {
|
||||
loggedError.Store(true)
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
// Test successful execution
|
||||
handler := func(_ context.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
logging := Logging(logger)
|
||||
wrapped := logging(handler)
|
||||
|
||||
if err := wrapped(withJobName(context.Background(), "test-job")); err != nil {
|
||||
t.Fatalf("handler failed: %v", err)
|
||||
}
|
||||
|
||||
if !loggedStart.Load() {
|
||||
t.Error("start was not logged")
|
||||
}
|
||||
if !loggedEnd.Load() {
|
||||
t.Error("end was not logged")
|
||||
}
|
||||
|
||||
// Test error handling
|
||||
handlerErr := func(_ context.Context) error {
|
||||
return errors.New("job error")
|
||||
}
|
||||
|
||||
wrappedErr := logging(handlerErr)
|
||||
_ = wrappedErr(withJobName(context.Background(), "test-job-error"))
|
||||
|
||||
if !loggedError.Load() {
|
||||
t.Error("error was not logged")
|
||||
}
|
||||
}
|
||||
|
||||
type testLogger struct {
|
||||
onInfo func(msg string, args ...interface{})
|
||||
onError func(msg string, args ...interface{})
|
||||
}
|
||||
|
||||
func (l *testLogger) Info(msg string, args ...interface{}) {
|
||||
if l.onInfo != nil {
|
||||
l.onInfo(msg, args...)
|
||||
}
|
||||
}
|
||||
|
||||
func (l *testLogger) Error(msg string, args ...interface{}) {
|
||||
if l.onError != nil {
|
||||
l.onError(msg, args...)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
// Schedule represents a parsed cron expression.
|
||||
type Schedule struct {
|
||||
seconds fieldMatcher // 0-59 (optional, for 6-field format)
|
||||
minutes fieldMatcher // 0-59
|
||||
hours fieldMatcher // 0-23
|
||||
days fieldMatcher // 1-31
|
||||
months fieldMatcher // 1-12
|
||||
weekdays fieldMatcher // 0-7 (0 and 7 are Sunday)
|
||||
hasSecs bool
|
||||
}
|
||||
|
||||
// fieldMatcher determines if a field value matches.
|
||||
type fieldMatcher interface {
|
||||
matches(value int) bool
|
||||
}
|
||||
|
||||
// ParseCronExpression parses a cron expression and returns a Schedule.
|
||||
// Supports both 5-field (minute hour day month weekday) and 6-field (second minute hour day month weekday) formats.
|
||||
func ParseCronExpression(expr string) (*Schedule, error) {
|
||||
if expr == "" {
|
||||
return nil, errors.New("empty cron expression")
|
||||
}
|
||||
|
||||
fields := strings.Fields(expr)
|
||||
if len(fields) != 5 && len(fields) != 6 {
|
||||
return nil, errors.Errorf("invalid cron expression: expected 5 or 6 fields, got %d", len(fields))
|
||||
}
|
||||
|
||||
s := &Schedule{
|
||||
hasSecs: len(fields) == 6,
|
||||
}
|
||||
|
||||
var err error
|
||||
offset := 0
|
||||
|
||||
// Parse seconds (if 6-field format)
|
||||
if s.hasSecs {
|
||||
s.seconds, err = parseField(fields[0], 0, 59)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "invalid seconds field")
|
||||
}
|
||||
offset = 1
|
||||
} else {
|
||||
s.seconds = &exactMatcher{value: 0} // Default to 0 seconds
|
||||
}
|
||||
|
||||
// Parse minutes
|
||||
s.minutes, err = parseField(fields[offset], 0, 59)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "invalid minutes field")
|
||||
}
|
||||
|
||||
// Parse hours
|
||||
s.hours, err = parseField(fields[offset+1], 0, 23)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "invalid hours field")
|
||||
}
|
||||
|
||||
// Parse days
|
||||
s.days, err = parseField(fields[offset+2], 1, 31)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "invalid days field")
|
||||
}
|
||||
|
||||
// Parse months
|
||||
s.months, err = parseField(fields[offset+3], 1, 12)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "invalid months field")
|
||||
}
|
||||
|
||||
// Parse weekdays (0-7, where both 0 and 7 represent Sunday)
|
||||
s.weekdays, err = parseField(fields[offset+4], 0, 7)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "invalid weekdays field")
|
||||
}
|
||||
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// Next returns the next time the schedule should run after the given time.
|
||||
func (s *Schedule) Next(from time.Time) time.Time {
|
||||
// Start from the next second/minute
|
||||
if s.hasSecs {
|
||||
from = from.Add(1 * time.Second).Truncate(time.Second)
|
||||
} else {
|
||||
from = from.Add(1 * time.Minute).Truncate(time.Minute)
|
||||
}
|
||||
|
||||
// Cap search at 4 years to prevent infinite loops
|
||||
maxTime := from.AddDate(4, 0, 0)
|
||||
|
||||
for from.Before(maxTime) {
|
||||
if s.matches(from) {
|
||||
return from
|
||||
}
|
||||
|
||||
// Advance to next potential match
|
||||
if s.hasSecs {
|
||||
from = from.Add(1 * time.Second)
|
||||
} else {
|
||||
from = from.Add(1 * time.Minute)
|
||||
}
|
||||
}
|
||||
|
||||
// Should never reach here with valid cron expressions
|
||||
return time.Time{}
|
||||
}
|
||||
|
||||
// matches checks if the given time matches the schedule.
|
||||
func (s *Schedule) matches(t time.Time) bool {
|
||||
return s.seconds.matches(t.Second()) &&
|
||||
s.minutes.matches(t.Minute()) &&
|
||||
s.hours.matches(t.Hour()) &&
|
||||
s.months.matches(int(t.Month())) &&
|
||||
(s.days.matches(t.Day()) || s.weekdays.matches(int(t.Weekday())))
|
||||
}
|
||||
|
||||
// parseField parses a single cron field (supports *, ranges, lists, steps).
|
||||
func parseField(field string, min, max int) (fieldMatcher, error) {
|
||||
// Wildcard
|
||||
if field == "*" {
|
||||
return &wildcardMatcher{}, nil
|
||||
}
|
||||
|
||||
// Step values (*/N)
|
||||
if strings.HasPrefix(field, "*/") {
|
||||
step, err := strconv.Atoi(field[2:])
|
||||
if err != nil || step < 1 || step > max {
|
||||
return nil, errors.Errorf("invalid step value: %s", field)
|
||||
}
|
||||
return &stepMatcher{step: step, min: min, max: max}, nil
|
||||
}
|
||||
|
||||
// List (1,2,3)
|
||||
if strings.Contains(field, ",") {
|
||||
parts := strings.Split(field, ",")
|
||||
values := make([]int, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
val, err := strconv.Atoi(strings.TrimSpace(p))
|
||||
if err != nil || val < min || val > max {
|
||||
return nil, errors.Errorf("invalid list value: %s", p)
|
||||
}
|
||||
values = append(values, val)
|
||||
}
|
||||
return &listMatcher{values: values}, nil
|
||||
}
|
||||
|
||||
// Range (1-5)
|
||||
if strings.Contains(field, "-") {
|
||||
parts := strings.Split(field, "-")
|
||||
if len(parts) != 2 {
|
||||
return nil, errors.Errorf("invalid range: %s", field)
|
||||
}
|
||||
start, err1 := strconv.Atoi(strings.TrimSpace(parts[0]))
|
||||
end, err2 := strconv.Atoi(strings.TrimSpace(parts[1]))
|
||||
if err1 != nil || err2 != nil || start < min || end > max || start > end {
|
||||
return nil, errors.Errorf("invalid range: %s", field)
|
||||
}
|
||||
return &rangeMatcher{start: start, end: end}, nil
|
||||
}
|
||||
|
||||
// Exact value
|
||||
val, err := strconv.Atoi(field)
|
||||
if err != nil || val < min || val > max {
|
||||
return nil, errors.Errorf("invalid value: %s (must be between %d and %d)", field, min, max)
|
||||
}
|
||||
return &exactMatcher{value: val}, nil
|
||||
}
|
||||
|
||||
// wildcardMatcher matches any value.
|
||||
type wildcardMatcher struct{}
|
||||
|
||||
func (*wildcardMatcher) matches(_ int) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// exactMatcher matches a specific value.
|
||||
type exactMatcher struct {
|
||||
value int
|
||||
}
|
||||
|
||||
func (m *exactMatcher) matches(value int) bool {
|
||||
return value == m.value
|
||||
}
|
||||
|
||||
// rangeMatcher matches values in a range.
|
||||
type rangeMatcher struct {
|
||||
start, end int
|
||||
}
|
||||
|
||||
func (m *rangeMatcher) matches(value int) bool {
|
||||
return value >= m.start && value <= m.end
|
||||
}
|
||||
|
||||
// listMatcher matches any value in a list.
|
||||
type listMatcher struct {
|
||||
values []int
|
||||
}
|
||||
|
||||
func (m *listMatcher) matches(value int) bool {
|
||||
return slices.Contains(m.values, value)
|
||||
}
|
||||
|
||||
// stepMatcher matches values at regular intervals.
|
||||
type stepMatcher struct {
|
||||
step, min, max int
|
||||
}
|
||||
|
||||
func (m *stepMatcher) matches(value int) bool {
|
||||
if value < m.min || value > m.max {
|
||||
return false
|
||||
}
|
||||
return (value-m.min)%m.step == 0
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestParseCronExpression(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
expr string
|
||||
wantErr bool
|
||||
}{
|
||||
// Standard 5-field format
|
||||
{"every minute", "* * * * *", false},
|
||||
{"hourly", "0 * * * *", false},
|
||||
{"daily midnight", "0 0 * * *", false},
|
||||
{"weekly sunday", "0 0 * * 0", false},
|
||||
{"monthly", "0 0 1 * *", false},
|
||||
{"specific time", "30 14 * * *", false}, // 2:30 PM daily
|
||||
{"range", "0 9-17 * * *", false}, // Every hour 9 AM - 5 PM
|
||||
{"step", "*/15 * * * *", false}, // Every 15 minutes
|
||||
{"list", "0 8,12,18 * * *", false}, // 8 AM, 12 PM, 6 PM
|
||||
|
||||
// 6-field format with seconds
|
||||
{"with seconds", "0 * * * * *", false},
|
||||
{"every 30 seconds", "*/30 * * * * *", false},
|
||||
|
||||
// Invalid expressions
|
||||
{"empty", "", true},
|
||||
{"too few fields", "* * *", true},
|
||||
{"too many fields", "* * * * * * *", true},
|
||||
{"invalid minute", "60 * * * *", true},
|
||||
{"invalid hour", "0 24 * * *", true},
|
||||
{"invalid day", "0 0 32 * *", true},
|
||||
{"invalid month", "0 0 1 13 *", true},
|
||||
{"invalid weekday", "0 0 * * 8", true},
|
||||
{"garbage", "not a cron expression", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
schedule, err := ParseCronExpression(tt.expr)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("ParseCronExpression(%q) error = %v, wantErr %v", tt.expr, err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if !tt.wantErr && schedule == nil {
|
||||
t.Errorf("ParseCronExpression(%q) returned nil schedule without error", tt.expr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestScheduleNext(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
expr string
|
||||
from time.Time
|
||||
expected time.Time
|
||||
}{
|
||||
{
|
||||
name: "every minute from start of hour",
|
||||
expr: "* * * * *",
|
||||
from: time.Date(2025, 1, 1, 10, 0, 0, 0, time.UTC),
|
||||
expected: time.Date(2025, 1, 1, 10, 1, 0, 0, time.UTC),
|
||||
},
|
||||
{
|
||||
name: "hourly at minute 30",
|
||||
expr: "30 * * * *",
|
||||
from: time.Date(2025, 1, 1, 10, 0, 0, 0, time.UTC),
|
||||
expected: time.Date(2025, 1, 1, 10, 30, 0, 0, time.UTC),
|
||||
},
|
||||
{
|
||||
name: "hourly at minute 30 (already past)",
|
||||
expr: "30 * * * *",
|
||||
from: time.Date(2025, 1, 1, 10, 45, 0, 0, time.UTC),
|
||||
expected: time.Date(2025, 1, 1, 11, 30, 0, 0, time.UTC),
|
||||
},
|
||||
{
|
||||
name: "daily at 2 AM",
|
||||
expr: "0 2 * * *",
|
||||
from: time.Date(2025, 1, 1, 10, 0, 0, 0, time.UTC),
|
||||
expected: time.Date(2025, 1, 2, 2, 0, 0, 0, time.UTC),
|
||||
},
|
||||
{
|
||||
name: "every 15 minutes",
|
||||
expr: "*/15 * * * *",
|
||||
from: time.Date(2025, 1, 1, 10, 7, 0, 0, time.UTC),
|
||||
expected: time.Date(2025, 1, 1, 10, 15, 0, 0, time.UTC),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
schedule, err := ParseCronExpression(tt.expr)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse expression: %v", err)
|
||||
}
|
||||
|
||||
next := schedule.Next(tt.from)
|
||||
if !next.Equal(tt.expected) {
|
||||
t.Errorf("Next(%v) = %v, expected %v", tt.from, next, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestScheduleNextWithTimezone(t *testing.T) {
|
||||
nyc, _ := time.LoadLocation("America/New_York")
|
||||
|
||||
// Schedule for 9 AM in New York
|
||||
schedule, err := ParseCronExpression("0 9 * * *")
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse expression: %v", err)
|
||||
}
|
||||
|
||||
// Current time: 8 AM in New York
|
||||
from := time.Date(2025, 1, 1, 8, 0, 0, 0, nyc)
|
||||
next := schedule.Next(from)
|
||||
|
||||
// Should be 9 AM same day in New York
|
||||
expected := time.Date(2025, 1, 1, 9, 0, 0, 0, nyc)
|
||||
if !next.Equal(expected) {
|
||||
t.Errorf("Next(%v) = %v, expected %v", from, next, expected)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
// Scheduler manages scheduled jobs.
|
||||
type Scheduler struct {
|
||||
jobs map[string]*registeredJob
|
||||
jobsMu sync.RWMutex
|
||||
timezone *time.Location
|
||||
middleware Middleware
|
||||
running bool
|
||||
runningMu sync.RWMutex
|
||||
stopCh chan struct{}
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// registeredJob wraps a Job with runtime state.
|
||||
type registeredJob struct {
|
||||
job *Job
|
||||
cancelFn context.CancelFunc
|
||||
}
|
||||
|
||||
// Option configures a Scheduler.
|
||||
type Option func(*Scheduler)
|
||||
|
||||
// WithTimezone sets the default timezone for all jobs.
|
||||
func WithTimezone(tz string) Option {
|
||||
return func(s *Scheduler) {
|
||||
loc, err := time.LoadLocation(tz)
|
||||
if err != nil {
|
||||
// Default to UTC on invalid timezone
|
||||
loc = time.UTC
|
||||
}
|
||||
s.timezone = loc
|
||||
}
|
||||
}
|
||||
|
||||
// WithMiddleware sets middleware to wrap all job handlers.
|
||||
func WithMiddleware(mw ...Middleware) Option {
|
||||
return func(s *Scheduler) {
|
||||
if len(mw) > 0 {
|
||||
s.middleware = Chain(mw...)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// New creates a new Scheduler with optional configuration.
|
||||
func New(opts ...Option) *Scheduler {
|
||||
s := &Scheduler{
|
||||
jobs: make(map[string]*registeredJob),
|
||||
timezone: time.UTC,
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(s)
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
// Register adds a job to the scheduler.
|
||||
// Jobs must be registered before calling Start().
|
||||
func (s *Scheduler) Register(job *Job) error {
|
||||
if job == nil {
|
||||
return errors.New("job cannot be nil")
|
||||
}
|
||||
|
||||
if err := job.Validate(); err != nil {
|
||||
return errors.Wrap(err, "invalid job")
|
||||
}
|
||||
|
||||
s.jobsMu.Lock()
|
||||
defer s.jobsMu.Unlock()
|
||||
|
||||
if _, exists := s.jobs[job.Name]; exists {
|
||||
return errors.Errorf("job with name %q already registered", job.Name)
|
||||
}
|
||||
|
||||
s.jobs[job.Name] = ®isteredJob{job: job}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Start begins executing scheduled jobs.
|
||||
func (s *Scheduler) Start() error {
|
||||
s.runningMu.Lock()
|
||||
defer s.runningMu.Unlock()
|
||||
|
||||
if s.running {
|
||||
return errors.New("scheduler already running")
|
||||
}
|
||||
|
||||
s.jobsMu.RLock()
|
||||
defer s.jobsMu.RUnlock()
|
||||
|
||||
// Parse and schedule all jobs
|
||||
for _, rj := range s.jobs {
|
||||
schedule, err := ParseCronExpression(rj.job.Schedule)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "failed to parse schedule for job %q", rj.job.Name)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
rj.cancelFn = cancel
|
||||
|
||||
s.wg.Add(1)
|
||||
go s.runJobWithSchedule(ctx, rj, schedule)
|
||||
}
|
||||
|
||||
s.running = true
|
||||
return nil
|
||||
}
|
||||
|
||||
// runJobWithSchedule executes a job according to its cron schedule.
|
||||
func (s *Scheduler) runJobWithSchedule(ctx context.Context, rj *registeredJob, schedule *Schedule) {
|
||||
defer s.wg.Done()
|
||||
|
||||
// Apply middleware to handler
|
||||
handler := rj.job.Handler
|
||||
if s.middleware != nil {
|
||||
handler = s.middleware(handler)
|
||||
}
|
||||
|
||||
for {
|
||||
// Calculate next run time
|
||||
now := time.Now()
|
||||
if rj.job.Timezone != "" {
|
||||
loc, err := time.LoadLocation(rj.job.Timezone)
|
||||
if err == nil {
|
||||
now = now.In(loc)
|
||||
}
|
||||
} else if s.timezone != nil {
|
||||
now = now.In(s.timezone)
|
||||
}
|
||||
|
||||
next := schedule.Next(now)
|
||||
duration := time.Until(next)
|
||||
|
||||
timer := time.NewTimer(duration)
|
||||
|
||||
select {
|
||||
case <-timer.C:
|
||||
// Add job name to context and execute
|
||||
jobCtx := withJobName(ctx, rj.job.Name)
|
||||
if err := handler(jobCtx); err != nil {
|
||||
// Error already handled by middleware (if any)
|
||||
_ = err
|
||||
}
|
||||
case <-ctx.Done():
|
||||
// Stop the timer to prevent it from firing. The timer will be garbage collected.
|
||||
timer.Stop()
|
||||
return
|
||||
case <-s.stopCh:
|
||||
// Stop the timer to prevent it from firing. The timer will be garbage collected.
|
||||
timer.Stop()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Stop gracefully shuts down the scheduler.
|
||||
// It waits for all running jobs to complete or until the context is canceled.
|
||||
func (s *Scheduler) Stop(ctx context.Context) error {
|
||||
s.runningMu.Lock()
|
||||
if !s.running {
|
||||
s.runningMu.Unlock()
|
||||
return errors.New("scheduler not running")
|
||||
}
|
||||
s.running = false
|
||||
s.runningMu.Unlock()
|
||||
|
||||
// Cancel all job contexts
|
||||
s.jobsMu.RLock()
|
||||
for _, rj := range s.jobs {
|
||||
if rj.cancelFn != nil {
|
||||
rj.cancelFn()
|
||||
}
|
||||
}
|
||||
s.jobsMu.RUnlock()
|
||||
|
||||
// Signal stop and wait for jobs to finish
|
||||
close(s.stopCh)
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
s.wg.Wait()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,165 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSchedulerCreation(t *testing.T) {
|
||||
s := New()
|
||||
if s == nil {
|
||||
t.Fatal("New() returned nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerWithTimezone(t *testing.T) {
|
||||
s := New(WithTimezone("America/New_York"))
|
||||
if s == nil {
|
||||
t.Fatal("New() with timezone returned nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestJobRegistration(t *testing.T) {
|
||||
s := New()
|
||||
|
||||
job := &Job{
|
||||
Name: "test-registration",
|
||||
Schedule: "0 * * * *",
|
||||
Handler: func(_ context.Context) error { return nil },
|
||||
}
|
||||
|
||||
if err := s.Register(job); err != nil {
|
||||
t.Fatalf("failed to register valid job: %v", err)
|
||||
}
|
||||
|
||||
// Registering duplicate name should fail
|
||||
if err := s.Register(job); err == nil {
|
||||
t.Error("expected error when registering duplicate job name")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerStartStop(t *testing.T) {
|
||||
s := New()
|
||||
|
||||
var runCount atomic.Int32
|
||||
job := &Job{
|
||||
Name: "test-start-stop",
|
||||
Schedule: "* * * * * *", // Every second (6-field format)
|
||||
Handler: func(_ context.Context) error {
|
||||
runCount.Add(1)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
if err := s.Register(job); err != nil {
|
||||
t.Fatalf("failed to register job: %v", err)
|
||||
}
|
||||
|
||||
// Start scheduler
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatalf("failed to start scheduler: %v", err)
|
||||
}
|
||||
|
||||
// Let it run for 2.5 seconds
|
||||
time.Sleep(2500 * time.Millisecond)
|
||||
|
||||
// Stop scheduler
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.Stop(ctx); err != nil {
|
||||
t.Fatalf("failed to stop scheduler: %v", err)
|
||||
}
|
||||
|
||||
count := runCount.Load()
|
||||
// Should have run at least twice (at 0s and 1s, maybe 2s)
|
||||
if count < 2 {
|
||||
t.Errorf("expected job to run at least 2 times, ran %d times", count)
|
||||
}
|
||||
|
||||
// Verify it stopped (count shouldn't increase)
|
||||
finalCount := count
|
||||
time.Sleep(1500 * time.Millisecond)
|
||||
if runCount.Load() != finalCount {
|
||||
t.Error("scheduler did not stop - job continued running")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerWithMiddleware(t *testing.T) {
|
||||
var executionLog []string
|
||||
var logMu sync.Mutex
|
||||
|
||||
logger := &testLogger{
|
||||
onInfo: func(msg string, _ ...interface{}) {
|
||||
logMu.Lock()
|
||||
executionLog = append(executionLog, fmt.Sprintf("INFO: %s", msg))
|
||||
logMu.Unlock()
|
||||
},
|
||||
onError: func(msg string, _ ...interface{}) {
|
||||
logMu.Lock()
|
||||
executionLog = append(executionLog, fmt.Sprintf("ERROR: %s", msg))
|
||||
logMu.Unlock()
|
||||
},
|
||||
}
|
||||
|
||||
s := New(WithMiddleware(
|
||||
Recovery(func(jobName string, r interface{}) {
|
||||
logMu.Lock()
|
||||
executionLog = append(executionLog, fmt.Sprintf("PANIC: %s - %v", jobName, r))
|
||||
logMu.Unlock()
|
||||
}),
|
||||
Logging(logger),
|
||||
))
|
||||
|
||||
job := &Job{
|
||||
Name: "test-middleware",
|
||||
Schedule: "* * * * * *", // Every second
|
||||
Handler: func(_ context.Context) error {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
if err := s.Register(job); err != nil {
|
||||
t.Fatalf("failed to register job: %v", err)
|
||||
}
|
||||
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatalf("failed to start: %v", err)
|
||||
}
|
||||
|
||||
time.Sleep(1500 * time.Millisecond)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := s.Stop(ctx); err != nil {
|
||||
t.Fatalf("failed to stop: %v", err)
|
||||
}
|
||||
|
||||
logMu.Lock()
|
||||
defer logMu.Unlock()
|
||||
|
||||
// Should have at least one start and one completion log
|
||||
hasStart := false
|
||||
hasCompletion := false
|
||||
for _, log := range executionLog {
|
||||
if strings.Contains(log, "Job started") {
|
||||
hasStart = true
|
||||
}
|
||||
if strings.Contains(log, "Job completed") {
|
||||
hasCompletion = true
|
||||
}
|
||||
}
|
||||
|
||||
if !hasStart {
|
||||
t.Error("expected job start log")
|
||||
}
|
||||
if !hasCompletion {
|
||||
t.Error("expected job completion log")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package s3
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
|
||||
"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/pkg/errors"
|
||||
|
||||
storepb "github.com/usememos/memos/proto/gen/store"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
Client *s3.Client
|
||||
Bucket *string
|
||||
}
|
||||
|
||||
func NewClient(ctx context.Context, s3Config *storepb.StorageS3Config) (*Client, error) {
|
||||
loadOptions := []func(*config.LoadOptions) error{
|
||||
config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(s3Config.AccessKeyId, s3Config.AccessKeySecret, "")),
|
||||
config.WithRegion(s3Config.Region),
|
||||
}
|
||||
if s3Config.InsecureSkipTlsVerify {
|
||||
// Skip TLS certificate verification for endpoints using self-signed certificates.
|
||||
// This is opt-in and removes protection against man-in-the-middle attacks.
|
||||
httpClient := awshttp.NewBuildableClient().WithTransportOptions(func(tr *http.Transport) {
|
||||
tr.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402 -- opt-in for self-signed S3 endpoints
|
||||
})
|
||||
loadOptions = append(loadOptions, config.WithHTTPClient(httpClient))
|
||||
}
|
||||
|
||||
cfg, err := config.LoadDefaultConfig(ctx, loadOptions...)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to load s3 config")
|
||||
}
|
||||
|
||||
client := s3.NewFromConfig(cfg, func(o *s3.Options) {
|
||||
o.BaseEndpoint = aws.String(s3Config.Endpoint)
|
||||
o.UsePathStyle = s3Config.UsePathStyle
|
||||
o.RequestChecksumCalculation = aws.RequestChecksumCalculationWhenRequired
|
||||
o.ResponseChecksumValidation = aws.ResponseChecksumValidationWhenRequired
|
||||
})
|
||||
return &Client{
|
||||
Client: client,
|
||||
Bucket: aws.String(s3Config.Bucket),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// UploadObject uploads an object to S3.
|
||||
func (c *Client) UploadObject(ctx context.Context, key string, fileType string, content io.Reader) (string, error) {
|
||||
putInput := s3.PutObjectInput{
|
||||
Bucket: c.Bucket,
|
||||
Key: aws.String(key),
|
||||
ContentType: aws.String(fileType),
|
||||
Body: content,
|
||||
}
|
||||
if _, err := c.Client.PutObject(ctx, &putInput); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return key, nil
|
||||
}
|
||||
|
||||
// PresignGetObject presigns an object in S3.
|
||||
func (c *Client) PresignGetObject(ctx context.Context, key string) (string, error) {
|
||||
presignClient := s3.NewPresignClient(c.Client)
|
||||
presignResult, err := presignClient.PresignGetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: aws.String(*c.Bucket),
|
||||
Key: aws.String(key),
|
||||
}, func(opts *s3.PresignOptions) {
|
||||
// Set the expiration time of the presigned URL to 5 days.
|
||||
// Reference: https://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-query-string-auth.html
|
||||
opts.Expires = time.Duration(5 * 24 * time.Hour)
|
||||
})
|
||||
if err != nil {
|
||||
return "", errors.Wrap(err, "failed to presign get object")
|
||||
}
|
||||
return presignResult.URL, nil
|
||||
}
|
||||
|
||||
// GetObject retrieves an object from S3.
|
||||
func (c *Client) GetObject(ctx context.Context, key string) ([]byte, error) {
|
||||
output, err := c.Client.GetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: c.Bucket,
|
||||
Key: aws.String(key),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to download object")
|
||||
}
|
||||
defer output.Body.Close()
|
||||
data, err := io.ReadAll(output.Body)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to read object body")
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// GetObjectStream retrieves an object from S3 as a stream.
|
||||
func (c *Client) GetObjectStream(ctx context.Context, key string) (io.ReadCloser, error) {
|
||||
output, err := c.Client.GetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: c.Bucket,
|
||||
Key: aws.String(key),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to get object")
|
||||
}
|
||||
return output.Body, nil
|
||||
}
|
||||
|
||||
// DeleteObject deletes an object in S3.
|
||||
func (c *Client) DeleteObject(ctx context.Context, key string) error {
|
||||
_, err := c.Client.DeleteObject(ctx, &s3.DeleteObjectInput{
|
||||
Bucket: c.Bucket,
|
||||
Key: aws.String(key),
|
||||
})
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "failed to delete object")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
package testutil
|
||||
|
||||
// BuildMotionPhotoJPEG returns a minimal JPEG blob with Motion Photo metadata
|
||||
// and an embedded MP4 header for tests.
|
||||
func BuildMotionPhotoJPEG() []byte {
|
||||
return append(
|
||||
[]byte{
|
||||
0xFF, 0xD8, 0xFF, 0xE1,
|
||||
},
|
||||
append(
|
||||
[]byte(`<?xpacket begin=""?><rdf:Description GCamera:MotionPhoto="1" GCamera:MotionPhotoPresentationTimestampUs="123456"></rdf:Description>`),
|
||||
[]byte{
|
||||
0xFF, 0xD9,
|
||||
0x00, 0x00, 0x00, 0x10, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm', 0x00, 0x00, 0x00, 0x00,
|
||||
}...,
|
||||
)...,
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package util //nolint:revive // util namespace is intentional for shared helpers
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"math/big"
|
||||
"net/mail"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ConvertStringToInt32 converts a string to int32.
|
||||
func ConvertStringToInt32(src string) (int32, error) {
|
||||
parsed, err := strconv.ParseInt(src, 10, 32)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return int32(parsed), nil
|
||||
}
|
||||
|
||||
// HasPrefixes returns true if the string s has any of the given prefixes.
|
||||
func HasPrefixes(src string, prefixes ...string) bool {
|
||||
for _, prefix := range prefixes {
|
||||
if strings.HasPrefix(src, prefix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ValidateEmail validates the email.
|
||||
func ValidateEmail(email string) bool {
|
||||
if _, err := mail.ParseAddress(email); err != nil {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func GenUUID() string {
|
||||
return uuid.New().String()
|
||||
}
|
||||
|
||||
var letters = []rune("0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
|
||||
|
||||
// RandomString returns a random string with length n.
|
||||
func RandomString(n int) (string, error) {
|
||||
var sb strings.Builder
|
||||
sb.Grow(n)
|
||||
for i := 0; i < n; i++ {
|
||||
// The reason for using crypto/rand instead of math/rand is that
|
||||
// the former relies on hardware to generate random numbers and
|
||||
// thus has a stronger source of random numbers.
|
||||
randNum, err := rand.Int(rand.Reader, big.NewInt(int64(len(letters))))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := sb.WriteRune(letters[randNum.Uint64()]); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
return sb.String(), nil
|
||||
}
|
||||
|
||||
// ReplaceString replaces all occurrences of old in slice with new.
|
||||
func ReplaceString(slice []string, old, new string) []string {
|
||||
for i, s := range slice {
|
||||
if s == old {
|
||||
slice[i] = new
|
||||
}
|
||||
}
|
||||
return slice
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package util //nolint:revive // util is an appropriate package name for utility functions
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestValidateEmail(t *testing.T) {
|
||||
tests := []struct {
|
||||
email string
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
email: "t@gmail.com",
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
email: "@usememos.com",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
email: "1@gmail",
|
||||
want: true,
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
result := ValidateEmail(test.email)
|
||||
if result != test.want {
|
||||
t.Errorf("Validate Email %s: got result %v, want %v.", test.email, result, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package version
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/mod/semver"
|
||||
)
|
||||
|
||||
// Version is set by release builds and defaults to a development build marker.
|
||||
var Version = "dev"
|
||||
|
||||
// Commit is set by CI builds and defaults to an unknown revision marker.
|
||||
var Commit = "unknown"
|
||||
|
||||
func GetCurrentVersion() string {
|
||||
return Version
|
||||
}
|
||||
|
||||
// GetMinorVersion extracts the minor version (e.g., "0.25") from a full version string (e.g., "0.25.1").
|
||||
// Returns the minor version string or empty string if the version format is invalid.
|
||||
// Version format should be "major.minor.patch" (e.g., "0.25.1").
|
||||
func GetMinorVersion(version string) string {
|
||||
versionList := strings.Split(version, ".")
|
||||
if len(versionList) < 2 {
|
||||
return ""
|
||||
}
|
||||
// Return major.minor only (first two components)
|
||||
return versionList[0] + "." + versionList[1]
|
||||
}
|
||||
|
||||
// IsVersionGreaterOrEqualThan returns true if version is greater than or equal to target.
|
||||
func IsVersionGreaterOrEqualThan(version, target string) bool {
|
||||
return semver.Compare(fmt.Sprintf("v%s", version), fmt.Sprintf("v%s", target)) > -1
|
||||
}
|
||||
|
||||
// IsVersionGreaterThan returns true if version is greater than target.
|
||||
func IsVersionGreaterThan(version, target string) bool {
|
||||
return semver.Compare(fmt.Sprintf("v%s", version), fmt.Sprintf("v%s", target)) > 0
|
||||
}
|
||||
|
||||
type SortVersion []string
|
||||
|
||||
func (s SortVersion) Len() int {
|
||||
return len(s)
|
||||
}
|
||||
|
||||
func (s SortVersion) Swap(i, j int) {
|
||||
s[i], s[j] = s[j], s[i]
|
||||
}
|
||||
|
||||
func (s SortVersion) Less(i, j int) bool {
|
||||
v1 := fmt.Sprintf("v%s", s[i])
|
||||
v2 := fmt.Sprintf("v%s", s[j])
|
||||
return semver.Compare(v1, v2) == -1
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
package version
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"golang.org/x/mod/semver"
|
||||
)
|
||||
|
||||
func TestIsVersionGreaterOrEqualThan(t *testing.T) {
|
||||
tests := []struct {
|
||||
version string
|
||||
target string
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
version: "0.9.1",
|
||||
target: "0.9.1",
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
version: "0.10.0",
|
||||
target: "0.9.1",
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
version: "0.9.0",
|
||||
target: "0.9.1",
|
||||
want: false,
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
result := IsVersionGreaterOrEqualThan(test.version, test.target)
|
||||
if result != test.want {
|
||||
t.Errorf("got result %v, want %v.", result, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsVersionGreaterThan(t *testing.T) {
|
||||
tests := []struct {
|
||||
version string
|
||||
target string
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
version: "0.9.1",
|
||||
target: "0.9.1",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
version: "0.10.0",
|
||||
target: "0.8.0",
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
version: "0.23",
|
||||
target: "0.22",
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
version: "0.8.0",
|
||||
target: "0.10.0",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
version: "0.9.0",
|
||||
target: "0.9.1",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
version: "0.22",
|
||||
target: "0.22",
|
||||
want: false,
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
result := IsVersionGreaterThan(test.version, test.target)
|
||||
if result != test.want {
|
||||
t.Errorf("got result %v, want %v.", result, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSortVersion(t *testing.T) {
|
||||
tests := []struct {
|
||||
versionList []string
|
||||
want []string
|
||||
}{
|
||||
{
|
||||
versionList: []string{"0.9.1", "0.10.0", "0.8.0"},
|
||||
want: []string{"0.8.0", "0.9.1", "0.10.0"},
|
||||
},
|
||||
{
|
||||
versionList: []string{"1.9.1", "0.9.1", "0.10.0", "0.8.0"},
|
||||
want: []string{"0.8.0", "0.9.1", "0.10.0", "1.9.1"},
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
slices.SortFunc(test.versionList, func(a, b string) int {
|
||||
return semver.Compare("v"+a, "v"+b)
|
||||
})
|
||||
assert.Equal(t, test.versionList, test.want)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package webhook
|
||||
|
||||
import (
|
||||
"net"
|
||||
"net/url"
|
||||
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
// reservedCIDRs lists IP ranges that must never be targeted by outbound webhook requests.
|
||||
// Covers loopback, RFC-1918 private, link-local (including cloud IMDS at 169.254.169.254),
|
||||
// and their IPv6 equivalents.
|
||||
var reservedCIDRs = []string{
|
||||
"127.0.0.0/8", // IPv4 loopback
|
||||
"10.0.0.0/8", // RFC-1918 class A
|
||||
"172.16.0.0/12", // RFC-1918 class B
|
||||
"192.168.0.0/16", // RFC-1918 class C
|
||||
"169.254.0.0/16", // Link-local / cloud IMDS
|
||||
"::1/128", // IPv6 loopback
|
||||
"fc00::/7", // IPv6 unique local
|
||||
"fe80::/10", // IPv6 link-local
|
||||
}
|
||||
|
||||
// reservedNetworks is the parsed form of reservedCIDRs, built once at startup.
|
||||
var reservedNetworks []*net.IPNet
|
||||
|
||||
func init() {
|
||||
for _, cidr := range reservedCIDRs {
|
||||
_, network, err := net.ParseCIDR(cidr)
|
||||
if err != nil {
|
||||
panic("webhook: invalid reserved CIDR " + cidr + ": " + err.Error())
|
||||
}
|
||||
reservedNetworks = append(reservedNetworks, network)
|
||||
}
|
||||
}
|
||||
|
||||
// AllowPrivateIPs controls whether webhook URLs may resolve to reserved/private
|
||||
// IP addresses. When true, the SSRF protection is disabled. This is useful for
|
||||
// self-hosted deployments where webhooks target services on the local network.
|
||||
var AllowPrivateIPs bool
|
||||
|
||||
// isReservedIP reports whether ip falls within any reserved/private range.
|
||||
func isReservedIP(ip net.IP) bool {
|
||||
if AllowPrivateIPs {
|
||||
return false
|
||||
}
|
||||
for _, network := range reservedNetworks {
|
||||
if network.Contains(ip) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ValidateURL checks that rawURL:
|
||||
// 1. Parses as a valid absolute URL.
|
||||
// 2. Uses the http or https scheme.
|
||||
// 3. Does not resolve to a reserved/private IP address.
|
||||
//
|
||||
// It returns a gRPC InvalidArgument status error so callers can return it directly.
|
||||
func ValidateURL(rawURL string) error {
|
||||
u, err := url.ParseRequestURI(rawURL)
|
||||
if err != nil {
|
||||
return status.Errorf(codes.InvalidArgument, "invalid webhook URL: %v", err)
|
||||
}
|
||||
if u.Scheme != "http" && u.Scheme != "https" {
|
||||
return status.Errorf(codes.InvalidArgument, "webhook URL must use http or https scheme, got %q", u.Scheme)
|
||||
}
|
||||
|
||||
ips, err := net.LookupHost(u.Hostname())
|
||||
if err != nil {
|
||||
return status.Errorf(codes.InvalidArgument, "webhook URL hostname could not be resolved: %v", err)
|
||||
}
|
||||
|
||||
for _, ipStr := range ips {
|
||||
ip := net.ParseIP(ipStr)
|
||||
if ip != nil && isReservedIP(ip) {
|
||||
return status.Errorf(codes.InvalidArgument, "webhook URL must not resolve to a reserved or private IP address")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateSigningSecret checks that secret is either empty (allowed) or contains
|
||||
// only printable ASCII characters (0x20–0x7E), excluding all control characters
|
||||
// such as \r and \n, which would corrupt the webhook signature headers. When the
|
||||
// secret uses the Standard Webhooks "whsec_<base64>" serialization, the base64
|
||||
// body must decode cleanly so signing cannot silently fall back to the wrong key.
|
||||
func ValidateSigningSecret(secret string) error {
|
||||
if secret == "" {
|
||||
return nil
|
||||
}
|
||||
for _, r := range secret {
|
||||
if r < 0x20 || r > 0x7E {
|
||||
return status.Errorf(codes.InvalidArgument, "signing secret contains invalid character")
|
||||
}
|
||||
}
|
||||
if _, err := resolveSigningKey(secret); err != nil {
|
||||
return status.Errorf(codes.InvalidArgument, "%v", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
package webhook
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/hmac"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/pkg/errors"
|
||||
|
||||
v1pb "github.com/usememos/memos/proto/gen/api/v1"
|
||||
)
|
||||
|
||||
var (
|
||||
// timeout is the timeout for webhook request. Default to 30 seconds.
|
||||
timeout = 30 * time.Second
|
||||
|
||||
// safeClient is the shared HTTP client used for all webhook dispatches.
|
||||
// Its Transport guards against SSRF by blocking connections to reserved/private
|
||||
// IP addresses at dial time, which also defeats DNS rebinding attacks.
|
||||
safeClient = &http.Client{
|
||||
Timeout: timeout,
|
||||
Transport: &http.Transport{
|
||||
DialContext: safeDialContext,
|
||||
},
|
||||
}
|
||||
|
||||
asyncPostQueue = make(chan *WebhookRequestPayload, 128)
|
||||
)
|
||||
|
||||
func init() {
|
||||
for range 4 {
|
||||
go func() {
|
||||
for payload := range asyncPostQueue {
|
||||
if err := Post(payload); err != nil {
|
||||
slog.Warn("Failed to dispatch webhook asynchronously",
|
||||
slog.String("url", payload.URL),
|
||||
slog.String("activityType", payload.ActivityType),
|
||||
slog.Any("err", err))
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
// safeDialContext is a net.Dialer.DialContext replacement that resolves the target
|
||||
// hostname and rejects any address that falls within a reserved/private IP range.
|
||||
func safeDialContext(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
host, port, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return nil, errors.Errorf("webhook: invalid address %q", addr)
|
||||
}
|
||||
|
||||
ips, err := net.DefaultResolver.LookupHost(ctx, host)
|
||||
if err != nil {
|
||||
return nil, errors.Wrapf(err, "webhook: failed to resolve host %q", host)
|
||||
}
|
||||
|
||||
for _, ipStr := range ips {
|
||||
if ip := net.ParseIP(ipStr); ip != nil && isReservedIP(ip) {
|
||||
return nil, errors.Errorf("webhook: connection to reserved/private IP address is not allowed")
|
||||
}
|
||||
}
|
||||
|
||||
return (&net.Dialer{}).DialContext(ctx, network, net.JoinHostPort(host, port))
|
||||
}
|
||||
|
||||
type WebhookRequestPayload struct {
|
||||
// The target URL for the webhook request.
|
||||
URL string `json:"url"`
|
||||
// The type of activity that triggered this webhook.
|
||||
ActivityType string `json:"activityType"`
|
||||
// The resource name of the creator. Format: users/{user}
|
||||
Creator string `json:"creator"`
|
||||
// The memo that triggered this webhook (if applicable).
|
||||
Memo *v1pb.Memo `json:"memo"`
|
||||
// Optional signing secret for HMAC-SHA256 signature. Not serialized to JSON.
|
||||
SigningSecret string `json:"-"`
|
||||
}
|
||||
|
||||
// resolveSigningKey returns the raw HMAC key for a signing secret. Secrets using
|
||||
// the Standard Webhooks "whsec_<base64>" serialization are base64-decoded to their
|
||||
// raw bytes; any other secret is used as-is. It returns an error when a whsec_-prefixed
|
||||
// secret is not valid base64, so callers fail loudly instead of silently signing with
|
||||
// the wrong key (which would make every signature unverifiable by the receiver).
|
||||
func resolveSigningKey(secret string) ([]byte, error) {
|
||||
if rest, ok := strings.CutPrefix(secret, "whsec_"); ok {
|
||||
decoded, err := base64.StdEncoding.DecodeString(rest)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "signing secret has whsec_ prefix but is not valid base64")
|
||||
}
|
||||
return decoded, nil
|
||||
}
|
||||
return []byte(secret), nil
|
||||
}
|
||||
|
||||
// GenerateSigningSecret returns a new Standard Webhooks signing secret in the
|
||||
// "whsec_<base64>" form, backed by 32 cryptographically-random bytes — comfortably
|
||||
// within the spec's 24–64 byte range.
|
||||
func GenerateSigningSecret() (string, error) {
|
||||
buf := make([]byte, 32)
|
||||
if _, err := rand.Read(buf); err != nil {
|
||||
return "", errors.Wrap(err, "failed to read random bytes for signing secret")
|
||||
}
|
||||
return "whsec_" + base64.StdEncoding.EncodeToString(buf), nil
|
||||
}
|
||||
|
||||
// Post posts the message to webhook endpoint.
|
||||
func Post(requestPayload *WebhookRequestPayload) error {
|
||||
body, err := json.Marshal(requestPayload)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "failed to marshal webhook request to %s", requestPayload.URL)
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("POST", requestPayload.URL, bytes.NewBuffer(body))
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "failed to construct webhook request to %s", requestPayload.URL)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
if requestPayload.SigningSecret != "" {
|
||||
key, err := resolveSigningKey(requestPayload.SigningSecret)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "failed to derive signing key for webhook to %s", requestPayload.URL)
|
||||
}
|
||||
|
||||
msgID := "msg_" + uuid.New().String()
|
||||
timestamp := strconv.FormatInt(time.Now().Unix(), 10)
|
||||
|
||||
mac := hmac.New(sha256.New, key)
|
||||
mac.Write([]byte(msgID + "." + timestamp + "."))
|
||||
mac.Write(body)
|
||||
signature := base64.StdEncoding.EncodeToString(mac.Sum(nil))
|
||||
|
||||
req.Header.Set("webhook-id", msgID)
|
||||
req.Header.Set("webhook-timestamp", timestamp)
|
||||
req.Header.Set("webhook-signature", "v1,"+signature)
|
||||
}
|
||||
|
||||
resp, err := safeClient.Do(req)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "failed to post webhook to %s", requestPayload.URL)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
b, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "failed to read webhook response from %s", requestPayload.URL)
|
||||
}
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode > 299 {
|
||||
return errors.Errorf("failed to post webhook %s, status code: %d", requestPayload.URL, resp.StatusCode)
|
||||
}
|
||||
|
||||
response := &struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}{}
|
||||
if err := json.Unmarshal(b, response); err != nil {
|
||||
return errors.Wrapf(err, "failed to unmarshal webhook response from %s", requestPayload.URL)
|
||||
}
|
||||
|
||||
if response.Code != 0 {
|
||||
return errors.Errorf("receive error code sent by webhook server, code %d, msg: %s", response.Code, response.Message)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// PostAsync posts the message to webhook endpoint asynchronously.
|
||||
// It enqueues the request for bounded asynchronous dispatch and does not wait for the response.
|
||||
func PostAsync(requestPayload *WebhookRequestPayload) {
|
||||
if requestPayload == nil {
|
||||
slog.Warn("Dropped webhook dispatch because payload is nil")
|
||||
return
|
||||
}
|
||||
select {
|
||||
case asyncPostQueue <- requestPayload:
|
||||
default:
|
||||
slog.Warn("Dropped webhook dispatch because the async queue is full",
|
||||
slog.String("url", requestPayload.URL),
|
||||
slog.String("activityType", requestPayload.ActivityType))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package webhook
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestPostAsyncNilPayloadDoesNotPanic(t *testing.T) {
|
||||
require.NotPanics(t, func() {
|
||||
PostAsync(nil)
|
||||
})
|
||||
}
|
||||
|
||||
func TestResolveSigningKey(t *testing.T) {
|
||||
rawKey := []byte("0123456789abcdef")
|
||||
whsec := "whsec_" + base64.StdEncoding.EncodeToString(rawKey)
|
||||
|
||||
t.Run("plain secret used as-is", func(t *testing.T) {
|
||||
key, err := resolveSigningKey("my-plain-secret")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []byte("my-plain-secret"), key)
|
||||
})
|
||||
|
||||
t.Run("whsec_ prefix is base64-decoded", func(t *testing.T) {
|
||||
key, err := resolveSigningKey(whsec)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, rawKey, key)
|
||||
})
|
||||
|
||||
t.Run("whsec_ with invalid base64 fails loudly", func(t *testing.T) {
|
||||
_, err := resolveSigningKey("whsec_not!valid!base64!")
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGenerateSigningSecret(t *testing.T) {
|
||||
secret, err := GenerateSigningSecret()
|
||||
require.NoError(t, err)
|
||||
require.True(t, strings.HasPrefix(secret, "whsec_"), "generated secret must use the whsec_ prefix")
|
||||
require.NoError(t, ValidateSigningSecret(secret), "generated secret must pass validation")
|
||||
|
||||
key, err := resolveSigningKey(secret)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, key, 32, "generated secret must decode to 32 raw bytes")
|
||||
|
||||
other, err := GenerateSigningSecret()
|
||||
require.NoError(t, err)
|
||||
require.NotEqual(t, secret, other, "each generated secret must be unique")
|
||||
}
|
||||
|
||||
func TestValidateSigningSecret(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
secret string
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "empty is allowed", secret: "", wantErr: false},
|
||||
{name: "printable ascii", secret: "abcDEF123!@#", wantErr: false},
|
||||
{name: "valid whsec_", secret: "whsec_" + base64.StdEncoding.EncodeToString([]byte("key")), wantErr: false},
|
||||
{name: "newline rejected", secret: "abc\ndef", wantErr: true},
|
||||
{name: "carriage return rejected", secret: "abc\rdef", wantErr: true},
|
||||
{name: "tab rejected", secret: "abc\tdef", wantErr: true},
|
||||
{name: "non-ascii rejected", secret: "abc€def", wantErr: true},
|
||||
{name: "whsec_ with invalid base64 rejected", secret: "whsec_not!base64", wantErr: true},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := ValidateSigningSecret(tc.secret)
|
||||
if tc.wantErr {
|
||||
require.Error(t, err)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestPostSignsRequest verifies the end-to-end Standard Webhooks signature so a
|
||||
// receiver following the documented verification recipe will accept our requests.
|
||||
func TestPostSignsRequest(t *testing.T) {
|
||||
// httptest listens on 127.0.0.1, which the SSRF guard blocks by default.
|
||||
prev := AllowPrivateIPs
|
||||
AllowPrivateIPs = true
|
||||
defer func() { AllowPrivateIPs = prev }()
|
||||
|
||||
rawKey := []byte("0123456789abcdef0123456789abcdef")
|
||||
cases := []struct {
|
||||
name string
|
||||
secret string
|
||||
key []byte
|
||||
}{
|
||||
{name: "plain secret", secret: "plain-secret-value", key: []byte("plain-secret-value")},
|
||||
{name: "whsec_ secret", secret: "whsec_" + base64.StdEncoding.EncodeToString(rawKey), key: rawKey},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var gotID, gotTimestamp, gotSignature string
|
||||
var gotBody []byte
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotID = r.Header.Get("webhook-id")
|
||||
gotTimestamp = r.Header.Get("webhook-timestamp")
|
||||
gotSignature = r.Header.Get("webhook-signature")
|
||||
gotBody, _ = io.ReadAll(r.Body)
|
||||
_, _ = w.Write([]byte(`{"code":0}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
err := Post(&WebhookRequestPayload{
|
||||
URL: server.URL,
|
||||
ActivityType: "memos.memo.created",
|
||||
Creator: "users/1",
|
||||
SigningSecret: tc.secret,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.True(t, strings.HasPrefix(gotID, "msg_"), "webhook-id should be prefixed with msg_")
|
||||
require.NotEmpty(t, gotTimestamp)
|
||||
require.True(t, strings.HasPrefix(gotSignature, "v1,"), "signature should carry the v1 version tag")
|
||||
|
||||
// Recompute the signature the way a receiver would and confirm it matches.
|
||||
mac := hmac.New(sha256.New, tc.key)
|
||||
mac.Write([]byte(gotID + "." + gotTimestamp + "."))
|
||||
mac.Write(gotBody)
|
||||
want := base64.StdEncoding.EncodeToString(mac.Sum(nil))
|
||||
require.Equal(t, "v1,"+want, gotSignature)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestPostWithoutSecretSetsNoSignatureHeaders ensures unsigned webhooks stay unsigned.
|
||||
func TestPostWithoutSecretSetsNoSignatureHeaders(t *testing.T) {
|
||||
prev := AllowPrivateIPs
|
||||
AllowPrivateIPs = true
|
||||
defer func() { AllowPrivateIPs = prev }()
|
||||
|
||||
var hasSignatureHeaders bool
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hasSignatureHeaders = r.Header.Get("webhook-id") != "" ||
|
||||
r.Header.Get("webhook-timestamp") != "" ||
|
||||
r.Header.Get("webhook-signature") != ""
|
||||
_, _ = w.Write([]byte(`{"code":0}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
err := Post(&WebhookRequestPayload{
|
||||
URL: server.URL,
|
||||
ActivityType: "memos.memo.created",
|
||||
Creator: "users/1",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.False(t, hasSignatureHeaders, "no signature headers should be set when no secret is configured")
|
||||
}
|
||||
Reference in New Issue
Block a user