specstream/pkg/pipeline/pipeline.go

// Package pipeline orchestrates the extraction flow:
//   1. Content type detection → select extractor
//   2. Extract structured content (typed, bounded output)
//   3. Filter for injection patterns (defense-in-depth)
//   4. Return clean ExtractedContent to caller
//
// The caller is the MCP server (HTTP JSON-RPC endpoint on port 3991).

package pipeline

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/sscoble/specstream/pkg/extract"
	"github.com/sscoble/specstream/pkg/filter"
	"github.com/sscoble/specstream/pkg/types"
)

// Pipeline is the top-level orchestrator.
// It detects content type, extracts structured data, and runs safety filters.
type Pipeline struct {
	extractors []extract.Extractor
	filters    []filter.FilterFunc
	maxSize    int64
	httpClient *http.Client
}

// New creates a pipeline with default extractors, filters, and a restricted HTTP client.
func New() *Pipeline {
	return &Pipeline{
		extractors: []extract.Extractor{
			extract.NewHTMLExtractor(),
			extract.NewMarkdownExtractor(),
			extract.NewTextExtractor(),
		},
		filters: filter.StandardFilters(),
		maxSize: 1 << 20, // 1MB
		httpClient: &http.Client{
			Timeout: 15 * time.Second,
			Transport: &http.Transport{
				DialContext: (&net.Dialer{
					Timeout:   5 * time.Second,
					KeepAlive: 0,
				}).DialContext,
				// Prevent redirects to private/internal IPs
				DisableKeepAlives: true,
			},
			CheckRedirect: func(req *http.Request, via []*http.Request) error {
				if len(via) >= 5 {
					return fmt.Errorf("too many redirects")
				}
				// Only allow http/https redirects
				if req.URL.Scheme != "http" && req.URL.Scheme != "https" {
					return fmt.Errorf("blocked redirect scheme: %s", req.URL.Scheme)
				}
				return nil
			},
		},
	}
}

// NewWithConfig creates a pipeline with custom extractors and filters.
func NewWithConfig(extractors []extract.Extractor, filters []filter.FilterFunc, maxSize int64, client *http.Client) *Pipeline {
	if maxSize <= 0 {
		maxSize = 1 << 20
	}
	if client == nil {
		client = &http.Client{Timeout: 15 * time.Second}
	}
	return &Pipeline{
		extractors: extractors,
		filters:    filters,
		maxSize:    maxSize,
		httpClient: client,
	}
}

// FetchAndExtract fetches a URL and runs it through the extraction pipeline.
func (p *Pipeline) FetchAndExtract(rawURL string, req types.ExtractRequest) (*types.ExtractedContent, error) {
	if rawURL == "" {
		return nil, fmt.Errorf("empty URL")
	}

	// Reject non-HTTP schemes
	if !strings.HasPrefix(rawURL, "http://") && !strings.HasPrefix(rawURL, "https://") {
		return nil, fmt.Errorf("blocked scheme: only http/https allowed")
	}

	// Fetch
	resp, err := p.httpClient.Get(rawURL)
	if err != nil {
		return nil, fmt.Errorf("http get: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("http %d: %s", resp.StatusCode, resp.Status)
	}

	// Read with size limit
	body, err := io.ReadAll(io.LimitReader(resp.Body, p.maxSize))
	if err != nil {
		return nil, fmt.Errorf("read body: %w", err)
	}

	req.URL = rawURL
	return p.Extract(body, req)
}

// Extract processes raw content bytes through the pipeline.
// Content-type detection is automatic. If filename is provided, it's used as a hint.
func (p *Pipeline) Extract(data []byte, req types.ExtractRequest) (*types.ExtractedContent, error) {
	if int64(len(data)) > p.maxSize {
		data = data[:p.maxSize]
	}

	// ── 1. Detect content type ──
	readerAt := bytes.NewReader(data)
	var selected extract.Extractor
	var selectedName string

	for _, ex := range p.extractors {
		ok, err := ex.Detect(readerAt, int64(len(data)))
		if err != nil {
			continue
		}
		if ok {
			selected = ex
			selectedName = ex.MIMETypes()[0]
			break
		}
	}

	// Fall back by filename extension
	if selected == nil && req.Filename != "" {
		ext := strings.ToLower(req.Filename)
		switch {
		case strings.HasSuffix(ext, ".html"), strings.HasSuffix(ext, ".htm"):
			selected = extract.NewHTMLExtractor()
			selectedName = "text/html"
		case strings.HasSuffix(ext, ".md"), strings.HasSuffix(ext, ".markdown"):
			selected = extract.NewMarkdownExtractor()
			selectedName = "text/markdown"
		default:
			selected = extract.NewTextExtractor()
			selectedName = "text/plain"
		}
	}

	// Final fallback: plain text extractor
	if selected == nil {
		selected = extract.NewTextExtractor()
		selectedName = "text/plain"
	}

	// ── 2. Extract structured content ──
	content, err := selected.Extract(bytes.NewReader(data), req)
	if err != nil {
		return nil, fmt.Errorf("extraction failed: %w", err)
	}

	// Fill in source info if missing
	if content.Source.MIMEType == "" {
		content.Source.MIMEType = selectedName
	}
	if content.Source.ByteSize == 0 {
		content.Source.ByteSize = int64(len(data))
	}

	// ── 3. Run safety filters ──
	for _, fn := range p.filters {
		warnings := fn(content)
		if len(warnings) > 0 {
			content.Metadata.Warnings = append(content.Metadata.Warnings, warnings...)
		}
	}

	return content, nil
}