usb-server/internal/relay/diag.go

package relay

import (
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"
)

// Diagnostics drop-off.
//
// Getting a report off an awkward machine — a headless NAS, a Windows box in
// the middle of driver debugging — is otherwise a matter of copying thousands
// of lines by hand. The relay is already reachable from every client, so it
// makes a convenient place to leave one.
//
// Reports are held in memory only, capped in size and count, and expire. The
// relay is not a storage service, and treating it like one is how it would
// become one.
const (
	// maxDiagReports bounds how many are kept; the oldest is dropped first.
	maxDiagReports = 32

	// maxDiagSize bounds one report.
	maxDiagSize = 4 << 20 // 4 MB

	// diagTTL is how long a report survives. Long enough to fetch and read,
	// short enough that machine details do not linger.
	diagTTL = 24 * time.Hour
)

// RetentionNote describes the retention policy for the client to print.
const RetentionNote = "24 hours"

type diagReport struct {
	data     []byte
	stored   time.Time
	fetched  int
	remoteIP string
}

type diagStore struct {
	mu      sync.Mutex
	reports map[string]*diagReport
}

func newDiagStore() *diagStore {
	return &diagStore{reports: make(map[string]*diagReport)}
}

// put stores a report, evicting the oldest if the store is full.
func (s *diagStore) put(id string, data []byte, remoteIP string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.expireLocked()

	if len(s.reports) >= maxDiagReports {
		var oldestID string
		var oldest time.Time
		for id, report := range s.reports {
			if oldestID == "" || report.stored.Before(oldest) {
				oldestID, oldest = id, report.stored
			}
		}
		delete(s.reports, oldestID)
	}

	s.reports[id] = &diagReport{
		data:     data,
		stored:   time.Now(),
		remoteIP: remoteIP,
	}
}

func (s *diagStore) get(id string) ([]byte, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.expireLocked()

	report, ok := s.reports[id]
	if !ok {
		return nil, false
	}
	report.fetched++
	return report.data, true
}

// expireLocked drops reports past their TTL. Callers must hold the lock.
func (s *diagStore) expireLocked() {
	cutoff := time.Now().Add(-diagTTL)
	for id, report := range s.reports {
		if report.stored.Before(cutoff) {
			delete(s.reports, id)
		}
	}
}

// handleDiag serves the diagnostics endpoint: PUT to store, GET to retrieve.
func (s *Server) handleDiag(w http.ResponseWriter, r *http.Request) {
	id := strings.TrimPrefix(r.URL.Path, "/diag/")
	if id == "" || strings.Contains(id, "/") {
		http.Error(w, "report ID required: /diag/<id>", http.StatusBadRequest)
		return
	}

	switch r.Method {
	case http.MethodPut, http.MethodPost:
		s.storeDiag(w, r, id)
	case http.MethodGet:
		s.fetchDiag(w, id)
	default:
		http.Error(w, "use PUT to store and GET to retrieve", http.StatusMethodNotAllowed)
	}
}

func (s *Server) storeDiag(w http.ResponseWriter, r *http.Request, id string) {
	// LimitReader rather than trusting Content-Length: a client can lie about
	// that, and this endpoint takes uploads from anyone who can reach it.
	data, err := io.ReadAll(io.LimitReader(r.Body, maxDiagSize+1))
	if err != nil {
		http.Error(w, "could not read the report", http.StatusBadRequest)
		return
	}
	if len(data) > maxDiagSize {
		http.Error(w, fmt.Sprintf("report exceeds the %d byte limit", maxDiagSize),
			http.StatusRequestEntityTooLarge)
		return
	}
	if len(data) == 0 {
		http.Error(w, "empty report", http.StatusBadRequest)
		return
	}

	s.diag.put(id, data, clientIP(r))

	w.WriteHeader(http.StatusCreated)
	fmt.Fprintf(w, "stored as %s, kept for %s\n", id, RetentionNote)
}

func (s *Server) fetchDiag(w http.ResponseWriter, id string) {
	data, ok := s.diag.get(id)
	if !ok {
		http.Error(w, "no such report (wrong ID, or it expired)", http.StatusNotFound)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(data)
}