mirror of https://github.com/fatedier/frp.git
support bandwith limit for one proxy
fatedier
5 years ago
parent
22a79710d8
commit
6da093a402
@ -0,0 +1,100 @@
|
||||
// Copyright 2019 fatedier, fatedier@gmail.com
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package config
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
MB = 1024 * 1024
|
||||
KB = 1024
|
||||
)
|
||||
|
||||
type BandwithQuantity struct {
|
||||
s string // MB or KB
|
||||
|
||||
i int64 // bytes
|
||||
}
|
||||
|
||||
func NewBandwithQuantity(s string) (BandwithQuantity, error) {
|
||||
q := BandwithQuantity{}
|
||||
err := q.UnmarshalString(s)
|
||||
if err != nil {
|
||||
return q, err
|
||||
}
|
||||
return q, nil
|
||||
}
|
||||
|
||||
func (q *BandwithQuantity) Equal(u *BandwithQuantity) bool {
|
||||
if q == nil && u == nil {
|
||||
return true
|
||||
}
|
||||
if q != nil && u != nil {
|
||||
return q.i == u.i
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (q *BandwithQuantity) String() string {
|
||||
return q.s
|
||||
}
|
||||
|
||||
func (q *BandwithQuantity) UnmarshalString(s string) error {
|
||||
q.s = strings.TrimSpace(s)
|
||||
if q.s == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
var (
|
||||
base int64
|
||||
f float64
|
||||
err error
|
||||
)
|
||||
if strings.HasSuffix(s, "MB") {
|
||||
base = MB
|
||||
s = strings.TrimSuffix(s, "MB")
|
||||
f, err = strconv.ParseFloat(s, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
} else if strings.HasSuffix(s, "KB") {
|
||||
base = KB
|
||||
s = strings.TrimSuffix(s, "KB")
|
||||
f, err = strconv.ParseFloat(s, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
return errors.New("unit not support")
|
||||
}
|
||||
|
||||
q.i = int64(f * float64(base))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (q *BandwithQuantity) UnmarshalJSON(b []byte) error {
|
||||
return q.UnmarshalString(string(b))
|
||||
}
|
||||
|
||||
func (q *BandwithQuantity) MarshalJSON() ([]byte, error) {
|
||||
return []byte(q.s), nil
|
||||
}
|
||||
|
||||
func (q *BandwithQuantity) Bytes() int64 {
|
||||
return q.i
|
||||
}
|
||||
@ -0,0 +1,51 @@
|
||||
// Copyright 2019 fatedier, fatedier@gmail.com
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package limit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
type Reader struct {
|
||||
r io.Reader
|
||||
limiter *rate.Limiter
|
||||
}
|
||||
|
||||
func NewReader(r io.Reader, limiter *rate.Limiter) *Reader {
|
||||
return &Reader{
|
||||
r: r,
|
||||
limiter: limiter,
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Reader) Read(p []byte) (n int, err error) {
|
||||
b := r.limiter.Burst()
|
||||
if b < len(p) {
|
||||
p = p[:b]
|
||||
}
|
||||
n, err = r.r.Read(p)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
err = r.limiter.WaitN(context.Background(), n)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
@ -0,0 +1,60 @@
|
||||
// Copyright 2019 fatedier, fatedier@gmail.com
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package limit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
type Writer struct {
|
||||
w io.Writer
|
||||
limiter *rate.Limiter
|
||||
}
|
||||
|
||||
func NewWriter(w io.Writer, limiter *rate.Limiter) *Writer {
|
||||
return &Writer{
|
||||
w: w,
|
||||
limiter: limiter,
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Writer) Write(p []byte) (n int, err error) {
|
||||
var nn int
|
||||
b := w.limiter.Burst()
|
||||
for {
|
||||
end := len(p)
|
||||
if end == 0 {
|
||||
break
|
||||
}
|
||||
if b < len(p) {
|
||||
end = b
|
||||
}
|
||||
err = w.limiter.WaitN(context.Background(), end)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
nn, err = w.w.Write(p[:end])
|
||||
n += nn
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
p = p[end:]
|
||||
}
|
||||
return
|
||||
}
|
||||
Loading…
Reference in New Issue