1package policy
  2
  3import (
  4	"context"
  5	"crypto/sha256"
  6	"fmt"
  7	"log"
  8	"log/slog"
  9	"net/netip"
 10	"os"
 11	"sync/atomic"
 12	"time"
 13
 14	"postern/internal/db"
 15
 16	"go.starlark.net/lib/json"
 17	sltime "go.starlark.net/lib/time"
 18	"go.starlark.net/starlark"
 19	"go.starlark.net/starlarkstruct"
 20	"go.starlark.net/syntax"
 21)
 22
 23type RelayConfig struct {
 24	Host     string
 25	Port     int
 26	Username string
 27	Password string
 28}
 29
 30// engineState holds the reloadable, per-load Starlark environment for a
 31// single evaluation of policy.star. The active snapshot is swapped in
 32// atomically so in-flight hook calls keep a consistent view.
 33type engineState struct {
 34	globals starlark.StringDict
 35	thread  *starlark.Thread
 36}
 37
 38type Engine struct {
 39	Listeners  []ListenerConfig
 40	policyFile string
 41	database   *db.DB
 42	state      atomic.Pointer[engineState]
 43}
 44
 45type ListenerConfig struct {
 46	Address     netip.AddrPort
 47	Role        string
 48	TLSCertFile string
 49	TLSKeyFile  string
 50}
 51
 52func NewEngine(policyFile string, database *db.DB) (*Engine, error) {
 53	e := &Engine{
 54		policyFile: policyFile,
 55		database:   database,
 56	}
 57
 58	state, listeners, err := e.load()
 59	if err != nil {
 60		return nil, err
 61	}
 62	e.Listeners = listeners
 63	e.state.Store(state)
 64	return e, nil
 65}
 66
 67// load reads and evaluates policyFile, returning a fresh engineState along
 68// with the parsed listeners produced by this evaluation.
 69func (e *Engine) load() (*engineState, []ListenerConfig, error) {
 70	script, err := os.ReadFile(e.policyFile)
 71	if err != nil {
 72		log.Fatal(err)
 73	}
 74
 75	// Create an 'os' module containing our getenv function
 76	osModule := &starlarkstruct.Module{
 77		Name: "os",
 78		Members: starlark.StringDict{
 79			"getenv": starlark.NewBuiltin("getenv", starlarkGetenv),
 80		},
 81	}
 82
 83	// Expose the "os", "mda", "users", and "http" modules to the Starlark global environment
 84	m := &MDA{
 85		parsedListeners: make([]ListenerConfig, 0),
 86	}
 87
 88	predeclared := starlark.StringDict{
 89		"exec":  &ExecModule{},
 90		"http":  &HTTPModule{},
 91		"json":  json.Module,
 92		"kv":    NewKVModule(e.database),
 93		"mda":   m,
 94		"os":    osModule,
 95		"time":  sltime.Module,
 96		"users": NewUsersModule(e.database),
 97	}
 98
 99	// Create a Starlark thread. We customize the Print function
100	// so the Starlark `print()` built-in logs cleanly in Go.
101	thread := &starlark.Thread{
102		Name: "mda-main",
103		Print: func(_ *starlark.Thread, msg string) {
104			fmt.Printf("[init.star]: %s\n", msg)
105		},
106	}
107	thread.SetLocal("context", context.Background())
108
109	// 1. Evaluate the script
110	globals, err := starlark.ExecFileOptions(&syntax.FileOptions{}, thread, "policy.star", script, predeclared)
111	if err != nil {
112		return nil, nil, err
113	}
114
115	// 2. Call the init() function
116	initVal, ok := globals["init"]
117	if !ok {
118		return nil, nil, fmt.Errorf("init not found in policy")
119	}
120
121	_, err = starlark.Call(thread, initVal, nil, nil)
122	if err != nil {
123		return nil, nil, err
124	}
125
126	state := &engineState{
127		globals: globals,
128		thread:  thread,
129	}
130	return state, m.parsedListeners, nil
131}
132
133// reload re-evaluates policyFile and, on success, atomically swaps the new
134// engineState into the engine. On any error the previous state is kept
135// untouched so the last-known-good policy continues to serve.
136func (e *Engine) reload() error {
137	state, _, err := e.load()
138	if err != nil {
139		return err
140	}
141	e.state.Store(state)
142	return nil
143}
144
145// StartHotReload spawns a goroutine that polls the policy file's content hash
146// on the given interval and reloads whenever it changes. It is safe to call
147// once at startup; the goroutine exits when ctx is canceled.
148func (e *Engine) StartHotReload(ctx context.Context, interval time.Duration) {
149	go func() {
150		ticker := time.NewTicker(interval)
151		defer ticker.Stop()
152
153		var lastHash [32]byte
154		if content, err := os.ReadFile(e.policyFile); err == nil {
155			lastHash = sha256.Sum256(content)
156		}
157
158		for {
159			select {
160			case <-ctx.Done():
161				return
162			case <-ticker.C:
163				content, err := os.ReadFile(e.policyFile)
164				if err != nil {
165					slog.Error("policy file read failed", "file", e.policyFile, "err", err)
166					continue
167				}
168				currentHash := sha256.Sum256(content)
169				if currentHash == lastHash {
170					continue
171				}
172				lastHash = currentHash
173				slog.Info("policy file changed, reloading", "file", e.policyFile)
174				err = e.reload()
175				if err != nil {
176					slog.Error("failed to reload", "file", e.policyFile, "err", err)
177				} else {
178					slog.Info("policy reloaded", "file", e.policyFile)
179				}
180			}
181		}
182	}()
183}
184
185func (e *Engine) OnMessageSubmit(user User, msg *Message) (RelayConfig, error) {
186	s := e.state.Load()
187	hookVal, ok := s.globals["on_message_submit"]
188	if !ok {
189		return RelayConfig{}, nil
190	}
191
192	result, err := starlark.Call(s.thread, hookVal, starlark.Tuple{user, msg}, nil)
193	if err != nil {
194		return RelayConfig{}, err
195	}
196
197	if result == starlark.None {
198		return RelayConfig{}, nil
199	}
200
201	dict, ok := result.(*starlark.Dict)
202	if !ok {
203		return RelayConfig{}, fmt.Errorf("on_message_submit must return a dict, got %s", result.Type())
204	}
205
206	relay := RelayConfig{
207		Host:     dictString(dict, "host"),
208		Port:     dictInt(dict, "port"),
209		Username: dictString(dict, "username"),
210		Password: dictString(dict, "password"),
211	}
212	return relay, nil
213}
214
215func (e *Engine) OnMessageDeliver(conn Conn, user User, msg *Message) (string, error) {
216	defaultMailbox := "INBOX"
217	s := e.state.Load()
218	hookVal, ok := s.globals["on_message_deliver"]
219	if !ok {
220		return defaultMailbox, nil
221	}
222
223	result, err := starlark.Call(s.thread, hookVal, starlark.Tuple{conn, user, msg}, nil)
224	if err != nil {
225		return defaultMailbox, err
226	}
227
228	mailbox := defaultMailbox
229	if result != starlark.None {
230		if s, ok := result.(starlark.String); ok {
231			mailbox = string(s)
232		}
233	}
234
235	return mailbox, nil
236}
237
238func (e *Engine) OnMessageMoved(user User, msg *Message, sourceFolder, destFolder string) error {
239	s := e.state.Load()
240	hookVal, ok := s.globals["on_message_moved"]
241	if !ok {
242		return nil
243	}
244
245	_, err := starlark.Call(s.thread, hookVal, starlark.Tuple{user, msg, starlark.String(sourceFolder), starlark.String(destFolder)}, nil)
246	return err
247}
248
249// starlarkGetenv implements Python's os.getenv("KEY", "default")
250func starlarkGetenv(_ *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
251	var key string
252	var def starlark.Value = starlark.None
253
254	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "key", &key, "default?", &def); err != nil {
255		return nil, err
256	}
257
258	if val, exists := os.LookupEnv(key); exists {
259		return starlark.String(val), nil
260	}
261	return def, nil
262}
263
264func dictString(dict *starlark.Dict, key string) string {
265	v, _, err := dict.Get(starlark.String(key))
266	if err != nil || v == starlark.None {
267		return ""
268	}
269	if s, ok := v.(starlark.String); ok {
270		return string(s)
271	}
272	return ""
273}
274
275func dictInt(dict *starlark.Dict, key string) int {
276	v, _, err := dict.Get(starlark.String(key))
277	if err != nil || v == starlark.None {
278		return 0
279	}
280	if i, ok := v.(starlark.Int); ok {
281		if n, ok := i.Int64(); ok {
282			return int(n)
283		}
284	}
285	return 0
286}