gost-core/chain/node.go

package chain

import (
	"context"
	"regexp"
	"sync/atomic"

	"github.com/go-gost/core/auth"
	"github.com/go-gost/core/bypass"
	"github.com/go-gost/core/hosts"
	"github.com/go-gost/core/metadata"
	"github.com/go-gost/core/resolver"
	"github.com/go-gost/core/rewriter"
	"github.com/go-gost/core/routing"
	"github.com/go-gost/core/selector"
)

// NodeFilterSettings defines filtering criteria used to select nodes during routing.
type NodeFilterSettings struct {
	// Protocol filters by proxy protocol name (e.g. "http", "socks5").
	Protocol string
	// Host filters by hostname pattern.
	Host string
	// Path filters by URL path.
	Path string
}

// HTTPURLRewriteSetting defines an HTTP URL rewrite rule.
type HTTPURLRewriteSetting struct {
	// Pattern is the regex to match against the URL.
	Pattern *regexp.Regexp
	// Replacement is the replacement string.
	Replacement string
}

// FailCodes defines HTTP response status codes that mark a node as failed.
// The response is still relayed to the client as-is; marking only affects
// subsequent node selection (via the selector's FailFilter).
// Codes < 100 are wildcard hundred-level prefixes (e.g. 5 → 5xx).
type FailCodes []int

// Match reports whether statusCode matches any configured code.
// Exact codes are compared directly; codes < 100 match hundred-level
// groups (e.g. code 5 matches 500-599).
func (fc FailCodes) Match(statusCode int) bool {
	for _, code := range fc {
		if (code < 100 && statusCode/100 == code) || statusCode == code {
			return true
		}
	}
	return false
}

// HTTPBodyRewriteSettings defines an HTTP body rewrite rule.
type HTTPBodyRewriteSettings struct {
	// Type is the MIME type to match, e.g. "text/html".
	Type string
	// Pattern is the regex to match against the response body.
	Pattern *regexp.Regexp
	// Replacement is the replacement bytes.
	Replacement []byte
	// Rewriter is an optional plugin-based rewriter.
	// When set, Rewrite delegates to the plugin instead of using Pattern.
	Rewriter rewriter.Rewriter
	// MaxChunkSize limits buffering for chunked/unknown-length non-streaming
	// bodies. When > 0 and the body fits within this limit, the body is
	// buffered and rewritten; if exceeded, the body passes through unchanged.
	MaxChunkSize int
}

// HTTPHeaderRewriteSettings defines an HTTP header rewrite rule.
type HTTPHeaderRewriteSettings struct {
	// Name is the regex matched against the *lowercased* header name, so a
	// pattern with uppercase letters needs an explicit (?i) flag.
	// In plugin mode it gates whether the plugin is invoked.
	// Note: *regexp.Regexp (not string like HTTPBodyRewriteSettings.Type)
	// is intentional — header *names* need regex matching, while body rules
	// match a MIME prefix. Do not "align" this back to string.
	Name *regexp.Regexp
	// Pattern is the regex matched against each header value (regex mode);
	// in plugin mode it is unused.
	Pattern *regexp.Regexp
	// Replacement is the replacement bytes (regex mode).
	Replacement []byte
	// Rewriter is an optional plugin-based rewriter. When set, Rewrite
	// delegates to the plugin over the serialized header block. The plugin
	// receives a sorted, canonicalized header block (http.Header.Write
	// normalizes names and sorts keys), not the original wire bytes.
	Rewriter rewriter.Rewriter
}

// HTTPNodeSettings holds HTTP-level configuration for a node.
type HTTPNodeSettings struct {
	// Host is the HTTP Host header value. When HostPattern is set, Host is
	// treated as a replacement template: submatches captured by HostPattern
	// against req.URL.Path are expanded into it ($1, $2, ...). This restores
	// the upstream Host from the encoded path for mirror routing (e.g.
	// path "/microsoft.github.io/x" + HostPattern "^/([a-z0-9-]+\.github\.io)/"
	// + Host "$1" → Host "microsoft.github.io").
	Host string
	// HostPattern is the regex matched against req.URL.Path to derive the
	// Host. It only affects the Host header (never the dial target or TLS
	// SNI), so it is intended for Host-routed origins (GitHub Pages) rather
	// than SNI-routed origins. When nil, Host is used as a static value.
	HostPattern *regexp.Regexp
	// RequestHeader contains custom request headers to inject.
	RequestHeader map[string]string
	// ResponseHeader contains custom response headers to inject.
	ResponseHeader map[string]string
	// Auther is the HTTP authenticator for this node.
	Auther auth.Authenticator
	// RewriteURL holds the URL rewrite rules.
	RewriteURL []HTTPURLRewriteSetting
	// RewriteRequestBody holds the request body rewrite rules.
	RewriteRequestBody []HTTPBodyRewriteSettings
	// RewriteResponseBody holds the response body rewrite rules.
	RewriteResponseBody []HTTPBodyRewriteSettings
	// RewriteRequestHeader holds the request-header rewrite rules. The Host
	// header is not reachable here (net/http keeps it in Request.Host, not
	// Request.Header) — use Host above instead.
	RewriteRequestHeader []HTTPHeaderRewriteSettings
	// RewriteResponseHeader holds the response-header rewrite rules.
	RewriteResponseHeader []HTTPHeaderRewriteSettings
	// FailCodes lists response status codes that mark the node as failed.
	FailCodes FailCodes
}

// TLSNodeSettings holds TLS configuration for a node.
type TLSNodeSettings struct {
	// ServerName is the TLS SNI server name.
	ServerName string
	// Secure indicates whether a secure connection is used.
	Secure bool
	// Options specifies TLS version and cipher suite settings.
	Options struct {
		MinVersion   string
		MaxVersion   string
		CipherSuites []string
		ALPN         []string
	}
}

// NodeOptions holds the initialization parameters for a Node.
type NodeOptions struct {
	Network         string
	Transport       Transporter
	Bypass          bypass.Bypass
	Resolver        resolver.Resolver
	HostMapper      hosts.HostMapper
	Filter          *NodeFilterSettings
	HTTP            *HTTPNodeSettings
	TLS             *TLSNodeSettings
	Metadata        metadata.Metadata
	Matcher         routing.Matcher
	MatcherBodySize int
	Priority        int
}

// NodeOption is a functional option for configuring NodeOptions.
type NodeOption func(*NodeOptions)

// TransportNodeOption sets the Transporter for the node.
func TransportNodeOption(tr Transporter) NodeOption {
	return func(o *NodeOptions) {
		o.Transport = tr
	}
}

// BypassNodeOption sets the Bypass rule for the node.
func BypassNodeOption(bp bypass.Bypass) NodeOption {
	return func(o *NodeOptions) {
		o.Bypass = bp
	}
}

// ResolverNodeOption sets the Resolver for the node.
func ResolverNodeOption(resolver resolver.Resolver) NodeOption {
	return func(o *NodeOptions) {
		o.Resolver = resolver
	}
}

// HostMapperNodeOption sets the HostMapper for the node.
func HostMapperNodeOption(m hosts.HostMapper) NodeOption {
	return func(o *NodeOptions) {
		o.HostMapper = m
	}
}

// NetworkNodeOption sets the network type for the node.
func NetworkNodeOption(network string) NodeOption {
	return func(o *NodeOptions) {
		o.Network = network
	}
}

// NodeFilterOption sets the filter settings for the node.
func NodeFilterOption(filter *NodeFilterSettings) NodeOption {
	return func(o *NodeOptions) {
		o.Filter = filter
	}
}

// HTTPNodeOption sets the HTTP settings for the node.
func HTTPNodeOption(httpSettings *HTTPNodeSettings) NodeOption {
	return func(o *NodeOptions) {
		o.HTTP = httpSettings
	}
}

// TLSNodeOption sets the TLS settings for the node.
func TLSNodeOption(tlsSettings *TLSNodeSettings) NodeOption {
	return func(o *NodeOptions) {
		o.TLS = tlsSettings
	}
}

// MetadataNodeOption sets the Metadata for the node.
func MetadataNodeOption(md metadata.Metadata) NodeOption {
	return func(o *NodeOptions) {
		o.Metadata = md
	}
}

// MatcherNodeOption sets the routing Matcher for the node.
func MatcherNodeOption(matcher routing.Matcher) NodeOption {
	return func(o *NodeOptions) {
		o.Matcher = matcher
	}
}

// MatcherBodySizeNodeOption sets the max request body prefix (bytes) exposed
// to body matchers for this node. 0 disables body matching.
func MatcherBodySizeNodeOption(size int) NodeOption {
	return func(o *NodeOptions) {
		o.MatcherBodySize = size
	}
}

// PriorityNodeOption sets the priority for the node. Higher priority nodes
// are selected first during routing.
func PriorityNodeOption(priority int) NodeOption {
	return func(o *NodeOptions) {
		o.Priority = priority
	}
}

// Node represents a single hop in a proxy chain. Each node has a name,
// address, transport layer, and optional filters for protocol-specific routing.
// Nodes implement selector.Markable so that failed nodes can be temporarily
// excluded from selection.
type Node struct {
	Name    string
	Addr    string
	marker  selector.Marker
	options NodeOptions

	probeResult atomic.Value // *ProbeResult
	probeCancel context.CancelFunc
}

// NewNode creates a new Node with the given name and address.
func NewNode(name string, addr string, opts ...NodeOption) *Node {
	var options NodeOptions
	for _, opt := range opts {
		if opt != nil {
			opt(&options)
		}
	}

	return &Node{
		Name:    name,
		Addr:    addr,
		marker:  selector.NewFailMarker(),
		options: options,
	}
}

// Options returns the Node's configuration.
func (node *Node) Options() *NodeOptions {
	return &node.options
}

// Metadata returns the Node's metadata. Implements the metadata.Metadatable interface.
func (node *Node) Metadata() metadata.Metadata {
	return node.options.Metadata
}

// Marker returns the Node's failure marker. Implements the selector.Markable interface.
func (node *Node) Marker() selector.Marker {
	return node.marker
}

// Copy returns a shallow copy of the Node.
func (node *Node) Copy() *Node {
	n := &Node{}
	*n = *node
	return n
}

// ProbeResult implements chain.ProbeResultReader.
func (node *Node) ProbeResult() *ProbeResult {
	if v := node.probeResult.Load(); v != nil {
		return v.(*ProbeResult)
	}
	return nil
}

// SetProbeResult stores the latest probe result.
func (node *Node) SetProbeResult(r *ProbeResult) {
	node.probeResult.Store(r)
}

// SetProbeCancel stores the cancel function for the node's probe goroutine.
func (node *Node) SetProbeCancel(c context.CancelFunc) {
	node.probeCancel = c
}

// Close stops the node's probe goroutine if one is running.
func (node *Node) Close() error {
	if node.probeCancel != nil {
		node.probeCancel()
	}
	return nil
}