[Feature Implementation] IP / rate limiting through gRPC API

## Overview

This is an implementation of a limiter for IP (device) limiting and rate (speed) limiting.

## Files to change

**Modification:**

```
app/dispatcher/default.go
```

**Creation:**

```
app/limiter/limiter.go
app/limiter/rate.go
```

**Recompilation:**

```
common/protocol/user.go
common/protocol/user.pb.go
common/protocol/user.proto
```

## Implementation

**Modify following snippets in app/dispatcher/default.go**

```go
# ...

import (
	"context"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/xtls/xray-core/common"
	"github.com/xtls/xray-core/common/buf"
	"github.com/xtls/xray-core/common/errors"
	"github.com/xtls/xray-core/common/log"
	"github.com/xtls/xray-core/common/net"
	"github.com/xtls/xray-core/common/protocol"
	"github.com/xtls/xray-core/common/session"
	"github.com/xtls/xray-core/core"
	"github.com/xtls/xray-core/features/dns"
	"github.com/xtls/xray-core/features/outbound"
	"github.com/xtls/xray-core/features/policy"
	"github.com/xtls/xray-core/features/routing"
	routing_session "github.com/xtls/xray-core/features/routing/session"
	"github.com/xtls/xray-core/features/stats"
	"github.com/xtls/xray-core/transport"
	"github.com/xtls/xray-core/transport/pipe"

	// Start of IP limit and rate limit
	"github.com/xtls/xray-core/app/limiter"
	// End
)

# ...

type DefaultDispatcher struct {
	ohm    outbound.Manager
	router routing.Router
	policy policy.Manager
	stats  stats.Manager
	dns    dns.Client
	fdns   dns.FakeDNSEngine

	// Start of IP limit and rate limit
	limiter *limiter.Limiter
	// End
}

# ...

func (d *DefaultDispatcher) Init(config *Config, om outbound.Manager, router routing.Router, pm policy.Manager, sm stats.Manager, dns dns.Client) error {
	d.ohm = om
	d.router = router
	d.policy = pm
	d.stats = sm
	d.dns = dns

	// Start of IP limit and rate limit
	d.limiter = limiter.New()
	// End

	return nil
}

# ...

func (d *DefaultDispatcher) getLink(ctx context.Context) (*transport.Link, *transport.Link) {
	opt := pipe.OptionsFromContext(ctx)
	uplinkReader, uplinkWriter := pipe.New(opt...)
	downlinkReader, downlinkWriter := pipe.New(opt...)

	inboundLink := &transport.Link{
		Reader: downlinkReader,
		Writer: uplinkWriter,
	}

	outboundLink := &transport.Link{
		Reader: uplinkReader,
		Writer: downlinkWriter,
	}

	sessionInbound := session.InboundFromContext(ctx)
	var user *protocol.MemoryUser
	if sessionInbound != nil {
		user = sessionInbound.User
	}

	if user != nil && len(user.Email) > 0 {
		// Start of IP limit and rate limit
		bucket, ok, reject := d.limiter.Get(sessionInbound.Tag, user.Email, user.DeviceLimit, user.SpeedLimit, sessionInbound.Source.Address.IP().String())
		if reject {
			errors.LogWarning(ctx, "IP limit exceeded: ", user.Email)
			common.Close(outboundLink.Writer)
			common.Close(inboundLink.Writer)
			common.Interrupt(outboundLink.Reader)
			common.Interrupt(inboundLink.Reader)
			return inboundLink, outboundLink
		}
		if ok {
			inboundLink.Writer = d.limiter.RateWriter(inboundLink.Writer, bucket)
			outboundLink.Writer = d.limiter.RateWriter(outboundLink.Writer, bucket)
		}
		// End

		p := d.policy.ForLevel(user.Level)
		if p.Stats.UserUplink {
			name := "user>>>" + user.Email + ">>>traffic>>>uplink"
			if c, _ := stats.GetOrRegisterCounter(d.stats, name); c != nil {
				inboundLink.Writer = &SizeStatWriter{
					Counter: c,
					Writer:  inboundLink.Writer,
				}
			}
		}
		if p.Stats.UserDownlink {
			name := "user>>>" + user.Email + ">>>traffic>>>downlink"
			if c, _ := stats.GetOrRegisterCounter(d.stats, name); c != nil {
				outboundLink.Writer = &SizeStatWriter{
					Counter: c,
					Writer:  outboundLink.Writer,
				}
			}
		}
	}

	return inboundLink, outboundLink
}

# ...
```

****
**Create app/limiter/limiter.go**

```go
package limiter

import (
	"sync"
	"time"

	"golang.org/x/time/rate"
)

type Limiter struct {
	InboundInfo *sync.Map // Key: Tag, Value: *InboundInfo
	stopChan    chan struct{}
	wg          sync.WaitGroup
}

type InboundInfo struct {
	Tag          string
	UserIPs      *sync.Map // Key: Email, Value: *{Key: IP, Value: Timestamp}
	RateLimiters *sync.Map // Key: Email, Value: *rate.Limiter
}

func New() *Limiter {
	l := &Limiter{
		InboundInfo: new(sync.Map),
		stopChan:    make(chan struct{}),
	}

	l.wg.Add(1)

	go l.Start(5 * time.Minute)

	return l
}

func (l *Limiter) Start(interval time.Duration) {
	defer l.wg.Done()

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			l.Clean(10 * time.Minute)
		case <-l.stopChan:
			return
		}
	}
}

func (l *Limiter) Stop() {
	close(l.stopChan)
	l.wg.Wait()
}

func (l *Limiter) Get(tag string, email string, ipLimit uint32, rateLimit uint64, ip string) (*rate.Limiter, bool, bool) {
	inboundInfoValue, _ := l.InboundInfo.LoadOrStore(tag, &InboundInfo{
		Tag:          tag,
		UserIPs:      new(sync.Map),
		RateLimiters: new(sync.Map),
	})
	inboundInfo := inboundInfoValue.(*InboundInfo)

	ipMapsValue, _ := inboundInfo.UserIPs.LoadOrStore(email, new(sync.Map))
	ipMaps := ipMapsValue.(*sync.Map)

	_, ipExists := ipMaps.LoadOrStore(ip, new(int64))

	// Record IP access
	timestamp := time.Now().Unix()
	ipMaps.Store(ip, timestamp)

	// Enforce IP limit
	if !ipExists && ipLimit > 0 {
		var ipCount uint32

		ipMaps.Range(func(_, _ interface{}) bool {
			ipCount++
			return true
		})

		if ipCount > ipLimit {
			ipMaps.Delete(ip)

			return nil, false, true
		}
	}

	// Enforce rate limit
	if rateLimit > 0 {
		if rateLimiter, emailExists := inboundInfo.RateLimiters.Load(email); emailExists {
			return rateLimiter.(*rate.Limiter), true, false
		}

		bucket := rate.NewLimiter(rate.Limit(rateLimit), int(rateLimit))
		inboundInfo.RateLimiters.Store(email, bucket)

		return bucket, true, false
	}

	return nil, false, false
}

func (l *Limiter) Clean(timeout time.Duration) {
	expirationTime := time.Now().Add(-timeout).Unix()

	l.InboundInfo.Range(func(_, value interface{}) bool {
		inboundInfo := value.(*InboundInfo)

		var emailsToDelete []string

		inboundInfo.UserIPs.Range(func(key, value interface{}) bool {
			email := key.(string)
			ipMaps := value.(*sync.Map)

			var ipsToDelete []string

			ipMaps.Range(func(key, value interface{}) bool {
				ip := key.(string)
				ipTimestamp := value.(int64)

				if ipTimestamp < expirationTime {
					ipsToDelete = append(ipsToDelete, ip)
				}

				return true
			})

			for _, ip := range ipsToDelete {
				ipMaps.Delete(ip)
			}

			var ipCount uint32

			ipMaps.Range(func(_, _ interface{}) bool {
				ipCount++
				return true
			})

			if ipCount == 0 {
				emailsToDelete = append(emailsToDelete, email)
			}

			return true
		})

		for _, email := range emailsToDelete {
			inboundInfo.UserIPs.Delete(email)
			inboundInfo.RateLimiters.Delete(email)
		}

		return true
	})
}
```

****
**Create app/limiter/rate.go**


```go
package limiter

import (
	"context"
	"io"

	"github.com/xtls/xray-core/common"
	"github.com/xtls/xray-core/common/buf"
	"golang.org/x/time/rate"
)

type Writer struct {
	writer  buf.Writer
	limiter *rate.Limiter
	w       io.Writer
}

func (l *Limiter) RateWriter(writer buf.Writer, limiter *rate.Limiter) buf.Writer {
	return &Writer{
		writer:  writer,
		limiter: limiter,
	}
}

func (w *Writer) Close() error {
	return common.Close(w.writer)
}

func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error {
	ctx := context.Background()
	w.limiter.WaitN(ctx, int(mb.Len()))
	return w.writer.WriteMultiBuffer(mb)
}
```

****
**Replace common/protocol/user.go with following**


```go
package protocol

import "github.com/xtls/xray-core/common/errors"

func (u *User) GetTypedAccount() (Account, error) {
	if u.GetAccount() == nil {
		return nil, errors.New("Account missing").AtWarning()
	}

	rawAccount, err := u.Account.GetInstance()
	if err != nil {
		return nil, err
	}
	if asAccount, ok := rawAccount.(AsAccount); ok {
		return asAccount.AsAccount()
	}
	if account, ok := rawAccount.(Account); ok {
		return account, nil
	}
	return nil, errors.New("Unknown account type: ", u.Account.Type)
}

func (u *User) ToMemoryUser() (*MemoryUser, error) {
	account, err := u.GetTypedAccount()
	if err != nil {
		return nil, err
	}
	return &MemoryUser{
		// Reserved for global device limit
		ID: u.ID,

		Account: account,
		Email:   u.Email,
		Level:   u.Level,

		// Device limit and speed limit
		DeviceLimit: u.DeviceLimit,
		SpeedLimit:  u.SpeedLimit,
	}, nil
}

// MemoryUser is a parsed form of User, to reduce number of parsing of Account proto.
type MemoryUser struct {
	// Reserved for global device limit
	ID uint32

	// Account is the parsed account of the protocol.
	Account Account
	Email   string
	Level   uint32

	// Device limit and speed limit
	DeviceLimit uint32
	SpeedLimit  uint64
}
```

****
**Replace common/protocol/user.pb.go with following**


```go
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// 	protoc-gen-go v1.25.0-devel
// 	protoc        v3.14.0
// source: common/protocol/user.proto

package protocol

import (
	serial "github.com/xtls/xray-core/common/serial"
	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
	reflect "reflect"
	sync "sync"
)

const (
	// Verify that this generated code is sufficiently up-to-date.
	_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
	// Verify that runtime/protoimpl is sufficiently up-to-date.
	_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
)

// User is a generic user for all procotols.
type User struct {
	state         protoimpl.MessageState
	sizeCache     protoimpl.SizeCache
	unknownFields protoimpl.UnknownFields

	// Device limit and speed limit
	SpeedLimit  uint64 `protobuf:"varint,1,opt,name=speed_limit,json=speedLimit,proto3" json:"speed_limit,omitempty"`
	DeviceLimit uint32 `protobuf:"varint,2,opt,name=device_limit,json=deviceLimit,proto3" json:"device_limit,omitempty"`

	Level       uint32 `protobuf:"varint,3,opt,name=level,proto3" json:"level,omitempty"`
	Email       string `protobuf:"bytes,4,opt,name=email,proto3" json:"email,omitempty"`
	// Protocol specific account information. Must be the account proto in one of
	// the proxies.
	Account *serial.TypedMessage `protobuf:"bytes,5,opt,name=account,proto3" json:"account,omitempty"`
	// Reserved for global device limit
	ID uint32 `protobuf:"varint,6,opt,name=id,proto3" json:"id,omitempty"`
}

func (x *User) Reset() {
	*x = User{}
	if protoimpl.UnsafeEnabled {
		mi := &file_common_protocol_user_proto_msgTypes[0]
		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
		ms.StoreMessageInfo(mi)
	}
}

func (x *User) String() string {
	return protoimpl.X.MessageStringOf(x)
}

func (*User) ProtoMessage() {}

func (x *User) ProtoReflect() protoreflect.Message {
	mi := &file_common_protocol_user_proto_msgTypes[0]
	if protoimpl.UnsafeEnabled && x != nil {
		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
		if ms.LoadMessageInfo() == nil {
			ms.StoreMessageInfo(mi)
		}
		return ms
	}
	return mi.MessageOf(x)
}

// Deprecated: Use User.ProtoReflect.Descriptor instead.
func (*User) Descriptor() ([]byte, []int) {
	return file_common_protocol_user_proto_rawDescGZIP(), []int{0}
}

func (x *User) GetSpeedLimit() uint64 {
	if x != nil {
		return x.SpeedLimit
	}
	return 0
}

func (x *User) GetDeviceLimit() uint32 {
	if x != nil {
		return x.DeviceLimit
	}
	return 0
}

func (x *User) GetLevel() uint32 {
	if x != nil {
		return x.Level
	}
	return 0
}

func (x *User) GetEmail() string {
	if x != nil {
		return x.Email
	}
	return ""
}

func (x *User) GetAccount() *serial.TypedMessage {
	if x != nil {
		return x.Account
	}
	return nil
}

func (x *User) GetID() uint32 {
	if x != nil {
		return x.ID
	}
	return 0
}

var File_common_protocol_user_proto protoreflect.FileDescriptor

var file_common_protocol_user_proto_rawDesc = []byte{
	0x0a, 0x1a, 0x63, 0x6f, 0x6d, 0x6d, 0x6f, 0x6e, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f,
	0x6c, 0x2f, 0x75, 0x73, 0x65, 0x72, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x14, 0x78, 0x72,
	0x61, 0x79, 0x2e, 0x63, 0x6f, 0x6d, 0x6d, 0x6f, 0x6e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63,
	0x6f, 0x6c, 0x1a, 0x21, 0x63, 0x6f, 0x6d, 0x6d, 0x6f, 0x6e, 0x2f, 0x73, 0x65, 0x72, 0x69, 0x61,
	0x6c, 0x2f, 0x74, 0x79, 0x70, 0x65, 0x64, 0x5f, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x2e,
	0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0xc2, 0x01, 0x0a, 0x04, 0x55, 0x73, 0x65, 0x72, 0x12, 0x1f,
	0x0a, 0x0b, 0x73, 0x70, 0x65, 0x65, 0x64, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x18, 0x01, 0x20,
	0x01, 0x28, 0x04, 0x52, 0x0a, 0x73, 0x70, 0x65, 0x65, 0x64, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x12,
	0x21, 0x0a, 0x0c, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x18,
	0x02, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x0b, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x4c, 0x69, 0x6d,
	0x69, 0x74, 0x12, 0x14, 0x0a, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x18, 0x03, 0x20, 0x01, 0x28,
	0x0d, 0x52, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x14, 0x0a, 0x05, 0x65, 0x6d, 0x61, 0x69,
	0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x65, 0x6d, 0x61, 0x69, 0x6c, 0x12, 0x3a,
	0x0a, 0x07, 0x61, 0x63, 0x63, 0x6f, 0x75, 0x6e, 0x74, 0x18, 0x05, 0x20, 0x01, 0x28, 0x0b, 0x32,
	0x20, 0x2e, 0x78, 0x72, 0x61, 0x79, 0x2e, 0x63, 0x6f, 0x6d, 0x6d, 0x6f, 0x6e, 0x2e, 0x73, 0x65,
	0x72, 0x69, 0x61, 0x6c, 0x2e, 0x54, 0x79, 0x70, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67,
	0x65, 0x52, 0x07, 0x61, 0x63, 0x63, 0x6f, 0x75, 0x6e, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64,
	0x18, 0x06, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x02, 0x69, 0x64, 0x42, 0x5e, 0x0a, 0x18, 0x63, 0x6f,
	0x6d, 0x2e, 0x78, 0x72, 0x61, 0x79, 0x2e, 0x63, 0x6f, 0x6d, 0x6d, 0x6f, 0x6e, 0x2e, 0x70, 0x72,
	0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x50, 0x01, 0x5a, 0x29, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62,
	0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x78, 0x74, 0x6c, 0x73, 0x2f, 0x78, 0x72, 0x61, 0x79, 0x2d, 0x63,
	0x6f, 0x72, 0x65, 0x2f, 0x63, 0x6f, 0x6d, 0x6d, 0x6f, 0x6e, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f,
	0x63, 0x6f, 0x6c, 0xaa, 0x02, 0x14, 0x58, 0x72, 0x61, 0x79, 0x2e, 0x43, 0x6f, 0x6d, 0x6d, 0x6f,
	0x6e, 0x2e, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74,
	0x6f, 0x33,
}

var (
	file_common_protocol_user_proto_rawDescOnce sync.Once
	file_common_protocol_user_proto_rawDescData = file_common_protocol_user_proto_rawDesc
)

func file_common_protocol_user_proto_rawDescGZIP() []byte {
	file_common_protocol_user_proto_rawDescOnce.Do(func() {
		file_common_protocol_user_proto_rawDescData = protoimpl.X.CompressGZIP(file_common_protocol_user_proto_rawDescData)
	})
	return file_common_protocol_user_proto_rawDescData
}

var file_common_protocol_user_proto_msgTypes = make([]protoimpl.MessageInfo, 1)
var file_common_protocol_user_proto_goTypes = []interface{}{
	(*User)(nil),                // 0: xray.common.protocol.User
	(*serial.TypedMessage)(nil), // 1: xray.common.serial.TypedMessage
}
var file_common_protocol_user_proto_depIdxs = []int32{
	1, // 0: xray.common.protocol.User.account:type_name -> xray.common.serial.TypedMessage
	1, // [1:1] is the sub-list for method output_type
	1, // [1:1] is the sub-list for method input_type
	1, // [1:1] is the sub-list for extension type_name
	1, // [1:1] is the sub-list for extension extendee
	0, // [0:1] is the sub-list for field type_name
}

func init() { file_common_protocol_user_proto_init() }
func file_common_protocol_user_proto_init() {
	if File_common_protocol_user_proto != nil {
		return
	}
	if !protoimpl.UnsafeEnabled {
		file_common_protocol_user_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
			switch v := v.(*User); i {
			case 0:
				return &v.state
			case 1:
				return &v.sizeCache
			case 2:
				return &v.unknownFields
			default:
				return nil
			}
		}
	}
	type x struct{}
	out := protoimpl.TypeBuilder{
		File: protoimpl.DescBuilder{
			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
			RawDescriptor: file_common_protocol_user_proto_rawDesc,
			NumEnums:      0,
			NumMessages:   1,
			NumExtensions: 0,
			NumServices:   0,
		},
		GoTypes:           file_common_protocol_user_proto_goTypes,
		DependencyIndexes: file_common_protocol_user_proto_depIdxs,
		MessageInfos:      file_common_protocol_user_proto_msgTypes,
	}.Build()
	File_common_protocol_user_proto = out.File
	file_common_protocol_user_proto_rawDesc = nil
	file_common_protocol_user_proto_goTypes = nil
	file_common_protocol_user_proto_depIdxs = nil
}
```

****
**Replace common/protocol/user.proto with following**


```go
syntax = "proto3";

package xray.common.protocol;
option csharp_namespace = "Xray.Common.Protocol";
option go_package = "github.com/xtls/xray-core/common/protocol";
option java_package = "com.xray.common.protocol";
option java_multiple_files = true;

import "common/serial/typed_message.proto";

// User is a generic user for all procotols.
message User {
  // Device limit and speed limit
  uint64 speed_limit = 1;
  uint32 device_limit = 2;
  
  uint32 level = 3;
  string email = 4;

  // Protocol specific account information. Must be the account proto in one of
  // the proxies.
  xray.common.serial.TypedMessage account = 5;

  // Reserved for global device limit
  uint32 id = 6;
}
```

## Credits

  - [XTLS/Xray-core](https://github.com/XTLS/Xray-core)
  - [XrayR-project/XrayR](https://github.com/XrayR-project/XrayR)
  - [crossfw/Xray-API-documents](https://github.com/crossfw/Xray-API-documents)
  - [调用 V2Ray 提供的 API 接口进行用户增删及流量统计查询操作](https://tachyondevel.medium.com/%E8%B0%83%E7%94%A8-v2ray-%E6%8F%90%E4%BE%9B%E7%9A%84-api-%E6%8E%A5%E5%8F%A3%E8%BF%9B%E8%A1%8C%E7%94%A8%E6%88%B7%E5%A2%9E%E5%88%A0%E6%93%8D%E4%BD%9C-adf9ff972973)
  - [利用v2ray api实现流量统计以及增删用户](https://www.mebi.me/1684)


Navigation Menu

Uh oh!

[Feature Implementation] IP / rate limiting through gRPC API #3667

Overview

Files to change

Implementation

Credits

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Uh oh!

[Feature Implementation] IP / rate limiting through gRPC API #3667

Description

Overview

Files to change

Implementation

Credits

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions