dnswatcher/internal/state/state.go

// Package state provides JSON file-based state persistence.
package state

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"slices"
	"sync"
	"time"

	"go.uber.org/fx"

	"sneak.berlin/go/dnswatcher/internal/config"
	"sneak.berlin/go/dnswatcher/internal/logger"
)

// filePermissions for the state file.
const filePermissions = 0o600

// dirPermissions for the data directory.
const dirPermissions = 0o700

// stateVersion is the current state file format version.
const stateVersion = 1

// Params contains dependencies for State.
type Params struct {
	fx.In

	Logger *logger.Logger
	Config *config.Config
}

// DomainState holds the monitoring state for an apex domain.
// NameserverAddresses holds the sorted addresses each nameserver's name
// resolves to, by nameserver name. A state file written before it
// existed loads with it nil.
type DomainState struct {
	Nameservers         []string            `json:"nameservers"`
	NameserverAddresses map[string][]string `json:"nameserverAddresses"`
	LastChecked         time.Time           `json:"lastChecked"`
}

// NameserverRecordState holds one NS's response for a hostname.
type NameserverRecordState struct {
	Records     map[string][]string `json:"records"`
	Status      string              `json:"status"`
	Error       string              `json:"error,omitempty"`
	LastChecked time.Time           `json:"lastChecked"`
}

// HostnameState holds per-nameserver monitoring state for a hostname.
// CNAMEAddresses holds the sorted addresses at the end of the name's
// CNAME chain, found when its nameservers answered with a CNAME and no
// address; it is empty otherwise. It is nil when they are not known: a
// state file written before it existed loads with it nil.
type HostnameState struct {
	RecordsByNameserver map[string]*NameserverRecordState `json:"recordsByNameserver"`
	CNAMEAddresses      []string                          `json:"cnameAddresses"`
	LastChecked         time.Time                         `json:"lastChecked"`
}

// PortState holds the monitoring state for a port.
type PortState struct {
	Open        bool      `json:"open"`
	Hostnames   []string  `json:"hostnames"`
	LastChecked time.Time `json:"lastChecked"`
}

// UnmarshalJSON implements custom unmarshaling to handle both
// the old single-hostname format and the new multi-hostname
// format for backward compatibility with existing state files.
func (ps *PortState) UnmarshalJSON(data []byte) error {
	// Use an alias to prevent infinite recursion.
	type portStateAlias struct {
		Open        bool      `json:"open"`
		Hostnames   []string  `json:"hostnames"`
		LastChecked time.Time `json:"lastChecked"`
	}

	var alias portStateAlias

	err := json.Unmarshal(data, &alias)
	if err != nil {
		return fmt.Errorf("unmarshaling port state: %w", err)
	}

	ps.Open = alias.Open
	ps.Hostnames = alias.Hostnames
	ps.LastChecked = alias.LastChecked

	// If Hostnames is empty, try reading the old single-hostname
	// format for backward compatibility.
	if len(ps.Hostnames) == 0 {
		var old struct {
			Hostname string `json:"hostname"`
		}

		// Best-effort: ignore errors since the main unmarshal
		// already succeeded.
		if json.Unmarshal(data, &old) == nil && old.Hostname != "" {
			ps.Hostnames = []string{old.Hostname}
		}
	}

	return nil
}

// CertificateState holds TLS certificate monitoring state.
type CertificateState struct {
	CommonName              string    `json:"commonName"`
	Issuer                  string    `json:"issuer"`
	NotAfter                time.Time `json:"notAfter"`
	SubjectAlternativeNames []string  `json:"subjectAlternativeNames"`
	Status                  string    `json:"status"`
	Error                   string    `json:"error,omitempty"`
	LastChecked             time.Time `json:"lastChecked"`
}

// Snapshot is the complete monitoring state persisted to disk.
type Snapshot struct {
	Version      int                          `json:"version"`
	LastUpdated  time.Time                    `json:"lastUpdated"`
	Domains      map[string]*DomainState      `json:"domains"`
	Hostnames    map[string]*HostnameState    `json:"hostnames"`
	Ports        map[string]*PortState        `json:"ports"`
	Certificates map[string]*CertificateState `json:"certificates"`
}

// State manages the monitoring state with file persistence.
type State struct {
	mu       sync.RWMutex
	snapshot *Snapshot
	log      *slog.Logger
	config   *config.Config
}

// New creates a new State instance and loads existing state from disk.
func New(
	lifecycle fx.Lifecycle,
	params Params,
) (*State, error) {
	state := &State{
		log:    params.Logger.Get(),
		config: params.Config,
		snapshot: &Snapshot{
			Version:      stateVersion,
			Domains:      make(map[string]*DomainState),
			Hostnames:    make(map[string]*HostnameState),
			Ports:        make(map[string]*PortState),
			Certificates: make(map[string]*CertificateState),
		},
	}

	lifecycle.Append(fx.Hook{
		OnStart: func(_ context.Context) error {
			err := state.checkDataDirWritable()
			if err != nil {
				return err
			}

			return state.Load()
		},
		OnStop: func(_ context.Context) error {
			return state.Save()
		},
	})

	return state, nil
}

// Load reads the state from disk.
func (s *State) Load() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	path := s.config.StatePath()

	//nolint:gosec // path is from trusted config
	data, err := os.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			s.log.Info(
				"no existing state file, starting fresh",
				"path", path,
			)

			return nil
		}

		return fmt.Errorf("reading state file: %w", err)
	}

	var snapshot Snapshot

	err = json.Unmarshal(data, &snapshot)
	if err != nil {
		return fmt.Errorf("parsing state file: %w", err)
	}

	// A state file saved before each record value was stored once can
	// hold a hostname's CNAME once for every record type asked for.
	// Each value is kept once, so the first check does not see a
	// record change.
	for _, hs := range snapshot.Hostnames {
		for _, ns := range hs.RecordsByNameserver {
			for recordType, values := range ns.Records {
				slices.Sort(values)
				ns.Records[recordType] = slices.Compact(values)
			}
		}
	}

	s.snapshot = &snapshot
	s.log.Info("loaded state from disk", "path", path)

	return nil
}

// Save writes the current state to disk atomically.
func (s *State) Save() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.snapshot.LastUpdated = time.Now().UTC()

	data, err := json.MarshalIndent(s.snapshot, "", "  ")
	if err != nil {
		return fmt.Errorf("marshaling state: %w", err)
	}

	path := s.config.StatePath()

	err = os.MkdirAll(filepath.Dir(path), dirPermissions)
	if err != nil {
		return fmt.Errorf("creating data directory: %w", err)
	}

	// Atomic write: write to temp file, then rename
	tmpPath := path + ".tmp"

	err = os.WriteFile(tmpPath, data, filePermissions)
	if err != nil {
		return fmt.Errorf("writing temp state file: %w", err)
	}

	err = os.Rename(tmpPath, path)
	if err != nil {
		return fmt.Errorf("renaming state file: %w", err)
	}

	s.log.Debug("state saved to disk", "path", path)

	return nil
}

// GetSnapshot returns a copy of the current snapshot.
func (s *State) GetSnapshot() Snapshot {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return *s.snapshot
}

// SetDomainState updates the state for a domain.
func (s *State) SetDomainState(
	domain string,
	ds *DomainState,
) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.snapshot.Domains[domain] = ds
}

// GetDomainState returns the state for a domain.
func (s *State) GetDomainState(
	domain string,
) (*DomainState, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	ds, ok := s.snapshot.Domains[domain]

	return ds, ok
}

// SetHostnameState updates the state for a hostname.
func (s *State) SetHostnameState(
	hostname string,
	hs *HostnameState,
) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.snapshot.Hostnames[hostname] = hs
}

// GetHostnameState returns the state for a hostname.
func (s *State) GetHostnameState(
	hostname string,
) (*HostnameState, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	hs, ok := s.snapshot.Hostnames[hostname]

	return hs, ok
}

// SetPortState updates the state for a port.
func (s *State) SetPortState(key string, ps *PortState) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.snapshot.Ports[key] = ps
}

// GetPortState returns the state for a port.
func (s *State) GetPortState(key string) (*PortState, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	ps, ok := s.snapshot.Ports[key]

	return ps, ok
}

// DeletePortState removes a port state entry.
func (s *State) DeletePortState(key string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	delete(s.snapshot.Ports, key)
}

// GetAllPortKeys returns all port state keys.
func (s *State) GetAllPortKeys() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()

	keys := make([]string, 0, len(s.snapshot.Ports))
	for k := range s.snapshot.Ports {
		keys = append(keys, k)
	}

	return keys
}

// SetCertificateState updates the state for a certificate.
func (s *State) SetCertificateState(
	key string,
	cs *CertificateState,
) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.snapshot.Certificates[key] = cs
}

// GetCertificateState returns the state for a certificate.
func (s *State) GetCertificateState(
	key string,
) (*CertificateState, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	cs, ok := s.snapshot.Certificates[key]

	return cs, ok
}

// checkDataDirWritable creates the data directory if needed, then writes
// and removes the temp file that Save uses. It runs at startup so that an
// unwritable directory stops the process, instead of the process running
// with every save failing and only logged.
func (s *State) checkDataDirWritable() error {
	dir := s.config.DataDir
	tmpPath := s.config.StatePath() + ".tmp"

	err := os.MkdirAll(dir, dirPermissions)
	if err == nil {
		err = os.WriteFile(tmpPath, nil, filePermissions)
	}

	if err == nil {
		err = os.Remove(tmpPath)
	}

	if err != nil {
		return fmt.Errorf("data directory %s is not writable: %w", dir, err)
	}

	return nil
}