From 45b7ac2021fedecd571ef58c74f9ad24e1417951 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Wed, 14 Sep 2022 00:14:20 +0800
Subject: [PATCH] add conn limiter

---
 common/net/dialer/dialer.go      |  2 +-
 limiter/conn/limiter.go          | 10 ++++++++++
 limiter/{ => traffic}/limiter.go |  4 ++--
 listener/option.go               | 34 ++++++++++++++++++++------------
 4 files changed, 34 insertions(+), 16 deletions(-)
 create mode 100644 limiter/conn/limiter.go
 rename limiter/{ => traffic}/limiter.go (84%)

diff --git a/common/net/dialer/dialer.go b/common/net/dialer/dialer.go
index cb7e7aa..c86a494 100644
--- a/common/net/dialer/dialer.go
+++ b/common/net/dialer/dialer.go
@@ -11,7 +11,7 @@ import (
 )
 
 const (
-	DefaultTimeout = 15 * time.Second
+	DefaultTimeout = 10 * time.Second
 )
 
 var (
diff --git a/limiter/conn/limiter.go b/limiter/conn/limiter.go
new file mode 100644
index 0000000..6c53e73
--- /dev/null
+++ b/limiter/conn/limiter.go
@@ -0,0 +1,10 @@
+package conn
+
+type Limiter interface {
+	Allow(n int) bool
+	Limit() int
+}
+
+type ConnLimiter interface {
+	Limiter(key string) Limiter
+}
diff --git a/limiter/limiter.go b/limiter/traffic/limiter.go
similarity index 84%
rename from limiter/limiter.go
rename to limiter/traffic/limiter.go
index d1b0f83..b439b58 100644
--- a/limiter/limiter.go
+++ b/limiter/traffic/limiter.go
@@ -1,4 +1,4 @@
-package limiter
+package traffic
 
 import "context"
 
@@ -9,7 +9,7 @@ type Limiter interface {
 	Limit() int
 }
 
-type RateLimiter interface {
+type TrafficLimiter interface {
 	In(key string) Limiter
 	Out(key string) Limiter
 }
diff --git a/listener/option.go b/listener/option.go
index 1ae15f9..d3b05f1 100644
--- a/listener/option.go
+++ b/listener/option.go
@@ -7,21 +7,23 @@ import (
 	"github.com/go-gost/core/admission"
 	"github.com/go-gost/core/auth"
 	"github.com/go-gost/core/chain"
-	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/conn"
+	"github.com/go-gost/core/limiter/traffic"
 	"github.com/go-gost/core/logger"
 )
 
 type Options struct {
-	Addr          string
-	Auther        auth.Authenticator
-	Auth          *url.Userinfo
-	TLSConfig     *tls.Config
-	Admission     admission.Admission
-	RateLimiter   limiter.RateLimiter
-	Chain         chain.Chainer
-	Logger        logger.Logger
-	Service       string
-	ProxyProtocol int
+	Addr           string
+	Auther         auth.Authenticator
+	Auth           *url.Userinfo
+	TLSConfig      *tls.Config
+	Admission      admission.Admission
+	TrafficLimiter traffic.TrafficLimiter
+	ConnLimiter    conn.ConnLimiter
+	Chain          chain.Chainer
+	Logger         logger.Logger
+	Service        string
+	ProxyProtocol  int
 }
 
 type Option func(opts *Options)
@@ -56,9 +58,15 @@ func AdmissionOption(admission admission.Admission) Option {
 	}
 }
 
-func RateLimiterOption(rlimiter limiter.RateLimiter) Option {
+func TrafficLimiterOption(limiter traffic.TrafficLimiter) Option {
 	return func(opts *Options) {
-		opts.RateLimiter = rlimiter
+		opts.TrafficLimiter = limiter
+	}
+}
+
+func ConnLimiterOption(limiter conn.ConnLimiter) Option {
+	return func(opts *Options) {
+		opts.ConnLimiter = limiter
 	}
 }