Add Phase 1 LAN bind and token-guarded agent API

2026-05-18 11:02:04 +08:00
parent 1adabd2c84
commit 2eb9fd9906
2 changed files with 455 additions and 2 deletions
--- a/agent/main.go
+++ b/agent/main.go
@@ -0,0 +1,387 @@
+package main
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log"
+	"net"
+	"net/http"
+	"net/netip"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+const brokerURL = "wss://memory.clouderized.com/agent"
+
+var defaultAllowedSubnets = []string{
+	"10.0.0.0/8",
+	"172.16.0.0/12",
+	"192.168.0.0/16",
+	"127.0.0.0/8",
+	"::1/128",
+	"fc00::/7",
+	"fe80::/10",
+}
+
+type Config struct {
+	StorageDir     string   `json:"storage_dir"`
+	BindHost       string   `json:"bind_host"`
+	BindPort       int      `json:"bind_port"`
+	AllowedSubnets []string `json:"allowed_subnets"`
+	PairingToken   string   `json:"pairing_token"`
+}
+
+type Message struct {
+	Type      string `json:"type"`
+	Filename  string `json:"filename,omitempty"`
+	Data      string `json:"data,omitempty"`
+	RequestID string `json:"request_id,omitempty"`
+	Error     string `json:"error,omitempty"`
+}
+
+type AgentState struct {
+	mu              sync.RWMutex
+	storageDir      string
+	connected       bool
+	lastError       string
+	uploadsHandled  int64
+	fetchesHandled  int64
+	startedAt       time.Time
+	allowedPrefixes []netip.Prefix
+	token           string
+	bindAddr        string
+}
+
+type StatusResponse struct {
+	Status          string   `json:"status"`
+	BrokerConnected bool     `json:"broker_connected"`
+	StorageDir      string   `json:"storage_dir"`
+	BindAddr        string   `json:"bind_addr"`
+	AllowedSubnets  []string `json:"allowed_subnets"`
+	UploadsHandled  int64    `json:"uploads_handled"`
+	FetchesHandled  int64    `json:"fetches_handled"`
+	LastError       string   `json:"last_error,omitempty"`
+	StartedAt       string   `json:"started_at"`
+}
+
+func defaultStorageDir() string {
+	home, err := os.UserHomeDir()
+	if err != nil || home == "" {
+		return "./photos"
+	}
+	return filepath.Join(home, "memoric", "photos")
+}
+
+func configPath() string {
+	d, err := os.UserConfigDir()
+	if err != nil || d == "" {
+		return "memoric-agent.json"
+	}
+	return filepath.Join(d, "memoric-agent", "config.json")
+}
+
+func loadConfig() Config {
+	cfg := Config{
+		StorageDir:     defaultStorageDir(),
+		BindHost:       "0.0.0.0",
+		BindPort:       5051,
+		AllowedSubnets: append([]string(nil), defaultAllowedSubnets...),
+	}
+
+	if env := strings.TrimSpace(os.Getenv("STORAGE_DIR")); env != "" {
+		cfg.StorageDir = env
+	}
+	if env := strings.TrimSpace(os.Getenv("AGENT_BIND_HOST")); env != "" {
+		cfg.BindHost = env
+	}
+	if env := strings.TrimSpace(os.Getenv("AGENT_BIND_PORT")); env != "" {
+		var port int
+		if _, err := fmt.Sscanf(env, "%d", &port); err == nil && port > 0 && port < 65536 {
+			cfg.BindPort = port
+		}
+	}
+	if env := strings.TrimSpace(os.Getenv("AGENT_ALLOWED_SUBNETS")); env != "" {
+		parts := strings.Split(env, ",")
+		var out []string
+		for _, p := range parts {
+			p = strings.TrimSpace(p)
+			if p != "" {
+				out = append(out, p)
+			}
+		}
+		if len(out) > 0 {
+			cfg.AllowedSubnets = out
+		}
+	}
+	if env := strings.TrimSpace(os.Getenv("AGENT_PAIRING_TOKEN")); env != "" {
+		cfg.PairingToken = env
+	}
+
+	b, err := os.ReadFile(configPath())
+	if err != nil {
+		return cfg
+	}
+	var fileCfg Config
+	if err := json.Unmarshal(b, &fileCfg); err != nil {
+		log.Printf("config parse error (%s): %v", configPath(), err)
+		return cfg
+	}
+	if strings.TrimSpace(fileCfg.StorageDir) != "" {
+		cfg.StorageDir = strings.TrimSpace(fileCfg.StorageDir)
+	}
+	if strings.TrimSpace(fileCfg.BindHost) != "" {
+		cfg.BindHost = strings.TrimSpace(fileCfg.BindHost)
+	}
+	if fileCfg.BindPort > 0 && fileCfg.BindPort < 65536 {
+		cfg.BindPort = fileCfg.BindPort
+	}
+	if len(fileCfg.AllowedSubnets) > 0 {
+		cfg.AllowedSubnets = fileCfg.AllowedSubnets
+	}
+	if strings.TrimSpace(fileCfg.PairingToken) != "" {
+		cfg.PairingToken = strings.TrimSpace(fileCfg.PairingToken)
+	}
+	return cfg
+}
+
+func parseAllowedSubnets(values []string) ([]netip.Prefix, error) {
+	prefixes := make([]netip.Prefix, 0, len(values))
+	for _, raw := range values {
+		p, err := netip.ParsePrefix(strings.TrimSpace(raw))
+		if err != nil {
+			return nil, fmt.Errorf("invalid subnet %q: %w", raw, err)
+		}
+		prefixes = append(prefixes, p)
+	}
+	if len(prefixes) == 0 {
+		return nil, errors.New("allowed_subnets must not be empty")
+	}
+	return prefixes, nil
+}
+
+func (s *AgentState) setConnected(v bool) {
+	s.mu.Lock()
+	s.connected = v
+	s.mu.Unlock()
+}
+
+func (s *AgentState) setError(err error) {
+	if err == nil {
+		return
+	}
+	s.mu.Lock()
+	s.lastError = err.Error()
+	s.mu.Unlock()
+}
+
+func (s *AgentState) incUpload() {
+	s.mu.Lock()
+	s.uploadsHandled++
+	s.mu.Unlock()
+}
+
+func (s *AgentState) incFetch() {
+	s.mu.Lock()
+	s.fetchesHandled++
+	s.mu.Unlock()
+}
+
+func (s *AgentState) status(allowed []string) StatusResponse {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return StatusResponse{
+		Status:          "ok",
+		BrokerConnected: s.connected,
+		StorageDir:      s.storageDir,
+		BindAddr:        s.bindAddr,
+		AllowedSubnets:  append([]string(nil), allowed...),
+		UploadsHandled:  s.uploadsHandled,
+		FetchesHandled:  s.fetchesHandled,
+		LastError:       s.lastError,
+		StartedAt:       s.startedAt.Format(time.RFC3339),
+	}
+}
+
+func (s *AgentState) writeJSON(w http.ResponseWriter, code int, v any) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(code)
+	_ = json.NewEncoder(w).Encode(v)
+}
+
+func (s *AgentState) clientAllowed(r *http.Request) bool {
+	host, _, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		host = r.RemoteAddr
+	}
+	ip, err := netip.ParseAddr(strings.Trim(host, "[]"))
+	if err != nil {
+		return false
+	}
+	for _, p := range s.allowedPrefixes {
+		if p.Contains(ip) {
+			return true
+		}
+	}
+	return false
+}
+
+func (s *AgentState) requireSubnet(next http.HandlerFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		if !s.clientAllowed(r) {
+			s.writeJSON(w, http.StatusForbidden, map[string]string{"error": "client not in allowed_subnets"})
+			return
+		}
+		next(w, r)
+	}
+}
+
+func (s *AgentState) requireMutableAuth(next http.HandlerFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		token := strings.TrimSpace(s.token)
+		if token == "" {
+			s.writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "mutable endpoints disabled: pairing token not set"})
+			return
+		}
+		auth := strings.TrimSpace(r.Header.Get("Authorization"))
+		xToken := strings.TrimSpace(r.Header.Get("X-Agent-Token"))
+		if strings.HasPrefix(strings.ToLower(auth), "bearer ") {
+			auth = strings.TrimSpace(auth[7:])
+		}
+		if auth != token && xToken != token {
+			s.writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unauthorized"})
+			return
+		}
+		next(w, r)
+	}
+}
+
+func (s *AgentState) startHTTPServer(allowedSubnets []string) {
+	mux := http.NewServeMux()
+
+	mux.HandleFunc("/healthz", s.requireSubnet(func(w http.ResponseWriter, r *http.Request) {
+		s.writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
+	}))
+
+	mux.HandleFunc("/api/status", s.requireSubnet(func(w http.ResponseWriter, r *http.Request) {
+		s.writeJSON(w, http.StatusOK, s.status(allowedSubnets))
+	}))
+
+	// Mutable placeholder endpoint for Phase 1 auth enforcement.
+	mux.HandleFunc("/api/control/ping", s.requireSubnet(s.requireMutableAuth(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			s.writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
+			return
+		}
+		s.writeJSON(w, http.StatusOK, map[string]string{"status": "ok", "message": "authorized mutable endpoint"})
+	})))
+
+	srv := &http.Server{
+		Addr:              s.bindAddr,
+		Handler:           mux,
+		ReadHeaderTimeout: 5 * time.Second,
+	}
+	log.Printf("agent UI/API listening on http://%s", s.bindAddr)
+	if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+		s.setError(err)
+		log.Fatalf("http server: %v", err)
+	}
+}
+
+func main() {
+	cfg := loadConfig()
+	prefixes, err := parseAllowedSubnets(cfg.AllowedSubnets)
+	if err != nil {
+		log.Fatalf("subnet config error: %v", err)
+	}
+
+	if err := os.MkdirAll(cfg.StorageDir, 0755); err != nil {
+		log.Fatal(err)
+	}
+
+	state := &AgentState{
+		storageDir:      cfg.StorageDir,
+		startedAt:       time.Now().UTC(),
+		allowedPrefixes: prefixes,
+		token:           strings.TrimSpace(cfg.PairingToken),
+		bindAddr:        fmt.Sprintf("%s:%d", cfg.BindHost, cfg.BindPort),
+	}
+
+	log.Println("photo storage directory:", cfg.StorageDir)
+	log.Println("allowed subnets:", strings.Join(cfg.AllowedSubnets, ", "))
+	if state.token == "" {
+		log.Println("pairing token not set: mutable endpoints disabled")
+	}
+
+	go state.startHTTPServer(cfg.AllowedSubnets)
+
+	conn, _, err := websocket.DefaultDialer.Dial(brokerURL, nil)
+	if err != nil {
+		log.Fatal("connect:", err)
+	}
+	defer conn.Close()
+	state.setConnected(true)
+	log.Println("connected to broker at", brokerURL)
+
+	for {
+		_, raw, err := conn.ReadMessage()
+		if err != nil {
+			state.setConnected(false)
+			state.setError(err)
+			log.Fatal("read:", err)
+		}
+
+		var msg Message
+		if err := json.Unmarshal(raw, &msg); err != nil {
+			state.setError(err)
+			log.Println("bad message:", err)
+			continue
+		}
+
+		switch msg.Type {
+		case "upload":
+			data, err := base64.StdEncoding.DecodeString(msg.Data)
+			if err != nil {
+				state.setError(err)
+				log.Println("decode error:", err)
+				continue
+			}
+			path := filepath.Join(cfg.StorageDir, msg.Filename)
+			if err := os.WriteFile(path, data, 0644); err != nil {
+				state.setError(err)
+				log.Println("write error:", err)
+				continue
+			}
+			state.incUpload()
+			log.Printf("saved %s (%d bytes)", msg.Filename, len(data))
+
+		case "fetch":
+			path := filepath.Join(cfg.StorageDir, msg.Filename)
+			data, err := os.ReadFile(path)
+			if err != nil {
+				log.Printf("fetch %s: not found", msg.Filename)
+				resp, _ := json.Marshal(Message{
+					Type:      "fetch_response",
+					RequestID: msg.RequestID,
+					Error:     "not found",
+				})
+				conn.WriteMessage(websocket.TextMessage, resp)
+				continue
+			}
+			resp, _ := json.Marshal(Message{
+				Type:      "fetch_response",
+				RequestID: msg.RequestID,
+				Data:      base64.StdEncoding.EncodeToString(data),
+			})
+			conn.WriteMessage(websocket.TextMessage, resp)
+			state.incFetch()
+			log.Printf("served %s (%d bytes)", msg.Filename, len(data))
+		}
+	}
+}