first commit
This commit is contained in:
142
pkg/proxy/io.go
Normal file
142
pkg/proxy/io.go
Normal file
@ -0,0 +1,142 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"golang.org/x/time/rate"
|
||||
"io"
|
||||
"net"
|
||||
"runtime/debug"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const burstLimit = 1000 * 1000 * 1000
|
||||
|
||||
type Reader struct {
|
||||
r io.Reader
|
||||
limiter *rate.Limiter
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
func NewReader(r io.Reader) *Reader {
|
||||
return &Reader{
|
||||
r: r,
|
||||
ctx: context.Background(),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Reader) SetRateLimit(bytesPerSec float64) {
|
||||
s.limiter = rate.NewLimiter(rate.Limit(bytesPerSec), burstLimit)
|
||||
s.limiter.AllowN(time.Now(), burstLimit) // spend initial burst
|
||||
}
|
||||
|
||||
func (s *Reader) Read(p []byte) (int, error) {
|
||||
if s.limiter == nil {
|
||||
return s.r.Read(p)
|
||||
}
|
||||
n, err := s.r.Read(p)
|
||||
if err != nil {
|
||||
return n, err
|
||||
}
|
||||
if err := s.limiter.WaitN(s.ctx, n); err != nil {
|
||||
return n, err
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func ConnectHost(hostAndPort string, timeout int) (conn net.Conn, err error) {
|
||||
conn, err = net.DialTimeout("tcp", hostAndPort, time.Duration(timeout)*time.Millisecond)
|
||||
return
|
||||
}
|
||||
|
||||
func CloseConn(conn net.Conn) {
|
||||
if conn != nil {
|
||||
_ = conn.SetDeadline(time.Now().Add(time.Millisecond))
|
||||
_ = conn.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func IoBind(dst io.ReadWriter, src io.ReadWriter, fn func(isSrcErr bool, err error), cfn func(count int, isPositive bool), bytesPreSec float64) {
|
||||
var one = &sync.Once{}
|
||||
go func() {
|
||||
defer func() {
|
||||
if e := recover(); e != nil {
|
||||
logger.Sugar().Errorf("IoBind crashed , err : %s , \ntrace:%s", e, string(debug.Stack()))
|
||||
}
|
||||
}()
|
||||
var err error
|
||||
var isSrcErr bool
|
||||
if bytesPreSec > 0 {
|
||||
newReader := NewReader(src)
|
||||
newReader.SetRateLimit(bytesPreSec)
|
||||
_, isSrcErr, err = IoCopy(dst, newReader, func(c int) {
|
||||
cfn(c, false)
|
||||
})
|
||||
|
||||
} else {
|
||||
_, isSrcErr, err = IoCopy(dst, src, func(c int) {
|
||||
cfn(c, false)
|
||||
})
|
||||
}
|
||||
if err != nil {
|
||||
one.Do(func() {
|
||||
fn(isSrcErr, err)
|
||||
})
|
||||
}
|
||||
}()
|
||||
go func() {
|
||||
defer func() {
|
||||
if e := recover(); e != nil {
|
||||
logger.Sugar().Errorf("IoBind crashed , err : %s , \ntrace:%s", e, string(debug.Stack()))
|
||||
}
|
||||
}()
|
||||
var err error
|
||||
var isSrcErr bool
|
||||
if bytesPreSec > 0 {
|
||||
newReader := NewReader(dst)
|
||||
newReader.SetRateLimit(bytesPreSec)
|
||||
_, isSrcErr, err = IoCopy(src, newReader, func(c int) {
|
||||
cfn(c, true)
|
||||
})
|
||||
} else {
|
||||
_, isSrcErr, err = IoCopy(src, dst, func(c int) {
|
||||
cfn(c, true)
|
||||
})
|
||||
}
|
||||
if err != nil {
|
||||
one.Do(func() {
|
||||
fn(isSrcErr, err)
|
||||
})
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func IoCopy(dst io.Writer, src io.Reader, fn ...func(count int)) (written int64, isSrcErr bool, err error) {
|
||||
buf := make([]byte, 32*1024)
|
||||
for {
|
||||
nr, er := src.Read(buf)
|
||||
if nr > 0 {
|
||||
nw, ew := dst.Write(buf[0:nr])
|
||||
if nw > 0 {
|
||||
written += int64(nw)
|
||||
if len(fn) == 1 {
|
||||
fn[0](nw)
|
||||
}
|
||||
}
|
||||
if ew != nil {
|
||||
err = ew
|
||||
break
|
||||
}
|
||||
if nr != nw {
|
||||
err = io.ErrShortWrite
|
||||
break
|
||||
}
|
||||
}
|
||||
if er != nil {
|
||||
err = er
|
||||
isSrcErr = true
|
||||
break
|
||||
}
|
||||
}
|
||||
return written, isSrcErr, err
|
||||
}
|
||||
105
pkg/proxy/service.go
Normal file
105
pkg/proxy/service.go
Normal file
@ -0,0 +1,105 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"go.uber.org/zap"
|
||||
"net"
|
||||
"runtime/debug"
|
||||
"sync"
|
||||
)
|
||||
|
||||
var (
|
||||
servicesMap = new(sync.Map)
|
||||
logger *zap.Logger
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
TCPConn TCP
|
||||
Name string
|
||||
}
|
||||
|
||||
func (s *Service) Stop() {
|
||||
servicesMap.Delete(s.Name)
|
||||
s.TCPConn.Close()
|
||||
}
|
||||
|
||||
func Run(name string, args TCPArgs, zapLogger *zap.Logger) *Service {
|
||||
logger = zapLogger
|
||||
service := &Service{
|
||||
TCPConn: &tcp{cfg: args},
|
||||
Name: name,
|
||||
}
|
||||
store, loaded := servicesMap.LoadOrStore(name, service)
|
||||
if loaded {
|
||||
service = store.(*Service)
|
||||
}
|
||||
|
||||
go func() {
|
||||
defer func() {
|
||||
recoverErr := recover()
|
||||
if recoverErr != nil {
|
||||
logger.Sugar().Errorf("%s servcie crashed, ERR: %s\ntrace:%s", name, recoverErr, string(debug.Stack()))
|
||||
}
|
||||
}()
|
||||
startErr := service.TCPConn.Start()
|
||||
if startErr != nil {
|
||||
logger.Sugar().Errorf("%s servcie fail, ERR: %s", name, startErr)
|
||||
}
|
||||
}()
|
||||
|
||||
return service
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
|
||||
type Listener struct {
|
||||
ip string
|
||||
port int
|
||||
Listener net.Listener
|
||||
errAcceptHandler func(err error)
|
||||
}
|
||||
|
||||
func NewListener(ip string, port int) Listener {
|
||||
return Listener{
|
||||
ip: ip,
|
||||
port: port,
|
||||
errAcceptHandler: func(err error) {
|
||||
logger.Sugar().Errorf("accept error , ERR:%s", err)
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (sc *Listener) ListenTCP(fn func(conn net.Conn)) (err error) {
|
||||
sc.Listener, err = net.Listen("tcp", fmt.Sprintf("%s:%d", sc.ip, sc.port))
|
||||
if err == nil {
|
||||
go func() {
|
||||
defer func() {
|
||||
if e := recover(); e != nil {
|
||||
logger.Sugar().Infof("ListenTCP crashed , err : %s , \ntrace:%s", e, string(debug.Stack()))
|
||||
}
|
||||
}()
|
||||
for {
|
||||
var conn net.Conn
|
||||
conn, err = sc.Listener.Accept()
|
||||
if err == nil {
|
||||
go func() {
|
||||
defer func() {
|
||||
if e := recover(); e != nil {
|
||||
logger.Sugar().Infof("connection handler crashed , err : %s , \ntrace:%s", e, string(debug.Stack()))
|
||||
}
|
||||
}()
|
||||
fn(conn)
|
||||
}()
|
||||
} else {
|
||||
sc.errAcceptHandler(err)
|
||||
break
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (sc *Listener) CloseListen() error {
|
||||
return sc.Listener.Close()
|
||||
}
|
||||
94
pkg/proxy/tcp.go
Normal file
94
pkg/proxy/tcp.go
Normal file
@ -0,0 +1,94 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"net"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
var _ TCP = (*tcp)(nil)
|
||||
|
||||
type TCPArgs struct {
|
||||
Local string //监听地址
|
||||
Parent string //被代理地址
|
||||
Timeout int //拨号超时(毫秒)
|
||||
OutCallback func() bool //回调 //是否允许代理
|
||||
}
|
||||
|
||||
type TCP interface {
|
||||
Start() (err error)
|
||||
Close()
|
||||
|
||||
callback(inConn net.Conn)
|
||||
outToTCP(inConn net.Conn) (err error)
|
||||
}
|
||||
|
||||
type tcp struct {
|
||||
inConn net.Conn
|
||||
outConn net.Conn
|
||||
listen Listener
|
||||
cfg TCPArgs
|
||||
}
|
||||
|
||||
func (s *tcp) Start() (err error) {
|
||||
host, port, _ := net.SplitHostPort(s.cfg.Local)
|
||||
p, _ := strconv.Atoi(port)
|
||||
s.listen = NewListener(host, p)
|
||||
err = s.listen.ListenTCP(s.callback)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (s *tcp) Close() {
|
||||
if s.inConn != nil {
|
||||
CloseConn(s.inConn)
|
||||
}
|
||||
if s.outConn != nil {
|
||||
CloseConn(s.outConn)
|
||||
}
|
||||
_ = s.listen.CloseListen()
|
||||
}
|
||||
|
||||
func (s *tcp) callback(inConn net.Conn) {
|
||||
defer func() {
|
||||
if err := recover(); err != nil {
|
||||
logger.Sugar().Infof("conn handler crashed with err : %s", err)
|
||||
}
|
||||
}()
|
||||
|
||||
if s.cfg.OutCallback != nil {
|
||||
if !s.cfg.OutCallback() {
|
||||
CloseConn(inConn)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
err := s.outToTCP(inConn)
|
||||
if err != nil {
|
||||
CloseConn(inConn)
|
||||
}
|
||||
|
||||
s.inConn = inConn
|
||||
}
|
||||
|
||||
func (s *tcp) outToTCP(inConn net.Conn) error {
|
||||
outConn, err := ConnectHost(s.cfg.Parent, s.cfg.Timeout)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
inAddr := inConn.RemoteAddr().String()
|
||||
inLocalAddr := inConn.LocalAddr().String()
|
||||
outAddr := outConn.RemoteAddr().String()
|
||||
outLocalAddr := outConn.LocalAddr().String()
|
||||
IoBind(inConn, outConn, func(isSrcErr bool, err error) {
|
||||
CloseConn(inConn)
|
||||
CloseConn(outConn)
|
||||
logger.Sugar().Infof("conn %s - %s - %s -%s released", inAddr, inLocalAddr, outLocalAddr, outAddr)
|
||||
}, func(n int, d bool) {}, 0)
|
||||
|
||||
logger.Sugar().Infof("conn %s - %s - %s -%s connected", inAddr, inLocalAddr, outLocalAddr, outAddr)
|
||||
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user