Merge pull request #89 from sonroyaalmerol/improve-errors

iron out some bugs with websockets and improve store mutex locking

cmd/pbs_plus/main.go

@@ -220,6 +220,7 @@ func main() {
 
 	// Agent auth routes
 	mux.HandleFunc("/plus/agent/bootstrap", mw.CORS(storeInstance, agents.AgentBootstrapHandler(storeInstance)))
+	mux.HandleFunc("/plus/agent/renew", mw.AgentOnly(storeInstance, mw.CORS(storeInstance, agents.AgentRenewHandler(storeInstance))))
 
 	server := &http.Server{
 		Addr:           serverConfig.Address,

cmd/windows_agent/main.go

@@ -14,6 +14,7 @@ import (
 	_ "net/http/pprof"
 
 	"github.com/kardianos/service"
+	"github.com/sonroyaalmerol/pbs-plus/internal/store/constants"
 	"github.com/sonroyaalmerol/pbs-plus/internal/syslog"
 	"golang.org/x/sys/windows/registry"
 )
@@ -90,6 +91,8 @@ func (w *watchdogService) Stop(s service.Service) error {
 }
 
 func main() {
+	constants.Version = Version
+
 	go func() {
 		log.Println(http.ListenAndServe("localhost:6060", nil))
 	}()

cmd/windows_agent/service.go

@@ -50,11 +50,25 @@ func (p *agentService) Start(s service.Service) error {
 	p.svc = s
 	p.ctx, p.cancel = context.WithCancel(context.Background())
 
-	p.wg.Add(1)
+	p.wg.Add(2)
 	go func() {
 		defer p.wg.Done()
 		p.run()
 	}()
+	go func() {
+		defer p.wg.Done()
+		for {
+			select {
+			case <-p.ctx.Done():
+				return
+			case <-time.After(time.Hour):
+				err := agent.CheckAndRenewCertificate()
+				if err != nil {
+					syslog.L.Errorf("Certificate renewal manager: %w", err)
+				}
+			}
+		}
+	}()
 
 	return nil
 }
@@ -110,7 +124,7 @@ func (p *agentService) waitForServerURL() error {
 }
 
 func (p *agentService) waitForBootstrap() error {
-	ticker := time.NewTicker(5 * time.Second)
+	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
 
 	for {
@@ -119,10 +133,16 @@ func (p *agentService) waitForBootstrap() error {
 		priv, _ := registry.GetEntry(registry.AUTH, "Priv", true)
 
 		if serverCA != nil && cert != nil && priv != nil {
-			return nil
+			err := agent.CheckAndRenewCertificate()
+			if err == nil {
+				return nil
+			}
+			syslog.L.Errorf("Renewal error: %v", err)
 		} else {
 			err := agent.Bootstrap()
-			syslog.L.Errorf("Bootstrap error: %v", err)
+			if err != nil {
+				syslog.L.Errorf("Bootstrap error: %v", err)
+			}
 		}
 
 		select {
@@ -201,27 +221,7 @@ func (p *agentService) connectWebSocket() error {
 			return err
 		}
 
-		client, err := websockets.NewWSClient(p.ctx, config, tlsConfig)
-		if err != nil {
-			syslog.L.Errorf("WS client init error: %s", err)
-			select {
-			case <-p.ctx.Done():
-				return fmt.Errorf("context cancelled while connecting to WebSocket")
-			case <-time.After(5 * time.Second):
-				continue
-			}
-		}
-
-		err = client.Connect()
-		if err != nil {
-			syslog.L.Errorf("WS client connect error: %s", err)
-			select {
-			case <-p.ctx.Done():
-				return fmt.Errorf("context cancelled while connecting to WebSocket")
-			case <-time.After(5 * time.Second):
-				continue
-			}
-		}
+		client := websockets.NewWSClient(p.ctx, config, tlsConfig)
 
 		client.RegisterHandler("backup_start", controllers.BackupStartHandler(client))
 		client.RegisterHandler("backup_close", controllers.BackupCloseHandler(client))

internal/agent/http.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/sonroyaalmerol/pbs-plus/internal/agent/registry"
+	"github.com/sonroyaalmerol/pbs-plus/internal/store/constants"
 )
 
 var httpClient *http.Client
@@ -40,6 +41,7 @@ func ProxmoxHTTPRequest(method, url string, body io.Reader, respBody any) (io.Re
 
 	req.Header.Add("Content-Type", "application/json")
 	req.Header.Add("X-PBS-Agent", hostname)
+	req.Header.Add("X-PBS-Plus-Version", constants.Version)
 
 	tlsConfig, err := GetTLSConfig()
 	if err != nil {

internal/agent/status_manager.go

@@ -9,7 +9,7 @@ import (
 // BackupStatus manages the state of ongoing backups
 type BackupStatus struct {
 	activeBackups sync.Map
-	mu             sync.RWMutex
+	mu            sync.RWMutex
 }
 
 var (
@@ -44,5 +44,3 @@ func (bs *BackupStatus) HasActiveBackups() bool {
 	})
 	return hasActive
 }
-
-

internal/agent/tls_config.go

@@ -3,11 +3,20 @@
 package agent
 
 import (
+	"bytes"
 	"crypto/tls"
 	"crypto/x509"
+	"encoding/base64"
+	"encoding/json"
+	"encoding/pem"
 	"fmt"
+	"net/http"
+	"os"
+	"time"
 
 	"github.com/sonroyaalmerol/pbs-plus/internal/agent/registry"
+	"github.com/sonroyaalmerol/pbs-plus/internal/auth/certificates"
+	"github.com/sonroyaalmerol/pbs-plus/internal/utils"
 )
 
 func GetTLSConfig() (*tls.Config, error) {
@@ -45,3 +54,116 @@ func GetTLSConfig() (*tls.Config, error) {
 		RootCAs:      rootCAs,
 	}, nil
 }
+
+func CheckAndRenewCertificate() error {
+	const renewalWindow = 30 * 24 * time.Hour // Renew if certificate expires in less than 30 days
+
+	certReg, err := registry.GetEntry(registry.AUTH, "Cert", true)
+	if err != nil {
+		return fmt.Errorf("CheckAndRenewCertificate: failed to retrieve certificate - %w", err)
+	}
+
+	block, _ := pem.Decode([]byte(certReg.Value))
+	if block == nil {
+		return fmt.Errorf("CheckAndRenewCertificate: failed to decode PEM block")
+	}
+
+	cert, err := x509.ParseCertificate(block.Bytes)
+	if err != nil {
+		return fmt.Errorf("CheckAndRenewCertificate: failed to parse certificate - %w", err)
+	}
+
+	now := time.Now()
+	timeUntilExpiry := cert.NotAfter.Sub(now)
+
+	switch {
+	case cert.NotAfter.Before(now):
+		_ = registry.DeleteEntry(registry.AUTH, "Cert")
+		_ = registry.DeleteEntry(registry.AUTH, "Priv")
+
+		return fmt.Errorf("Certificate has expired. This agent needs to be bootstrapped again.")
+	case timeUntilExpiry < renewalWindow:
+		fmt.Printf("Certificate expires in %v hours. Renewing...\n", timeUntilExpiry.Hours())
+		return renewCertificate()
+	default:
+		fmt.Printf("Certificate valid for %v days. No renewal needed.\n", timeUntilExpiry.Hours()/24)
+		return nil
+	}
+}
+
+func renewCertificate() error {
+	hostname, _ := os.Hostname()
+
+	csr, privKey, err := certificates.GenerateCSR(hostname, 2048)
+	if err != nil {
+		return fmt.Errorf("Bootstrap: generating csr failed -> %w", err)
+	}
+
+	encodedCSR := base64.StdEncoding.EncodeToString(csr)
+
+	reqBody, err := json.Marshal(&BootstrapRequest{
+		Hostname: hostname,
+		Drives:   utils.GetLocalDrives(),
+		CSR:      encodedCSR,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to marshal bootstrap request: %w", err)
+	}
+
+	renewResp := &BootstrapResponse{}
+
+	_, err = ProxmoxHTTPRequest(http.MethodPost, "/plus/agent/renew", bytes.NewBuffer(reqBody), &renewResp)
+	if err != nil {
+		return fmt.Errorf("failed to fetch renewed certificate: %w", err)
+	}
+
+	decodedCA, err := base64.StdEncoding.DecodeString(renewResp.CA)
+	if err != nil {
+		return fmt.Errorf("Renew: error decoding ca content (%s) -> %w", string(renewResp.CA), err)
+	}
+
+	decodedCert, err := base64.StdEncoding.DecodeString(renewResp.Cert)
+	if err != nil {
+		return fmt.Errorf("Renew: error decoding cert content (%s) -> %w", string(renewResp.Cert), err)
+	}
+
+	privKeyPEM := certificates.EncodeKeyPEM(privKey)
+
+	caEntry := registry.RegistryEntry{
+		Key:      "ServerCA",
+		Value:    string(decodedCA),
+		Path:     registry.AUTH,
+		IsSecret: true,
+	}
+
+	certEntry := registry.RegistryEntry{
+		Key:      "Cert",
+		Value:    string(decodedCert),
+		Path:     registry.AUTH,
+		IsSecret: true,
+	}
+
+	privEntry := registry.RegistryEntry{
+		Key:      "Priv",
+		Value:    string(privKeyPEM),
+		Path:     registry.AUTH,
+		IsSecret: true,
+	}
+
+	err = registry.CreateEntry(&caEntry)
+	if err != nil {
+		return fmt.Errorf("Renew: error storing ca to registry -> %w", err)
+	}
+
+	err = registry.CreateEntry(&certEntry)
+	if err != nil {
+		return fmt.Errorf("Renew: error storing cert to registry -> %w", err)
+	}
+
+	err = registry.CreateEntry(&privEntry)
+	if err != nil {
+		return fmt.Errorf("Renew: error storing priv to registry -> %w", err)
+	}
+
+	return nil
+}

internal/backend/mount/mount.go

@@ -106,6 +106,7 @@ func Mount(storeInstance *store.Store, target *types.Target) (*AgentMount, error
 
 	// If all retries failed, clean up and return error
 	agentMount.Unmount()
+	agentMount.CloseMount()
 	return nil, fmt.Errorf("Mount: error mounting NFS share after %d attempts -> %w", maxRetries, lastErr)
 }
 

internal/config/mutex_manager.go

@@ -0,0 +1,77 @@
+package config
+
+import (
+	"path/filepath"
+	"sync"
+)
+
+type FileMutexManager struct {
+	mu     sync.RWMutex
+	locks  map[string]*sync.RWMutex
+	refCnt map[string]int
+}
+
+func NewFileMutexManager() *FileMutexManager {
+	return &FileMutexManager{
+		locks:  make(map[string]*sync.RWMutex),
+		refCnt: make(map[string]int),
+	}
+}
+
+func (fm *FileMutexManager) getLock(path string) *sync.RWMutex {
+	fm.mu.Lock()
+	defer fm.mu.Unlock()
+
+	absPath, err := filepath.Abs(path)
+	if err != nil {
+		absPath = path
+	}
+
+	if lock, exists := fm.locks[absPath]; exists {
+		fm.refCnt[absPath]++
+		return lock
+	}
+
+	lock := &sync.RWMutex{}
+	fm.locks[absPath] = lock
+	fm.refCnt[absPath] = 1
+	return lock
+}
+
+func (fm *FileMutexManager) releaseLock(path string) {
+	fm.mu.Lock()
+	defer fm.mu.Unlock()
+
+	absPath, err := filepath.Abs(path)
+	if err != nil {
+		absPath = path
+	}
+
+	if fm.refCnt[absPath] > 1 {
+		fm.refCnt[absPath]--
+		return
+	}
+
+	delete(fm.locks, absPath)
+	delete(fm.refCnt, absPath)
+}
+
+func (fm *FileMutexManager) WithReadLock(path string, fn func() error) error {
+	lock := fm.getLock(path)
+	lock.RLock()
+	defer func() {
+		lock.RUnlock()
+		fm.releaseLock(path)
+	}()
+	return fn()
+}
+
+func (fm *FileMutexManager) WithWriteLock(path string, fn func() error) error {
+	lock := fm.getLock(path)
+	lock.Lock()
+	defer func() {
+		lock.Unlock()
+		fm.releaseLock(path)
+	}()
+	return fn()
+}