agent.go

package main

import (
	"errors"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"
)

const (
	defaultTokenFileLoc = "/var/lib/wiresteward/token"
)

// Agent is the wirestward client instance that manages a set of network devices
// based on configuration generated by remote wiresteward servers.
type Agent struct {
	deviceManagers []*DeviceManager
	oa             *oauthTokenHandler
	renewMu        sync.Mutex // serializes renewAllLeases calls
}

// NewAgent creates an Agent from an AgentConfig. It generates a DeviceManager
// per device specified in the configuration, sets up and starts the associated
// resources.
func NewAgent(cfg *agentConfig) *Agent {
	agent := &Agent{}
	for _, dev := range cfg.Devices {
		urls := []string{}
		for _, peer := range dev.Peers {
			urls = append(urls, peer.URL)
		}
		dm := newDeviceManager(dev.Name, dev.MTU, urls, cfg.HTTPClient.Timeout, cfg.HealthCheck)
		if err := dm.Run(); err != nil {
			logger.Errorf("Error starting device `%s`: %v", dm.Name(), err)
			continue
		}
		agent.deviceManagers = append(agent.deviceManagers, dm)
	}
	tokenDir := filepath.Dir(defaultTokenFileLoc)
	err := os.MkdirAll(tokenDir, 0750)
	if err != nil {
		logger.Errorf("Unable to create directory=%s", tokenDir)
	}
	agent.oa = newOAuthTokenHandler(
		cfg.OAuth.AuthURL,
		cfg.OAuth.TokenURL,
		cfg.OAuth.ClientID,
		defaultTokenFileLoc,
		cfg.OAuth.RefreshBeforeExpiry.Duration,
	)
	return agent
}

// ListenAndServe sets up and starts the background token refresh loop.
func (a *Agent) ListenAndServe() {
	http.HandleFunc("/oauth2/callback", a.callbackHandler)
	http.HandleFunc("/renew", a.renewHandler)
	http.HandleFunc("/", a.mainHandler)

	logger.Verbosef("Starting agent at http://%s", *flagAgentAddress)

	go a.tokenRefreshLoop()

	if err := http.ListenAndServe(*flagAgentAddress, nil); err != nil {
		logger.Errorf("%v", err)
	}
}

// tokenRefreshLoop proactively checks the token status every minute.
func (a *Agent) tokenRefreshLoop() {
	// On startup, immediately renew leases if a valid cached token exists,
	// without waiting for the first refresh tick.
	token, _, err := a.oa.GetToken()
	if err != nil {
		logger.Verbosef("Could not fetch oauth token on startup %v", err)
	} else {
		a.renewAllLeases(token.AccessToken)
	}

	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()

	for range ticker.C {
		a.checkAndRefresh()
	}
}

func (a *Agent) checkAndRefresh() {
	token, refreshed, err := a.oa.GetToken()
	if err != nil {
		logger.Errorf("Could not fetch oauth token %v", err)
		return
	}

	if refreshed {
		logger.Verbosef("Token refreshed. Syncing new %v lease with remote server.", time.Until(token.Expiry).Round(time.Second))
		a.renewAllLeases(token.AccessToken)
	}
}

// Stop calls the Stop method on all DeviceManager instances that this Agent
// controls.
func (a *Agent) Stop() {
	for _, dm := range a.deviceManagers {
		dm.Stop()
	}
}

func (a *Agent) renewAllLeases(token string) {
	a.renewMu.Lock()
	defer a.renewMu.Unlock()
	logger.Verbosef("Running renew leases loop..")
	for _, dm := range a.deviceManagers {
		dm.setCachedToken(token)
	}
	for _, dm := range a.deviceManagers {
		dm.triggerLeaseRenewal()
	}
}

func (a *Agent) callbackHandler(w http.ResponseWriter, r *http.Request) {
	stateCookie, err := r.Cookie(oauthStateCookieName)
	if errors.Is(err, http.ErrNoCookie) {
		logger.Errorf("State cookie missing: %s", err)
		http.Error(w, "State cookie missing", 500)
		return
	}
	if err != nil {
		logger.Errorf("Failed to retrieve state cookie: %s", err)
		http.Error(w, "Failed to retrieve state cookie", 500)
		return
	}

	if r.FormValue("state") != stateCookie.Value {
		logger.Errorf("State token mismatch: expected=%s received=%s", stateCookie.Value, r.FormValue("state"))
		http.Error(w, "State token mismatch", 500)
		return
	}
	token, err := a.oa.ExchangeToken(r.FormValue("code"))
	if err != nil {
		fmt.Fprintf(w, "error fetching token from web: %v", err)
		return
	}
	a.renewAllLeases(token.AccessToken)
	// Redirect to / after renewing leases
	rootURL := fmt.Sprintf("http://%s/", r.Host)
	http.Redirect(w, r, rootURL, 302)
}

// renewHandler will initiate the auth challenge. Leases renewal will be handled
// in the callback handler.
func (a *Agent) renewHandler(w http.ResponseWriter, r *http.Request) {
	url, err := a.oa.prepareTokenWebChalenge(w)
	if err != nil {
		fmt.Fprintf(w, "error creating web url to get token: %v", err)
		return
	}
	http.Redirect(w, r, url, 302)
}

func (a *Agent) mainHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	token, refreshed, err := a.oa.GetToken()
	if err != nil {
		logger.Verbosef("cannot get a valid token: %v", err)
		statusHTTPWriter(w, r, a.deviceManagers, nil)
		return
	}
	if refreshed {
		logger.Verbosef("Token refreshed. Syncing new %v lease with remote server.", time.Until(token.Expiry).Round(time.Second))
		a.renewAllLeases(token.AccessToken)
	}
	statusHTTPWriter(w, r, a.deviceManagers, token)
}