celestia_server.go

package celestia

import (
	"bytes"
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/celestiaorg/op-alt-da/fallback"
	"github.com/celestiaorg/op-alt-da/metrics"
	altda "github.com/ethereum-optimism/optimism/op-alt-da"
	"github.com/ethereum/go-ethereum/common/hexutil"
	"github.com/ethereum/go-ethereum/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// securityHeadersMiddleware adds security headers to prevent MIME sniffing and clickjacking.
func securityHeadersMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("X-Frame-Options", "DENY")
		next.ServeHTTP(w, r)
	})
}

// sanitizeLogField removes newlines and carriage returns to prevent log injection.
func sanitizeLogField(s string) string {
	s = strings.ReplaceAll(s, "\n", "")
	s = strings.ReplaceAll(s, "\r", "")
	return s
}

// recoveryMiddleware recovers from panics to prevent server crashes.
// Logs the panic with request context and returns 500 to the client.
func recoveryMiddleware(next http.Handler, logger log.Logger) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if err := recover(); err != nil {
				// Check if client disconnected (context cancelled)
				if r.Context().Err() != nil {
					logger.Debug("Client disconnected during panic recovery",
						"path", sanitizeLogField(r.URL.Path),
						"method", r.Method)
					return
				}
				logger.Error("Handler panic recovered",
					"err", err,
					"path", sanitizeLogField(r.URL.Path),
					"method", r.Method,
					"remote", r.RemoteAddr)
				http.Error(w, "internal server error", http.StatusInternalServerError)
			}
		}()
		next.ServeHTTP(w, r)
	})
}

// Store defines the interface for blob storage operations.
// This allows mocking the storage layer in tests.
type Store interface {
	Get(ctx context.Context, key []byte) ([]byte, error)
	Put(ctx context.Context, data []byte) ([]byte, []byte, error)
	CreateCommitment(data []byte) ([]byte, error)
}

// CelestiaServer implements the HTTP server for the DA server.
// Supports optional fallback provider for redundant storage.
type CelestiaServer struct {
	log      log.Logger
	endpoint string
	host     string

	// Celestia storage layer
	store Store

	// Timeouts and settings
	submitTimeout time.Duration
	getTimeout    time.Duration
	maxBlobSize   int64 // DoS protection: limit request body size

	// HTTP server
	httpServer *http.Server
	listener   net.Listener

	// Metrics
	metricsEnabled bool
	metricsPort    int
	metrics        *metrics.CelestiaMetrics
	metricsServer  *http.Server // Separate server for graceful shutdown

	// Fallback provider for write-through and read-fallback
	fallback       fallback.Provider
	fallbackMode   string

	// Track pending async operations for graceful shutdown
	wg sync.WaitGroup
}

// NewCelestiaServer creates a new Celestia DA server.
func NewCelestiaServer(
	host string,
	port int,
	store Store,
	submitTimeout time.Duration,
	getTimeout time.Duration,
	httpReadTimeout time.Duration,
	httpWriteTimeout time.Duration,
	httpIdleTimeout time.Duration,
	maxBlobSize int64,
	metricsEnabled bool,
	metricsPort int,
	fallbackProvider fallback.Provider,
	fallbackMode string,
	log log.Logger,
) *CelestiaServer {
	endpoint := net.JoinHostPort(host, strconv.Itoa(port))

	if fallbackProvider == nil {
		fallbackProvider = &fallback.NoopProvider{}
	}

	server := &CelestiaServer{
		log:            log,
		endpoint:       endpoint,
		host:           host,
		store:          store,
		submitTimeout:  submitTimeout,
		getTimeout:     getTimeout,
		maxBlobSize:    maxBlobSize,
		metricsEnabled: metricsEnabled,
		metricsPort:    metricsPort,
		fallback:       fallbackProvider,
		fallbackMode:   fallback.NormalizeMode(fallbackMode),
		httpServer: &http.Server{
			Addr:           endpoint,
			ReadTimeout:    httpReadTimeout,
			WriteTimeout:   httpWriteTimeout,
			IdleTimeout:    httpIdleTimeout,
			MaxHeaderBytes: 1 << 16, // 64KB max headers
		},
	}

	if metricsEnabled {
		server.metrics = metrics.NewCelestiaMetrics(prometheus.DefaultRegisterer)
	}

	return server
}

func (d *CelestiaServer) Start(ctx context.Context) error {
	mux := http.NewServeMux()

	mux.HandleFunc("/get/", d.HandleGet)
	mux.HandleFunc("/put/", d.HandlePut)
	mux.HandleFunc("/put", d.HandlePut)
	mux.HandleFunc("/health", d.HandleHealth)

	// Middleware chain: recovery (outermost) -> security headers -> handler
	handler := recoveryMiddleware(securityHeadersMiddleware(mux), d.log)
	d.httpServer.Handler = handler

	listener, err := net.Listen("tcp", d.endpoint)
	if err != nil {
		return fmt.Errorf("failed to listen: %w", err)
	}
	d.listener = listener
	d.endpoint = listener.Addr().String()

	// Start main HTTP server
	errCh := make(chan error, 1)
	go func() {
		d.log.Info("Starting HTTP server", "addr", d.endpoint)
		if err := d.httpServer.Serve(d.listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
			errCh <- err
		}
	}()

	// Start metrics server if enabled
	if d.metricsEnabled {
		d.metricsServer = &http.Server{
			Addr:    fmt.Sprintf(":%d", d.metricsPort),
			Handler: promhttp.Handler(),
		}
		go func() {
			d.log.Info("Starting metrics server", "addr", d.metricsServer.Addr)
			if err := d.metricsServer.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
				d.log.Error("Metrics server failed", "err", err)
			}
		}()
	}

	// Wait briefly to catch immediate startup errors
	select {
	case err := <-errCh:
		return fmt.Errorf("http server failed: %w", err)
	case <-ctx.Done():
		return ctx.Err()
	case <-time.After(10 * time.Millisecond):
		return nil
	}
}

// HandleHealth returns 200 OK for health checks.
func (d *CelestiaServer) HandleHealth(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write([]byte("OK"))
}

// HandleGet retrieves a blob directly from Celestia.
// Request: GET /get/<hex-encoded-commitment>
// Response: Raw blob data on success, 404 on not found, 500 on error.
func (d *CelestiaServer) HandleGet(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	start := time.Now()
	d.log.Debug("GET", "url", r.URL)

	// Validate route
	if path.Dir(r.URL.Path) != "/get" {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	// Parse commitment
	comm, err := d.parseCommitment(r.URL.Path)
	if err != nil {
		d.log.Warn("Invalid commitment format", "err", err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	// Retrieve blob (Celestia first, fallback if not found)
	ctx := r.Context()
	data, err := d.getBlob(ctx, comm)
	if err != nil {
		d.handleGetError(w, comm, err, start)
		return
	}

	// Success
	d.recordGetSuccess(start, len(data))
	d.log.Info("Blob retrieved", "commitment", hex.EncodeToString(comm), "size", len(data), "duration", time.Since(start))

	w.Header().Set("Content-Type", "application/octet-stream")
	if _, err := w.Write(data); err != nil {
		d.log.Error("Failed to write response", "err", err)
	}
}

// parseCommitment extracts and validates commitment from URL path.
func (d *CelestiaServer) parseCommitment(urlPath string) ([]byte, error) {
	key := path.Base(urlPath)
	comm, err := hexutil.Decode(key)
	if err != nil {
		// Sanitize for logging
		safeKey := strings.ReplaceAll(strings.ReplaceAll(key, "\n", ""), "\r", "")
		return nil, fmt.Errorf("invalid hex %q: %w", safeKey, err)
	}
	return comm, nil
}

// getBlob retrieves blob from Celestia. If found in Celestia and fallback is available,
// it populates fallback asynchronously (read-through). For non-terminal Celestia errors,
// it may consult the fallback, but fallback data must verify against the requested commitment.
func (d *CelestiaServer) getBlob(ctx context.Context, comm []byte) ([]byte, error) {
	getCtx, cancel := context.WithTimeout(ctx, d.getTimeout)
	defer cancel()

	data, celestiaErr := d.store.Get(getCtx, comm)
	if celestiaErr == nil {
		// Success from Celestia - optionally read-through to fallback for future requests
		if d.fallback.Available() && fallback.WriteEnabled(d.fallbackMode) {
			d.wg.Add(1)
			go d.putFallback(context.Background(), comm, data)
		}
		return data, nil
	}

	if isNotFoundError(celestiaErr) {
		return nil, altda.ErrNotFound
	}

	// Celestia failed - try fallback if available
	if !d.fallback.Available() {
		d.log.Warn("Celestia failed and no fallback available", "err", celestiaErr)
		return nil, celestiaErr
	}

	// Log that we're falling back (INFO level so it's visible)
	d.log.Info("Celestia failed, trying S3 fallback", "commitment", hex.EncodeToString(comm), "celestia_err", celestiaErr)

	data, err := d.getFallback(ctx, comm)
	if err != nil {
		d.log.Warn("Fallback read also failed", "provider", d.fallback.Name(), "fallback_err", err, "original_celestia_err", celestiaErr)
		return nil, celestiaErr
	}

	if err := d.verifyFallbackCommitment(comm, data); err != nil {
		d.log.Error("Fallback data failed integrity verification",
			"provider", d.fallback.Name(),
			"commitment", hex.EncodeToString(comm),
			"err", err,
			"original_celestia_err", celestiaErr)
		return nil, err
	}

	d.log.Info("Blob retrieved from fallback", "provider", d.fallback.Name(), "commitment", hex.EncodeToString(comm))
	return data, nil
}

func (d *CelestiaServer) verifyFallbackCommitment(comm []byte, data []byte) error {
	blobID, err := parseBlobIDFromCommitment(comm)
	if err != nil {
		return fmt.Errorf("failed to parse requested commitment for fallback verification: %w", err)
	}

	derivedCommitment, err := d.store.CreateCommitment(data)
	if err != nil {
		return fmt.Errorf("failed to reconstruct blob commitment from fallback data: %w", err)
	}

	if !bytes.Equal(blobID.Commitment, derivedCommitment) {
		return fmt.Errorf("fallback data commitment mismatch")
	}

	return nil
}

// handleGetError handles errors from getBlob and writes appropriate HTTP response.
func (d *CelestiaServer) handleGetError(w http.ResponseWriter, comm []byte, err error, start time.Time) {
	if d.metrics != nil {
		d.metrics.RecordRetrievalError()
		d.metrics.RecordHTTPRequest("get", time.Since(start))
	}

	if isNotFoundError(err) {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	d.log.Error("Failed to retrieve blob", "commitment", hex.EncodeToString(comm), "err", err)
	w.WriteHeader(http.StatusInternalServerError)
}

// recordGetSuccess records metrics for successful blob retrieval.
func (d *CelestiaServer) recordGetSuccess(start time.Time, size int) {
	if d.metrics == nil {
		return
	}
	d.metrics.RecordRetrieval(time.Since(start), size)
	d.metrics.RecordHTTPRequest("get", time.Since(start))
}

// HandlePut submits a blob to Celestia and returns the commitment.
// Request: PUT /put with blob data in body
// Response: GenericCommitment on success, 500 on error (let op-batcher retry).
func (d *CelestiaServer) HandlePut(w http.ResponseWriter, r *http.Request) {
	// Accept both PUT and POST for compatibility
	if r.Method != http.MethodPut && r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	start := time.Now()
	d.log.Debug("PUT", "url", r.URL)

	route := path.Base(r.URL.Path)
	if route != "put" {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	// Limit request body to prevent memory exhaustion
	r.Body = http.MaxBytesReader(w, r.Body, d.maxBlobSize)

	input, err := io.ReadAll(r.Body)
	if err != nil {
		if err.Error() == "http: request body too large" {
			d.log.Warn("Request body too large", "limit", d.maxBlobSize)
			http.Error(w, "blob exceeds maximum size", http.StatusRequestEntityTooLarge)
			return
		}
		d.log.Error("Failed to read request body", "err", err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	// Validate non-empty blob
	if len(input) == 0 {
		d.log.Warn("Empty blob rejected")
		http.Error(w, "empty blob not allowed", http.StatusBadRequest)
		return
	}

	// Create context with configurable timeout
	ctx, cancel := context.WithTimeout(r.Context(), d.submitTimeout)
	defer cancel()

	// Submit blob directly to Celestia (stateless - synchronous, blocks until confirmed)
	commitment, _, err := d.store.Put(ctx, input)
	if err != nil {
		// Record metrics for failed submission
		if d.metrics != nil {
			d.metrics.RecordSubmissionError()
			d.metrics.RecordHTTPRequest("put", time.Since(start))
		}

		d.log.Error("Failed to submit blob to Celestia",
			"size", len(input),
			"err", err)
		// Return 500 with error message - let op-batcher retry
		http.Error(w, fmt.Sprintf("submission failed: %v", err), http.StatusInternalServerError)
		return
	}

	// Record metrics for successful submission
	duration := time.Since(start)
	if d.metrics != nil {
		d.metrics.RecordSubmission(duration, len(input))
		d.metrics.RecordHTTPRequest("put", duration)

		// Parse blob ID to get height for metrics
		var blobID CelestiaBlobID
		// Skip first 2 bytes which are frame version and altda version
		if err := blobID.UnmarshalBinary(commitment[2:]); err == nil {
			d.metrics.SetInclusionHeight(blobID.Height)
		}
	}

	d.log.Info("Blob submitted successfully",
		"commitment", hex.EncodeToString(commitment),
		"size", len(input),
		"duration", duration)

	// Write to fallback provider asynchronously (non-blocking)
	if d.fallback.Available() && fallback.WriteEnabled(d.fallbackMode) {
		d.wg.Add(1)
		go d.putFallback(context.Background(), commitment, input)
	}

	w.Header().Set("Content-Type", "application/octet-stream")
	if _, err := w.Write(commitment); err != nil {
		d.log.Error("Failed to write response", "err", err)
	}
}

func (d *CelestiaServer) Endpoint() string {
	return d.listener.Addr().String()
}

// getFallback retrieves data from fallback provider (synchronous).
func (d *CelestiaServer) getFallback(ctx context.Context, commitment []byte) ([]byte, error) {
	ctx, cancel := context.WithTimeout(ctx, d.fallback.Timeout())
	defer cancel()
	return d.fallback.Get(ctx, commitment)
}

// putFallback writes to fallback provider. Must be called with `go` after `d.wg.Add(1)`.
// Caller provides parent context (typically context.Background() for async operations).
func (d *CelestiaServer) putFallback(ctx context.Context, commitment []byte, data []byte) {
	defer d.wg.Done()

	ctx, cancel := context.WithTimeout(ctx, d.fallback.Timeout())
	defer cancel()

	if err := d.fallback.Put(ctx, commitment, data); err != nil {
		d.log.Warn("Fallback write failed", "provider", d.fallback.Name(), "err", err)
		if d.metrics != nil {
			d.metrics.RecordFallbackWriteError()
		}
	} else {
		d.log.Debug("Fallback write succeeded", "provider", d.fallback.Name())
		if d.metrics != nil {
			d.metrics.RecordFallbackWrite()
		}
	}
}

func (d *CelestiaServer) Stop() error {
	var errs []error

	// Shutdown HTTP server first
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := d.httpServer.Shutdown(ctx); err != nil {
		d.log.Error("HTTP server shutdown error", "err", err)
		errs = append(errs, fmt.Errorf("http server shutdown: %w", err))
	}

	// Shutdown metrics server if running
	if d.metricsServer != nil {
		if err := d.metricsServer.Shutdown(ctx); err != nil {
			d.log.Error("Metrics server shutdown error", "err", err)
			errs = append(errs, fmt.Errorf("metrics server shutdown: %w", err))
		}
	}

	// Wait for pending async operations (fallback writes) to complete
	done := make(chan struct{})
	go func() {
		d.wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		d.log.Info("All pending fallback operations completed")
	case <-time.After(10 * time.Second):
		d.log.Warn("Timeout waiting for pending fallback operations")
		errs = append(errs, fmt.Errorf("timeout waiting for pending fallback operations"))
	}

	return errors.Join(errs...)
}

// isNotFoundError checks if the error indicates the blob was not found.
func isNotFoundError(err error) bool {
	if err == nil {
		return false
	}
	return errors.Is(err, altda.ErrNotFound)
}