Xray-core/features/policy/policy.go

package policy

import (
	"context"
	"runtime"
	"time"

	"github.com/xtls/xray-core/common/platform"
	"github.com/xtls/xray-core/features"
)

// Timeout contains limits for connection timeout.
type Timeout struct {
	// Timeout for handshake phase in a connection.
	Handshake time.Duration
	// Timeout for connection being idle, i.e., there is no egress or ingress traffic in this connection.
	ConnectionIdle time.Duration
	// Timeout for an uplink only connection, i.e., the downlink of the connection has been closed.
	UplinkOnly time.Duration
	// Timeout for an downlink only connection, i.e., the uplink of the connection has been closed.
	DownlinkOnly time.Duration
}

// Stats contains settings for stats counters.
type Stats struct {
	// Whether or not to enable stat counter for user uplink traffic.
	UserUplink bool
	// Whether or not to enable stat counter for user downlink traffic.
	UserDownlink bool
}

// Buffer contains settings for internal buffer.
type Buffer struct {
	// Size of buffer per connection, in bytes. -1 for unlimited buffer.
	PerConnection int32
}

// Buffer contains settings for restriction such as ip restriction.
type Restriction struct {
	// Maximum allowed ips, -1 for unlimited
	MaxIPs int32
}

// SystemStats contains stat policy settings on system level.
type SystemStats struct {
	// Whether or not to enable stat counter for uplink traffic in inbound handlers.
	InboundUplink bool
	// Whether or not to enable stat counter for downlink traffic in inbound handlers.
	InboundDownlink bool
	// Whether or not to enable stat counter for uplink traffic in outbound handlers.
	OutboundUplink bool
	// Whether or not to enable stat counter for downlink traffic in outbound handlers.
	OutboundDownlink bool
}

// System contains policy settings at system level.
type System struct {
	Stats  SystemStats
	Buffer Buffer
}

// Session is session based settings for controlling Xray requests. It contains various settings (or limits) that may differ for different users in the context.
type Session struct {
	Timeouts    Timeout // Timeout settings
	Stats       Stats
	Buffer      Buffer
	Restriction Restriction
}

// Manager is a feature that provides Policy for the given user by its id or level.
//
// xray:api:stable
type Manager interface {
	features.Feature

	// ForLevel returns the Session policy for the given user level.
	ForLevel(level uint32) Session

	// ForSystem returns the System policy for Xray system.
	ForSystem() System
}

// ManagerType returns the type of Manager interface. Can be used to implement common.HasType.
//
// xray:api:stable
func ManagerType() interface{} {
	return (*Manager)(nil)
}

var defaultBufferSize int32

func init() {
	const defaultValue = -17
	size := platform.NewEnvFlag(platform.BufferSize).GetValueAsInt(defaultValue)

	switch size {
	case 0:
		defaultBufferSize = -1 // For pipe to use unlimited size
	case defaultValue: // Env flag not defined. Use default values per CPU-arch.
		switch runtime.GOARCH {
		case "arm", "mips", "mipsle":
			defaultBufferSize = 0
		case "arm64", "mips64", "mips64le":
			defaultBufferSize = 4 * 1024 // 4k cache for low-end devices
		default:
			defaultBufferSize = 512 * 1024
		}
	default:
		defaultBufferSize = int32(size) * 1024 * 1024
	}
}

func defaultBufferPolicy() Buffer {
	return Buffer{
		PerConnection: defaultBufferSize,
	}
}

func defaultRestrictionPolicy() Restriction {
	return Restriction{
		MaxIPs: -1,
	}
}

// SessionDefault returns the Policy when user is not specified.
func SessionDefault() Session {
	return Session{
		Timeouts: Timeout{
			// Align Handshake timeout with nginx client_header_timeout
			// So that this value will not indicate server identity
			Handshake:      time.Second * 60,
			ConnectionIdle: time.Second * 300,
			UplinkOnly:     time.Second * 1,
			DownlinkOnly:   time.Second * 1,
		},
		Stats: Stats{
			UserUplink:   false,
			UserDownlink: false,
		},
		Buffer:      defaultBufferPolicy(),
		Restriction: defaultRestrictionPolicy(),
	}
}

type policyKey int32

const (
	bufferPolicyKey policyKey = 0
)

func ContextWithBufferPolicy(ctx context.Context, p Buffer) context.Context {
	return context.WithValue(ctx, bufferPolicyKey, p)
}

func BufferPolicyFromContext(ctx context.Context) Buffer {
	pPolicy := ctx.Value(bufferPolicyKey)
	if pPolicy == nil {
		return defaultBufferPolicy()
	}
	return pPolicy.(Buffer)
}