steamcache2/steamcache/steamcache.go

// steamcache/steamcache.go
package steamcache

import (
	"context"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"s1d3sw1ped/SteamCache2/steamcache/logger"
	"s1d3sw1ped/SteamCache2/vfs"
	"s1d3sw1ped/SteamCache2/vfs/cache"
	"s1d3sw1ped/SteamCache2/vfs/disk"
	"s1d3sw1ped/SteamCache2/vfs/gc"
	"s1d3sw1ped/SteamCache2/vfs/memory"
	"strings"
	"sync"
	"time"

	"github.com/docker/go-units"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	requestsTotal = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "http_requests_total",
			Help: "Total number of HTTP requests",
		},
		[]string{"method", "status"},
	)

	cacheStatusTotal = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "cache_status_total",
			Help: "Total cache status counts",
		},
		[]string{"status"},
	)

	responseTime = promauto.NewHistogram(
		prometheus.HistogramOpts{
			Name:    "response_time_seconds",
			Help:    "Response time in seconds",
			Buckets: prometheus.DefBuckets,
		},
	)
)

type SteamCache struct {
	address  string
	upstream string

	vfs vfs.VFS

	memory *memory.MemoryFS
	disk   *disk.DiskFS

	memorygc *gc.GCFS
	diskgc   *gc.GCFS

	server *http.Server
	client *http.Client
	cancel context.CancelFunc
	wg     sync.WaitGroup
}

func New(address string, memorySize string, diskSize string, diskPath, upstream string) *SteamCache {
	memorysize, err := units.FromHumanSize(memorySize)
	if err != nil {
		panic(err)
	}

	disksize, err := units.FromHumanSize(diskSize)
	if err != nil {
		panic(err)
	}

	c := cache.New(
		gc.PromotionDecider,
	)

	var m *memory.MemoryFS
	var mgc *gc.GCFS
	if memorysize > 0 {
		m = memory.New(memorysize)
		mgc = gc.New(m, gc.LRUGC)
	}

	var d *disk.DiskFS
	var dgc *gc.GCFS
	if disksize > 0 {
		d = disk.New(diskPath, disksize)
		dgc = gc.New(d, gc.LRUGC)
	}

	// configure the cache to match the specified mode (memory only, disk only, or memory and disk) based on the provided sizes
	if disksize == 0 && memorysize != 0 {
		//memory only mode - no disk

		c.SetSlow(mgc)
	} else if disksize != 0 && memorysize == 0 {
		// disk only mode

		c.SetSlow(dgc)
	} else if disksize != 0 && memorysize != 0 {
		// memory and disk mode

		c.SetFast(mgc)
		c.SetSlow(dgc)
	} else {
		// no memory or disk isn't a valid configuration
		logger.Logger.Error().Bool("memory", false).Bool("disk", false).Msg("configuration invalid :( exiting")
		os.Exit(1)
	}

	transport := &http.Transport{
		MaxIdleConns:        100,
		MaxIdleConnsPerHost: 10,
		IdleConnTimeout:     90 * time.Second,
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		TLSHandshakeTimeout:   10 * time.Second,
		ResponseHeaderTimeout: 10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}

	client := &http.Client{
		Transport: transport,
		Timeout:   60 * time.Second,
	}

	sc := &SteamCache{
		upstream: upstream,
		address:  address,
		vfs:      c,
		memory:   m,
		disk:     d,
		memorygc: mgc,
		diskgc:   dgc,
		client:   client,
		server: &http.Server{
			Addr:         address,
			ReadTimeout:  5 * time.Second,
			WriteTimeout: 10 * time.Second,
			IdleTimeout:  120 * time.Second,
		},
	}

	if d != nil {
		if d.Size() > d.Capacity() {
			gc.LRUGC(d, uint(d.Size()-d.Capacity()))
		}
	}

	return sc
}

func (sc *SteamCache) Run() {
	if sc.upstream != "" {
		resp, err := sc.client.Get(sc.upstream)
		if err != nil || resp.StatusCode != http.StatusOK {
			logger.Logger.Error().Err(err).Str("upstream", sc.upstream).Msg("Failed to connect to upstream server")
			os.Exit(1)
		}
		resp.Body.Close()
	}

	sc.server.Handler = sc
	ctx, cancel := context.WithCancel(context.Background())
	sc.cancel = cancel

	sc.wg.Add(1)
	go func() {
		defer sc.wg.Done()
		err := sc.server.ListenAndServe()
		if err != nil && err != http.ErrServerClosed {
			logger.Logger.Error().Err(err).Msg("Failed to start SteamCache2")
			os.Exit(1)
		}
	}()

	<-ctx.Done()
	sc.server.Shutdown(ctx)
	sc.wg.Wait()
}

func (sc *SteamCache) Shutdown() {
	if sc.cancel != nil {
		sc.cancel()
	}
	sc.wg.Wait()
}

func (sc *SteamCache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/metrics" {
		promhttp.Handler().ServeHTTP(w, r)
		return
	}

	if r.Method != http.MethodGet {
		requestsTotal.WithLabelValues(r.Method, "405").Inc()
		logger.Logger.Warn().Str("method", r.Method).Msg("Only GET method is supported")
		http.Error(w, "Only GET method is supported", http.StatusMethodNotAllowed)
		return
	}

	if r.URL.String() == "/lancache-heartbeat" {
		w.Header().Add("X-LanCache-Processed-By", "SteamCache2")
		w.WriteHeader(http.StatusNoContent)
		w.Write(nil)
		return
	}

	if strings.HasPrefix(r.URL.String(), "/depot/") {
		// trim the query parameters from the URL path
		// this is necessary because the cache key should not include query parameters
		path := strings.Split(r.URL.String(), "?")[0]

		tstart := time.Now()
		defer func() { responseTime.Observe(time.Since(tstart).Seconds()) }()

		cacheKey := strings.ReplaceAll(path[1:], "\\", "/") // replace all backslashes with forward slashes shouldn't be necessary but just in case

		if cacheKey == "" {
			requestsTotal.WithLabelValues(r.Method, "400").Inc()
			logger.Logger.Warn().Str("url", path).Msg("Invalid URL")
			http.Error(w, "Invalid URL", http.StatusBadRequest)
			return
		}

		w.Header().Add("X-LanCache-Processed-By", "SteamCache2") // SteamPrefill uses this header to determine if the request was processed by the cache maybe steam uses it too

		reader, err := sc.vfs.Open(cacheKey)
		if err == nil {
			defer reader.Close()
			w.Header().Add("X-LanCache-Status", "HIT")

			io.Copy(w, reader)

			logger.Logger.Info().
				Str("key", cacheKey).
				Str("host", r.Host).
				Str("status", "HIT").
				Dur("duration", time.Since(tstart)).
				Msg("request")

			requestsTotal.WithLabelValues(r.Method, "200").Inc()
			cacheStatusTotal.WithLabelValues("HIT").Inc()

			return
		}

		var req *http.Request
		if sc.upstream != "" { // if an upstream server is configured, proxy the request to the upstream server
			ur, err := url.JoinPath(sc.upstream, path)
			if err != nil {
				requestsTotal.WithLabelValues(r.Method, "500").Inc()
				logger.Logger.Error().Err(err).Str("upstream", sc.upstream).Msg("Failed to join URL path")
				http.Error(w, "Failed to join URL path", http.StatusInternalServerError)
				return
			}

			req, err = http.NewRequest(http.MethodGet, ur, nil)
			if err != nil {
				requestsTotal.WithLabelValues(r.Method, "500").Inc()
				logger.Logger.Error().Err(err).Str("upstream", sc.upstream).Msg("Failed to create request")
				http.Error(w, "Failed to create request", http.StatusInternalServerError)
				return
			}
			req.Host = r.Host
		} else { // if no upstream server is configured, proxy the request to the host specified in the request
			host := r.Host
			if r.Header.Get("X-Sls-Https") == "enable" {
				host = "https://" + host
			} else {
				host = "http://" + host
			}

			ur, err := url.JoinPath(host, path)
			if err != nil {
				requestsTotal.WithLabelValues(r.Method, "500").Inc()
				logger.Logger.Error().Err(err).Str("host", host).Msg("Failed to join URL path")
				http.Error(w, "Failed to join URL path", http.StatusInternalServerError)
				return
			}

			req, err = http.NewRequest(http.MethodGet, ur, nil)
			if err != nil {
				requestsTotal.WithLabelValues(r.Method, "500").Inc()
				logger.Logger.Error().Err(err).Str("host", host).Msg("Failed to create request")
				http.Error(w, "Failed to create request", http.StatusInternalServerError)
				return
			}
		}

		// Copy headers from the original request to the new request
		for key, values := range r.Header {
			for _, value := range values {
				req.Header.Add(key, value)
			}
		}

		// Retry logic
		backoffSchedule := []time.Duration{1 * time.Second, 3 * time.Second, 10 * time.Second}
		var resp *http.Response
		for i, backoff := range backoffSchedule {
			resp, err = sc.client.Do(req)
			if err == nil && resp.StatusCode == http.StatusOK {
				break
			}
			if i < len(backoffSchedule)-1 {
				time.Sleep(backoff)
			}
		}
		if err != nil || resp.StatusCode != http.StatusOK {
			requestsTotal.WithLabelValues(r.Method, "500 upstream host "+r.Host).Inc()
			logger.Logger.Error().Err(err).Str("url", req.URL.String()).Msg("Failed to fetch the requested URL")
			http.Error(w, "Failed to fetch the requested URL", http.StatusInternalServerError)
			return
		}
		defer resp.Body.Close()

		size := resp.ContentLength

		// this is sortof not needed as we should always be able to get a writer from the cache as long as the gc is able to reclaim enough space aka the file is not bigger than the disk can handle
		ww := w.(io.Writer)                        // default writer to write to the response writer
		writer, _ := sc.vfs.Create(cacheKey, size) // create a writer to write to the cache
		if writer != nil {                         // if the writer is not nil, it means the cache is writable
			defer writer.Close()           // close the writer when done
			ww = io.MultiWriter(w, writer) // write to both the response writer and the cache writer
		}

		w.Header().Add("X-LanCache-Status", "MISS")

		io.Copy(ww, resp.Body)

		logger.Logger.Info().
			Str("key", cacheKey).
			Str("host", r.Host).
			Str("status", "MISS").
			Dur("duration", time.Since(tstart)).
			Msg("request")

		requestsTotal.WithLabelValues(r.Method, "200").Inc()
		cacheStatusTotal.WithLabelValues("MISS").Inc()

		return
	}

	if r.URL.Path == "/favicon.ico" {
		w.WriteHeader(http.StatusNoContent)
		return
	}

	if r.URL.Path == "/robots.txt" {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("User-agent: *\nDisallow: /\n"))
		return
	}

	requestsTotal.WithLabelValues(r.Method, "404").Inc()
	logger.Logger.Warn().Str("url", r.URL.String()).Msg("Not found")
	http.Error(w, "Not found", http.StatusNotFound)
}