edubox/agent/proxy.go

package main

import (
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

const (
	ProxyModeDisabled = "disabled"
	ProxyModeAuto     = "auto"
	ProxyModeEnabled  = "enabled"
)

// autoProxyLockDuration is the minimum time we stay in proxy mode once the
// agent automatically switched to it. This prevents flip-flopping on short
// network blips.
const autoProxyLockDuration = 5 * time.Minute

// proxyState tracks the runtime proxy decision in "auto" mode. It is guarded
// by proxyMu.
var (
	proxyMu          sync.RWMutex
	proxyActive      bool
	proxyLockedUntil time.Time
)

// proxyMode normalizes the configured proxy mode.
func proxyMode(cfg *AgentConfig) string {
	if cfg == nil {
		return ProxyModeDisabled
	}
	switch strings.ToLower(strings.TrimSpace(cfg.ProxyMode)) {
	case ProxyModeEnabled:
		return ProxyModeEnabled
	case ProxyModeAuto:
		return ProxyModeAuto
	default:
		return ProxyModeDisabled
	}
}

// IsProxyActive reports whether outbound requests should currently go through
// the configured proxy. In "enabled" mode it always returns true; in "auto"
// mode it reflects the last automatic decision.
func IsProxyActive() bool {
	proxyMu.RLock()
	defer proxyMu.RUnlock()
	return proxyActive
}

// setProxyActive updates the runtime proxy decision and, in auto mode, locks
// the decision for autoProxyLockDuration to avoid flip-flopping.
func setProxyActive(active bool) bool {
	proxyMu.Lock()
	defer proxyMu.Unlock()
	changed := proxyActive != active
	proxyActive = active
	if active {
		proxyLockedUntil = time.Now().Add(autoProxyLockDuration)
	}
	return changed
}

// resetProxyState disables the automatic proxy decision. Call this when the
// configuration changes.
func resetProxyState() {
	proxyMu.Lock()
	proxyActive = false
	proxyLockedUntil = time.Time{}
	proxyMu.Unlock()
}

// canRetryDirect reports whether enough time has passed to try a direct
// connection again while in auto-proxy mode.
func canRetryDirect() bool {
	proxyMu.RLock()
	defer proxyMu.RUnlock()
	return time.Now().After(proxyLockedUntil)
}

// proxyURL parses and validates the configured proxy URL.
func proxyURL(cfg *AgentConfig) *url.URL {
	if cfg == nil || cfg.ProxyURL == "" {
		return nil
	}
	u, err := url.Parse(cfg.ProxyURL)
	if err != nil {
		return nil
	}
	return u
}

// proxyFunc returns a proxy selection function for http.Transport. It returns
// nil when the proxy should not be used.
func proxyFunc(cfg *AgentConfig) func(*http.Request) (*url.URL, error) {
	mode := proxyMode(cfg)
	u := proxyURL(cfg)

	switch mode {
	case ProxyModeEnabled:
		if u == nil {
			return nil
		}
		return func(*http.Request) (*url.URL, error) { return u, nil }
	case ProxyModeAuto:
		if u == nil {
			return nil
		}
		if !IsProxyActive() {
			return nil
		}
		return func(*http.Request) (*url.URL, error) { return u, nil }
	default:
		return nil
	}
}

// websocketDialer returns a websocket.Dialer configured for the current proxy
// mode and state.
func websocketDialer(cfg *AgentConfig) *websocket.Dialer {
	d := websocket.DefaultDialer
	fn := proxyFunc(cfg)
	if fn == nil {
		return d
	}
	return &websocket.Dialer{
		Proxy:             fn,
		HandshakeTimeout:  d.HandshakeTimeout,
		ReadBufferSize:    d.ReadBufferSize,
		WriteBufferSize:   d.WriteBufferSize,
		EnableCompression: d.EnableCompression,
	}
}

// httpClientWithProxy returns an http.Client configured for the current proxy
// mode and state.
func httpClientWithProxy(cfg *AgentConfig) *http.Client {
	fn := proxyFunc(cfg)
	if fn == nil {
		return http.DefaultClient
	}
	return &http.Client{
		Transport: &http.Transport{Proxy: fn},
	}
}