180 lines
4.5 KiB
Go
180 lines
4.5 KiB
Go
//Copyright 2022 gdy, 272288813@qq.com
|
|
package base
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/gdy666/lucky/thirdlib/gdylib/pool"
|
|
)
|
|
|
|
type Proxy interface {
|
|
StartProxy()
|
|
StopProxy()
|
|
|
|
ReceiveDataCallback(int64)
|
|
SendDataCallback(int64)
|
|
GetProxyType() string
|
|
GetStatus() string
|
|
GetListenIP() string
|
|
GetListenPort() int
|
|
GetKey() string
|
|
GetCurrentConnections() int64
|
|
SetFromRule(string)
|
|
FromRule() string
|
|
String() string
|
|
GetTrafficIn() int64
|
|
GetTrafficOut() int64
|
|
SafeCheck(ip string) bool
|
|
}
|
|
|
|
type RelayRuleOptions struct {
|
|
UDPPackageSize int `json:"UDPPackageSize,omitempty"`
|
|
SingleProxyMaxConnections int64 `json:"SingleProxyMaxConnections,omitempty"`
|
|
UDPProxyPerformanceMode bool `json:"UDPProxyPerformanceMode,omitempty"`
|
|
UDPShortMode bool `json:"UDPShortMode,omitempty"`
|
|
SafeMode string `json:"SafeMode,omitempty"`
|
|
}
|
|
|
|
// Join two io.ReadWriteCloser and do some operations.
|
|
func (p *BaseProxyConf) relayData(targetServer io.ReadWriteCloser, client io.ReadWriteCloser) {
|
|
var wait sync.WaitGroup
|
|
pipe := func(to io.ReadWriteCloser, from io.ReadWriteCloser, writedataCallback func(int64)) {
|
|
defer to.Close()
|
|
defer from.Close()
|
|
defer wait.Done()
|
|
|
|
nw, _ := p.copyBuffer(to, from, nil, nil)
|
|
if writedataCallback != nil {
|
|
writedataCallback(nw)
|
|
}
|
|
|
|
// if p.TrafficMonitor {
|
|
// buf := pool.GetBuf(8 * 1024 * 1024)
|
|
// p.CopyBuffer(to, from, buf, writedataCallback)
|
|
// pool.PutBuf(buf)
|
|
// } else {
|
|
// nw, _ := p.copyBuffer(to, from, nil, nil)
|
|
// if writedataCallback != nil {
|
|
// writedataCallback(nw)
|
|
// }
|
|
// }
|
|
|
|
}
|
|
|
|
wait.Add(2)
|
|
go pipe(targetServer, client, p.ReceiveDataCallback)
|
|
go pipe(client, targetServer, p.SendDataCallback)
|
|
wait.Wait()
|
|
}
|
|
|
|
func (p *BaseProxyConf) CopyBuffer(dst io.Writer, src io.Reader, buf []byte, writedataCallback func(int64)) (written int64, err error) {
|
|
if buf != nil && len(buf) == 0 {
|
|
panic("empty buffer in CopyBuffer")
|
|
}
|
|
return p.copyBuffer(dst, src, buf, writedataCallback)
|
|
}
|
|
|
|
// copyBuffer is the actual implementation of Copy and CopyBuffer.
|
|
// if buf is nil, one is allocated.
|
|
func (p *BaseProxyConf) copyBuffer(dst io.Writer, src io.Reader, buf []byte, writedataCallback func(int64)) (written int64, err error) {
|
|
if buf == nil {
|
|
if wt, ok := src.(io.WriterTo); ok {
|
|
return wt.WriteTo(dst)
|
|
}
|
|
|
|
if rt, ok := dst.(io.ReaderFrom); ok {
|
|
return rt.ReadFrom(src)
|
|
}
|
|
|
|
size := 32 * 1024
|
|
if l, ok := src.(*io.LimitedReader); ok && int64(size) > l.N {
|
|
if l.N < 1 {
|
|
size = 1
|
|
} else {
|
|
size = int(l.N)
|
|
}
|
|
}
|
|
|
|
buf = pool.GetBuf(8 * size)
|
|
defer pool.PutBuf(buf)
|
|
}
|
|
|
|
for {
|
|
nr, er := src.Read(buf)
|
|
if nr > 0 {
|
|
nw, ew := dst.Write(buf[0:nr])
|
|
if nw < 0 || nr < nw {
|
|
nw = 0
|
|
if ew == nil {
|
|
ew = errors.New("invalid write result")
|
|
}
|
|
}
|
|
written += int64(nw)
|
|
|
|
if writedataCallback != nil {
|
|
writedataCallback(int64(nw))
|
|
}
|
|
|
|
if ew != nil {
|
|
err = ew
|
|
break
|
|
}
|
|
if nr != nw {
|
|
err = io.ErrShortWrite
|
|
break
|
|
}
|
|
}
|
|
if er != nil {
|
|
if er != io.EOF {
|
|
err = er
|
|
}
|
|
break
|
|
}
|
|
}
|
|
|
|
return written, err
|
|
}
|
|
|
|
func formatFileSize(fileSize int64) (size string) {
|
|
switch {
|
|
case fileSize < 1024:
|
|
return fmt.Sprintf("%.2fB", float64(fileSize)/float64(1))
|
|
case fileSize < (1024 * 1024):
|
|
return fmt.Sprintf("%.2fKB", float64(fileSize)/float64(1024))
|
|
case fileSize < (1024 * 1024 * 1024):
|
|
return fmt.Sprintf("%.2fMB", float64(fileSize)/float64(1024*1024))
|
|
case fileSize < (1024 * 1024 * 1024 * 1024):
|
|
return fmt.Sprintf("%.2fGB", float64(fileSize)/float64(1024*1024*1024))
|
|
case fileSize < (1024 * 1024 * 1024 * 1024 * 1024):
|
|
return fmt.Sprintf("%.2fTB", float64(fileSize)/float64(1024*1024*1024*1024))
|
|
default:
|
|
return fmt.Sprintf("%.2fEB", float64(fileSize)/float64(1024*1024*1024*1024*1024))
|
|
}
|
|
|
|
}
|
|
|
|
func CreateProxy(proxyType, listenIP, targetHost string, balanceTargetAddressList *[]string, listenPort, targetPort int, options *RelayRuleOptions) (p Proxy, err error) {
|
|
//key := GetProxyKey(proxyType, listenIP, listenPort)
|
|
switch {
|
|
case strings.HasPrefix(proxyType, "tcp"):
|
|
{
|
|
return CreateTCPProxy(proxyType, listenIP, targetHost, balanceTargetAddressList, listenPort, targetPort, options), nil
|
|
}
|
|
case strings.HasPrefix(proxyType, "udp"):
|
|
{
|
|
return CreateUDPProxy(proxyType, listenIP, targetHost, balanceTargetAddressList, listenPort, targetPort, options), nil
|
|
}
|
|
default:
|
|
return nil, fmt.Errorf("未支持的类型:%s", proxyType)
|
|
}
|
|
|
|
}
|
|
|
|
func GetProxyKey(proxyType, listenIP string, listenPort int) string {
|
|
return fmt.Sprintf("%s@%s:%d", proxyType, listenIP, listenPort)
|
|
}
|