fix speedlimit

cmd/geph-exit/handle.go

@@ -178,7 +178,7 @@ func handle(rawClient net.Conn) {
 			return
 		}
 	case 'R':
-		err = handleResumable(limiter, slowLimit, tssClient)
+		err = handleResumable(slowLimit, tssClient)
 		log.Println("handleResumable returned with", err)
 		if err != nil {
 			tssClient.Close()
@@ -200,7 +200,7 @@ type scEntry struct {
 var sessionCache = make(map[[32]byte]*scEntry)
 var sessionCacheLock sync.Mutex
 
-func handleResumable(limiter *rate.Limiter, slowLimit bool, tssClient net.Conn) (err error) {
+func handleResumable(slowLimit bool, tssClient net.Conn) (err error) {
 	log.Println("handling resumable from", tssClient.RemoteAddr())
 	tssClient.SetDeadline(time.Now().Add(time.Second * 10))
 	var clientHello struct {
@@ -243,10 +243,6 @@ func handleResumable(limiter *rate.Limiter, slowLimit bool, tssClient net.Conn)
 		handle:   btcp,
 		currConn: tssClient,
 	}
-	if slowLimit {
-		log.Printf("[%v] slow limit", tssClient.RemoteAddr())
-		limiter = getLimiter(clientHello.MetaSess)
-	}
 	go func() {
 		defer func() {
 			sessionCacheLock.Lock()
@@ -260,8 +256,8 @@ func handleResumable(limiter *rate.Limiter, slowLimit bool, tssClient net.Conn)
 			KeepAliveInterval: time.Minute * 20,
 			KeepAliveTimeout:  time.Minute * 40,
 			MaxFrameSize:      32768,
-			MaxReceiveBuffer:  100 * 1024 * 1024,
-			MaxStreamBuffer:   100 * 1024 * 1024,
+			MaxReceiveBuffer:  10 * 1024 * 1024,
+			MaxStreamBuffer:   10 * 1024 * 1024,
 		})
 		if err != nil {
 			return
@@ -270,6 +266,12 @@ func handleResumable(limiter *rate.Limiter, slowLimit bool, tssClient net.Conn)
 			n, e = muxSrv.AcceptStream()
 			return
 		}
+		var limiter *rate.Limiter
+		if slowLimit {
+			limiter = slowLimitFactory.getLimiter(clientHello.MetaSess)
+		} else {
+			limiter = fastLimitFactory.getLimiter(clientHello.MetaSess)
+		}
 		smuxLoop(fmt.Sprintf("%x", clientHello.MetaSess), limiter, acceptStream)
 	}()
 	return

cmd/geph-exit/limitercache.go

@@ -5,14 +5,28 @@ import (
 	"golang.org/x/time/rate"
 )
 
-var lcache, _ = lru.New(16384)
+type limitFactory struct {
+	lcache *lru.Cache
+	limit  rate.Limit
+}
 
-func getLimiter(sessid [32]byte) *rate.Limiter {
-	new := rate.NewLimiter(100*1024, 1000*1024)
-	prev, _, _ := lcache.PeekOrAdd(sessid, new)
+func (lf *limitFactory) getLimiter(sessid [32]byte) *rate.Limiter {
+	new := rate.NewLimiter(lf.limit, 1000*1024)
+	prev, _, _ := lf.lcache.PeekOrAdd(sessid, new)
 	if prev != nil {
 		return prev.(*rate.Limiter)
 	} else {
 		return new
 	}
 }
+
+func newLimitFactory(limit rate.Limit) *limitFactory {
+	l, _ := lru.New(65536)
+	return &limitFactory{
+		limit:  limit,
+		lcache: l,
+	}
+}
+
+var slowLimitFactory = newLimitFactory(100 * 1024)
+var fastLimitFactory *limitFactory

cmd/geph-exit/main.go

@@ -19,6 +19,7 @@ import (
 	"github.com/geph-official/geph2/libs/pseudotcp"
 	"github.com/patrickmn/go-cache"
 	log "github.com/sirupsen/logrus"
+	"golang.org/x/time/rate"
 )
 
 var keyfile string
@@ -53,6 +54,7 @@ func main() {
 	flag.IntVar(&speedLimit, "speedLimit", 12500, "per-session speed limit, in KB/s")
 	flag.StringVar(&singleHop, "singleHop", "", "if supplied, runs in single-hop mode. (for example, -singleHop :5000 would listen on port 5000)")
 	flag.Parse()
+	fastLimitFactory = newLimitFactory(rate.Limit(speedLimit * 1024))
 	go func() {
 		log.Println(http.ListenAndServe("localhost:6060", nil))
 	}()